单个管道中的多个 Spark DataFrame 突变

Question

Consider a Spark DataFrame df with the following schema:

root 
|-- date: timestamp (nullable = true) 
|-- customerID: string (nullable = true) 
|-- orderID: string (nullable = true) 
|-- productID: string (nullable = true)

One column should be cast to a different type, other columns should just have their white-space trimmed.

df.select(
  $"date",
  df("customerID").cast(IntegerType),
  $"orderID",
  $"productId")
  .withColumn("orderID", trim(col("orderID")))
  .withColumn("productID", trim(col("productID")))

The operations seem to require different syntax; casting is done via select , while trim is done via withColumn . I'm used to R and dplyr where all the above would be handled in a single mutate function, so mixing select and withColumn feels a bit cumbersome.

Is there a cleaner way to do this in a single pipe?

Answer 1

df.select(
  $"date",
  $"customerID".cast(IntegerType),
  trim($"orderID").as("orderID"),
  trim($"productID").as("productID"))

Answer 2

You can use either one. The difference is that withColumn will add (or replace if the same name is used) a new column to the dataframe while select will only keep the columns you specified. Depending on the situation, choose one to use.

The cast can be done using withColumn as follows:

df.withColumn("customerID", $"customerID".cast(IntegerType))
  .withColumn("orderID", trim($"orderID"))
  .withColumn("productID", trim($"productID"))

Note that you do not need to use withColumn on the date column above.

---

The trim functions can be done in a select as follows, here the column names are kept the same:

df.select(
  $"date",
  $"customerID".cast(IntegerType),
  trim($"orderID").as("orderID"),
  trim($"productId").as("productId"))