Replace multiple calls to withColumn with single select to simplify query plans

[earangol-stripe]

Summary

Refactors some of the join code to avoid multiple calls to withColumn and withColumnRenamed.

Why / Goal

Performance. Multiple calls to withColumn (and withColumnRenamed) are known to cause performance issues when done on too many columns, as it may generate very complex query plans. At Stripe, some of the jobs with too many columns on the RHS (e.g. ~2k) may fail due to StackOverflowErrors on the driver when generating the query plan.

The following code snippet illustrates the issue, given a DF with 2k columns:

val range = 0 until 2000
val columns = range.map(i => s"column$i")
val schema = columns.map(c => StructField(c, IntegerType, true))
val df = spark.createDataFrame(spark.sparkContext.emptyRDD[Row], StructType(schema))

On a notebook, this code took me ~3 minutes to run:

val renamedDf = df.columns.foldLeft(df) {
  case (partialDf, c) =>
    df.withColumnRenamed(c, s"renamed_$c")
}
renamedDf.explain() // slow

While this code took ~1 second:

val renamedColumns = df.columns.map(c => col(c).as(s"renamed_$c"))
val renamedDf = df.select(renamedColumns:_*)
renamedDf.explain() // fast

Test Plan

• Added Unit Tests
• Covered by existing CI
• Integration tested

Checklist

• Documentation update

Reviewers

@pengyu-hou @jbrooks-stripe

[nikhilsimha]

very cool stuff!! lgtm will let others stamp

[thomaschow]

lgtm

[earangol-stripe]

Thanks for the review!

[earangol-stripe]

@thomaschow -- need another stamp, I had to rebase main.

[nikhilsimha]

LGTM!

[earangol-stripe]

@thomaschow @nikhilsimha -- the branch got out of date and I had to merge main again so this needs another stamp. Thanks!

[thomaschow]

LGTM, purely a review