[SPARK-49022] Use Column Node API in Column

[hvanhovell]

What changes were proposed in this pull request?

This PR makes the org.apache.spark.sql.Column and friends use the recently introduced ColumnNode API. This is a stepping stone towards making the Column API implementation agnostic.

Most of the changes are fairly mechanical, and they are mostly caused by the removal of the Column(Expression) constructor.

Why are the changes needed?

We want to create unified Scala interface for Classic and Connect. A language agnostic Column API implementation is part of this.

Does this PR introduce any user-facing change?

No

How was this patch tested?

Existing tests.

Was this patch authored or co-authored using generative AI tooling?

No

[hvanhovell]

@HyukjinKwon I cannot invoke Column.fn(...) here. Do you know why? I added Column.pysparkFn(...) as a workaround.

[HyukjinKwon]

Ah, the Column.fn has to be decorated with scala.annotation.varargs (which will create the same signature for variant arguments for Java).

[hvanhovell]

This is waiting for #47746.

[HyukjinKwon]

not so related to this PR but I wonder if we should name the package a little bit differently, e.g., org.apache.spark.sql.internal.api. At least I am checking which package does individual class belongs to.

[hvanhovell]

You mean move all ColumnNode classes elsewhere?

[HyukjinKwon]

I think the error report here would be a bit weird via Origin because the top level of the function call changed. We get Thread.currentThread().getStackTrace, and set it to Origin so whichever function fails, it will always point it here.

Would be great if we can double check the code like the below still works fine:

val df = spark.range(10)

val df1 = df.withColumn("div_ten", df.col("id") / 10)
val df2 = df1.withColumn("plus_four", df.col("id") + 4)

// This is problematic divide operation that occurs DIVIDE_BY_ZERO.
val df3 = df2.withColumn("div_zero", df.col("id") / 0) // Error here
val df4 = df3.withColumn("minus_five", df.col("id") / 5)
df4.collect()

should report sth like:

org.apache.spark.SparkArithmeticException: [DIVIDE_BY_ZERO] Division by zero. Use `try_divide` to tolerate divisor being 0 and return NULL instead. If necessary set "spark.sql.ansi.enabled" to "false" to bypass this error. SQLSTATE: 22012
== DataFrame ==
"div" was called from
<init>(<console>:7)

[hvanhovell]

This is what get:

org.apache.spark.SparkArithmeticException: [DIVIDE_BY_ZERO] Division by zero. Use `try_divide` to tolerate divisor being 0 and return NULL instead. If necessary set "spark.sql.ansi.enabled" to "false" to bypass this error. SQLSTATE: 22012
== DataFrame ==
"div" was called from
<init>(<console>:1)

So that checks out. TBH this part is almost the same as before.

[HyukjinKwon]

Looks good otherwise. Mostly minor comments. Only one real comment from me is https://github.com/apache/spark/pull/47688/files#r1719271705 to make sure error reporting is still working fine.

[hvanhovell]

Tests have passed. I am merging this to master.

[HyukjinKwon]

Merged to master.

[EnricoMi]

@hvanhovell what is the recommended way to migrate user code like

new Column(MyExpression())

to this Column Node Api?

The following are all private[sql] / private[spark]:

Column(MyExpression())
new Column(ExpressionColumnNode(MyExpression()))

[hvanhovell]

@EnricoMi what are you exactly doing? Could you share an example? We can definitely open up the interface to allow for user extensions.

If it is basically adding functions, then I'd probably go with the Column.fn(...) approach in combination with a SparkSessionExtension that will register your custom expressions. This way we can easily make what you have created work with Spark Connect (and with SQL). I understand some folks are not interested in this, in that case we can still open up an expression way of creating a Column.

[EnricoMi]

... in that case we can still open up an expression way of creating a Column.

That would shortcut a lot of extra complexity and simplify backward compatibility.

[HyukjinKwon]

backward compatibility

Which backward compatibility? Showing an example would help a lot to understand. All those expressions are subject to be internal and private.

[EnricoMi]

The spark-extension packages provides some Dataset diff tooling. There, a user-defined comparison can simply be defined by implementing the scala.math.Equiv interface: https://github.com/G-Research/spark-extension/blob/master/src/main/scala/uk/co/gresearch/spark/diff/DiffComparators.scala#L41

That Equiv implementation is wrapped into an Expression (including codegen) and turned into a Comparatorthat is then used by the package to diff columns: given two columns left and right, return a Column that evaluates (compares the columns) to Boolean:

• EquivDiffComparator.scala:34
• EquivDiffComparator.scala:67

This obviously won't work for Spark connect, but with Column Node API this does not work for classic Spark client either.

That package supports Spark 3.0 - 3.5. Creating a Column from an Expression would allow for minimal changes to keep this working for Spark 4.0 with non-Connect client. This is what I meant with backward compatibility.

In order to support Spark Connect, there is no way around using the Spark Connect plugin / extensions.

[EnricoMi]

Currently working around this by gaining access to Column(Expression) in private[spark] package: https://github.com/G-Research/spark-extension/pull/256/files