Add source data type for connector cast expression

urosstan-db · May 15, 2024 · f03f08b · f03f08b
1 parent d0385c4
commit f03f08b
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 2 deletions.
diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/Cast.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/Cast.java
@@ -29,14 +29,31 @@
 @Evolving
 public class Cast extends ExpressionWithToString {
   private Expression expression;
+
+  /**
+   * Original data type of given expression
+   */
+  private DataType expressionDataType;
+
+  /**
+   * Target data type, i.e. data type in which expression will be cast
+   */
   private DataType dataType;
 
+  @Deprecated
   public Cast(Expression expression, DataType dataType) {
     this.expression = expression;
     this.dataType = dataType;
   }
 
+  public Cast(Expression expression, DataType expressionDataType, DataType targetDataType) {
+    this.expression = expression;
+    this.expressionDataType = expressionDataType;
+    this.dataType = targetDataType;
+  }
+
   public Expression expression() { return expression; }
+  public DataType expressionDataType() { return expressionDataType; }
   public DataType dataType() { return dataType; }
 
   @Override

diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/util/V2ExpressionSQLBuilder.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/util/V2ExpressionSQLBuilder.java
@@ -79,7 +79,7 @@ public String build(Expression expr) {
     } else if (expr instanceof NamedReference namedReference) {
       return visitNamedReference(namedReference);
     } else if (expr instanceof Cast cast) {
-      return visitCast(build(cast.expression()), cast.dataType());
+      return visitCast(build(cast.expression()), cast.expressionDataType(), cast.dataType());
     } else if (expr instanceof Extract extract) {
       return visitExtract(extract.field(), build(extract.source()));
     } else if (expr instanceof SortOrder sortOrder) {
@@ -231,10 +231,15 @@ protected String visitBinaryArithmetic(String name, String l, String r) {
     return l + " " + name + " " + r;
   }
 
+  @Deprecated
   protected String visitCast(String l, DataType dataType) {
     return "CAST(" + l + " AS " + dataType.typeName() + ")";
   }
 
+  protected String visitCast(String expr, DataType exprDataType, DataType targetDataType) {
+    return "CAST(" + expr + " AS " + targetDataType.typeName() + ")";
+  }
+
   protected String visitAnd(String name, String l, String r) {
     return "(" + l + ") " + name + " (" + r + ")";
   }

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/catalyst/util/V2ExpressionBuilder.scala b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/util/V2ExpressionBuilder.scala
@@ -95,7 +95,7 @@ class V2ExpressionBuilder(e: Expression, isPredicate: Boolean = false) {
       }
     case Cast(child, dataType, _, evalMode)
         if evalMode == EvalMode.ANSI || Cast.canUpCast(child.dataType, dataType) =>
-      generateExpression(child).map(v => new V2Cast(v, dataType))
+      generateExpression(child).map(v => new V2Cast(v, child.dataType, dataType))
     case AggregateExpression(aggregateFunction, Complete, isDistinct, None, _) =>
       generateAggregateFunc(aggregateFunction, isDistinct)
     case Abs(_, true) => generateExpressionWithName("ABS", expr, isPredicate)