Support Spark4 Column Node API

G-Research · Aug 20, 2024 · d47f0b7 · d47f0b7
1 parent 568eac0
commit d47f0b7
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/src/main/scala/uk/co/gresearch/spark/diff/comparator/EquivDiffComparator.scala b/src/main/scala/uk/co/gresearch/spark/diff/comparator/EquivDiffComparator.scala
@@ -16,6 +16,7 @@
 
 package uk.co.gresearch.spark.diff.comparator
 
+import org.apache.spark.extension.ExpressionExtension
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.encoders.encoderFor
 import org.apache.spark.sql.catalyst.expressions.codegen.Block.BlockHelper
@@ -31,8 +32,7 @@ trait EquivDiffComparator[T] extends DiffComparator {
 
 private trait ExpressionEquivDiffComparator[T] extends EquivDiffComparator[T] {
   def equiv(left: Expression, right: Expression): EquivExpression[T]
-  def equiv(left: Column, right: Column): Column =
-    new Column(equiv(left.expr, right.expr).asInstanceOf[Expression])
+  def equiv(left: Column, right: Column): Column = equiv(left.expr, right.expr).toColumn
 }
 
 trait TypedEquivDiffComparator[T] extends EquivDiffComparator[T] with TypedDiffComparator

diff --git a/src/main/scala/uk/co/gresearch/spark/package.scala b/src/main/scala/uk/co/gresearch/spark/package.scala
@@ -16,6 +16,7 @@
 
 package uk.co.gresearch
 
+import org.apache.spark.extension.ExpressionExtension
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql._
 import org.apache.spark.sql.catalyst.expressions.NamedExpression
@@ -271,7 +272,7 @@ package object spark extends Logging with SparkVersion with BuildVersion {
    *   result tick value column
    */
   def timestampToDotNetTicks(timestampColumn: Column): Column =
-    unixEpochTenthMicrosToDotNetTicks(new Column(UnixMicros.unixMicros(timestampColumn.expr)) * 10)
+    unixEpochTenthMicrosToDotNetTicks(UnixMicros.unixMicros(timestampColumn.expr).toColumn * 10)
 
   /**
    * Convert a Spark timestamp to a .Net `DateTime.Ticks` timestamp. The input column must be of TimestampType.