Spark 3.4: Cast type when calling projection, support recursive resolve

(cherry picked from commit 936a18a)

Author: Yxang

spark-3.4/clickhouse-spark-it/src/test/scala/org/apache/spark/sql/clickhouse/cluster/ClusterShardByTransformSuite.scala

@@ -94,16 +94,21 @@ class ClusterShardByTransformSuite extends SparkClickHouseClusterTest {
   }
 
   Seq(
+    // wait for SPARK-44180 to be fixed, then add implicit cast test cases
     ("toYear", Array("create_date")),
+//    ("toYear", Array("create_time")),
     ("toYYYYMM", Array("create_date")),
+//    ("toYYYYMM", Array("create_time")),
     ("toYYYYMMDD", Array("create_date")),
+//    ("toYYYYMMDD", Array("create_time")),
     ("toHour", Array("create_time")),
     ("xxHash64", Array("value")),
     ("murmurHash2_64", Array("value")),
     ("murmurHash2_32", Array("value")),
     ("murmurHash3_64", Array("value")),
     ("murmurHash3_32", Array("value")),
-    ("cityHash64", Array("value"))
+    ("cityHash64", Array("value")),
+    ("positiveModulo", Array("toYYYYMM(create_date)", "10"))
   ).foreach {
     case (func_name: String, func_args: Array[String]) =>
       test(s"shard by $func_name(${func_args.mkString(",")})")(runTest(func_name, func_args))

spark-3.4/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/ExprUtils.scala

@@ -15,9 +15,15 @@
 package org.apache.spark.sql.clickhouse
 
 import org.apache.spark.sql.AnalysisException
-import org.apache.spark.sql.catalyst.SQLConfHelper
 import org.apache.spark.sql.catalyst.analysis.NoSuchFunctionException
-import org.apache.spark.sql.catalyst.expressions.{BoundReference, Expression, TransformExpression}
+import org.apache.spark.sql.catalyst.{expressions, SQLConfHelper}
+import org.apache.spark.sql.catalyst.expressions.{
+  BoundReference,
+  Cast,
+  Expression,
+  TransformExpression,
+  V2ExpressionUtils
+}
 import org.apache.spark.sql.clickhouse.ClickHouseSQLConf.IGNORE_UNSUPPORTED_TRANSFORM
 import org.apache.spark.sql.connector.catalog.Identifier
 import org.apache.spark.sql.connector.catalog.functions.{BoundFunction, ScalarFunction, UnboundFunction}
@@ -89,6 +95,20 @@ object ExprUtils extends SQLConfHelper with Serializable {
     }
   }
 
+  def resolveTransformCatalyst(
+    catalystExpr: Expression,
+    timeZoneId: Option[String] = None
+  ): Expression = catalystExpr match {
+    case TransformExpression(function: ScalarFunction[_], args, _) =>
+      val resolvedArgs: Seq[Expression] = args.map(resolveTransformCatalyst(_, timeZoneId))
+      val castedArgs: Seq[Expression] = resolvedArgs.zip(function.inputTypes()).map {
+        case (arg, expectedType) if !arg.dataType.sameType(expectedType) => Cast(arg, expectedType, timeZoneId)
+        case (arg, _) => arg
+      }
+      V2ExpressionUtils.resolveScalarFunction(function, castedArgs)
+    case other => other
+  }
+
   def toCatalyst(
     v2Expr: V2Expression,
     fields: Array[StructField],
@@ -108,6 +128,7 @@ object ExprUtils extends SQLConfHelper with Serializable {
           .map(bound => TransformExpression(bound, catalystArgs)).getOrElse {
             throw CHClientException(s"Unsupported expression: $v2Expr")
           }
+      case literal: LiteralValue[Any] => expressions.Literal(literal.value)
       case _ => throw CHClientException(
           s"Unsupported expression: $v2Expr"
         )

spark-3.4/clickhouse-spark/src/main/scala/xenon/clickhouse/write/ClickHouseWriter.scala

@@ -17,16 +17,9 @@ package xenon.clickhouse.write
 import com.clickhouse.client.ClickHouseProtocol
 import com.clickhouse.data.ClickHouseCompression
 import org.apache.commons.io.IOUtils
-import org.apache.spark.sql.catalyst.expressions.{
-  BoundReference,
-  Expression,
-  SafeProjection,
-  TransformExpression,
-  V2ExpressionUtils
-}
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, Expression, SafeProjection, TransformExpression}
 import org.apache.spark.sql.catalyst.{expressions, InternalRow}
 import org.apache.spark.sql.clickhouse.ExprUtils
-import org.apache.spark.sql.connector.catalog.functions.ScalarFunction
 import org.apache.spark.sql.connector.metric.CustomTaskMetric
 import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage}
 import org.apache.spark.sql.types._
@@ -86,23 +79,21 @@ abstract class ClickHouseWriter(writeJob: WriteJobDescription)
 
   protected lazy val shardProjection: Option[expressions.Projection] = shardExpr
     .filter(_ => writeJob.writeOptions.convertDistributedToLocal)
-    .flatMap(expr =>
-      expr match {
-        case BoundReference(_, _, _) =>
-          Some(SafeProjection.create(Seq(expr)))
-        case TransformExpression(function, args, _) =>
-          val retType = function.resultType() match {
-            case ByteType => classOf[Byte]
-            case ShortType => classOf[Short]
-            case IntegerType => classOf[Int]
-            case LongType => classOf[Long]
-            case _ => throw CHClientException(s"Invalid return data type for function ${function.name()}," +
-                s"sharding field: ${function.resultType()}")
-          }
-          val expr = V2ExpressionUtils.resolveScalarFunction(function.asInstanceOf[ScalarFunction[retType.type]], args)
-          Some(SafeProjection.create(Seq(expr)))
-      }
-    )
+    .flatMap {
+      case expr: BoundReference =>
+        Some(SafeProjection.create(Seq(expr)))
+      case expr @ TransformExpression(function, _, _) =>
+        // result type must be integer class
+        function.resultType() match {
+          case ByteType => classOf[Byte]
+          case ShortType => classOf[Short]
+          case IntegerType => classOf[Int]
+          case LongType => classOf[Long]
+          case _ => throw CHClientException(s"Invalid return data type for function ${function.name()}," +
+              s"sharding field: ${function.resultType()}")
+        }
+        Some(SafeProjection.create(Seq(ExprUtils.resolveTransformCatalyst(expr, Some(writeJob.tz.getId)))))
+    }
 
   // put the node select strategy in executor side because we need to calculate shard and don't know the records
   // util DataWriter#write(InternalRow) invoked.