add type cast if the real type is different but compatible with encoder schema

Author: cloud-fan

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala

@@ -28,9 +28,10 @@ import org.apache.spark.sql.catalyst.analysis.{SimpleAnalyzer, UnresolvedExtract
 import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, Project}
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.codegen.{GenerateSafeProjection, GenerateUnsafeProjection}
+import org.apache.spark.sql.catalyst.optimizer.SimplifyCasts
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.ScalaReflection
-import org.apache.spark.sql.types.{StructField, ObjectType, StructType}
+import org.apache.spark.sql.types.{DataType, StructField, ObjectType, StructType}
 
 /**
  * A factory for constructing encoders that convert objects and primitives to and from the
@@ -210,26 +211,71 @@ case class ExpressionEncoder[T](
     })
   }
 
+  private def handleStruct(input: Expression, s: StructType): Expression = {
+    assert(input.isInstanceOf[NewInstance] || input.isInstanceOf[CreateExternalRow])
+    val children = input.children
+    assert(children.length == s.length)
+
+    val newChildren = children.zip(s.map(_.dataType)).map {
+      case (child, dt) => typeCast(child, dt)
+    }
+
+    input.withNewChildren(newChildren)
+  }
+
+  private def typeCast(input: Expression, expectedType: DataType): Expression = expectedType match {
+    case s: StructType =>
+      var continue = true
+      input transformDown {
+        case c: CreateExternalRow if continue =>
+          continue = false
+          handleStruct(c, s)
+        case n: NewInstance if continue =>
+          continue = false
+          handleStruct(n, s)
+      }
+
+    case _ =>
+      var continue = true
+      input transformDown {
+        case u: UnresolvedExtractValue if continue =>
+          continue = false
+          Cast(u, expectedType)
+        case g: GetInternalRowField if continue =>
+          continue = false
+          Cast(g, expectedType)
+        case u: UnresolvedAttribute if continue =>
+          continue = false
+          Cast(u, expectedType)
+        case a: AttributeReference if continue =>
+          continue = false
+          Cast(a, expectedType)
+      }
+  }
+
   /**
    * Returns a new copy of this encoder, where the expressions used by `fromRow` are resolved to the
    * given schema.
    */
   def resolve(
-      schema: Seq[Attribute],
+      attrs: Seq[Attribute],
       outerScopes: ConcurrentMap[String, AnyRef]): ExpressionEncoder[T] = {
-    val positionToAttribute = AttributeMap.toIndex(schema)
-    val unbound = fromRowExpression transform {
+    val positionToAttribute = AttributeMap.toIndex(attrs)
+    val unbound = fromRowExpression transformUp {
       case b: BoundReference => positionToAttribute(b.ordinal)
     }
 
-    val plan = Project(Alias(unbound, "")() :: Nil, LocalRelation(schema))
+    val withTypeCast = typeCast(unbound, if (flat) schema.head.dataType else schema)
+
+    val plan = Project(Alias(withTypeCast, "")() :: Nil, LocalRelation(attrs))
     val analyzedPlan = SimpleAnalyzer.execute(plan)
+    val optimizedPlan = SimplifyCasts(analyzedPlan)
 
     // In order to construct instances of inner classes (for example those declared in a REPL cell),
     // we need an instance of the outer scope.  This rule substitues those outer objects into
     // expressions that are missing them by looking up the name in the SQLContexts `outerScopes`
     // registry.
-    copy(fromRowExpression = analyzedPlan.expressions.head.children.head transform {
+    copy(fromRowExpression = optimizedPlan.expressions.head.children.head transform {
       case n: NewInstance if n.outerPointer.isEmpty && n.cls.isMemberClass =>
         val outer = outerScopes.get(n.cls.getDeclaringClass.getName)
         if (outer == null) {

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala

@@ -126,7 +126,7 @@ case class CreateStruct(children: Seq[Expression]) extends Expression {
 case class CreateNamedStruct(children: Seq[Expression]) extends Expression {
 
   /**
-   * Returns Aliased [[Expressions]] that could be used to construct a flattened version of this
+   * Returns Aliased [[Expression]]s that could be used to construct a flattened version of this
    * StructType.
    */
   def flatten: Seq[NamedExpression] = valExprs.zip(names).map {

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/EncoderResolveSuite.scala

@@ -0,0 +1,121 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.encoders
+
+import org.apache.spark.sql.catalyst.dsl.expressions._
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.plans.PlanTest
+import org.apache.spark.sql.types._
+
+case class StringLongClass(a: String, b: Long)
+
+case class ComplexClass(a: Long, b: StringLongClass)
+
+class EncoderResolveSuite extends PlanTest {
+  test("real type doesn't match encoder schema but they are compatible: product") {
+    val encoder = ExpressionEncoder[StringLongClass]
+    val cls = classOf[StringLongClass]
+
+    var attrs = Seq('a.string, 'b.int)
+    var fromRowExpr: Expression = encoder.resolve(attrs, null).fromRowExpression
+    var expected: Expression = NewInstance(
+      cls,
+      toExternalString('a.string) :: 'b.int.cast(LongType) :: Nil,
+      false,
+      ObjectType(cls))
+    compareExpressions(fromRowExpr, expected)
+
+    attrs = Seq('a.int, 'b.long)
+    fromRowExpr = encoder.resolve(attrs, null).fromRowExpression
+    expected = NewInstance(
+      cls,
+      toExternalString('a.int.cast(StringType)) :: 'b.long :: Nil,
+      false,
+      ObjectType(cls))
+    compareExpressions(fromRowExpr, expected)
+  }
+
+  test("real type doesn't match encoder schema but they are compatible: nested product") {
+    val encoder = ExpressionEncoder[ComplexClass]
+    val innerCls = classOf[StringLongClass]
+    val cls = classOf[ComplexClass]
+
+    val structType = new StructType().add("a", IntegerType).add("b", LongType)
+    val attrs = Seq('a.int, 'b.struct(structType))
+    val fromRowExpr: Expression = encoder.resolve(attrs, null).fromRowExpression
+    val expected: Expression = NewInstance(
+      cls,
+      Seq(
+        'a.int.cast(LongType),
+        If(
+          'b.struct(structType).isNull,
+          Literal.create(null, ObjectType(innerCls)),
+          NewInstance(
+            innerCls,
+            Seq(
+              toExternalString(GetStructField(
+                'b.struct(structType),
+                structType(0),
+                0).cast(StringType)),
+              GetStructField(
+                'b.struct(structType),
+                structType(1),
+                1)),
+            false,
+            ObjectType(innerCls))
+        )),
+      false,
+      ObjectType(cls))
+    compareExpressions(fromRowExpr, expected)
+  }
+
+  test("real type doesn't match encoder schema but they are compatible: tupled encoder") {
+    val encoder = ExpressionEncoder.tuple(
+      ExpressionEncoder[StringLongClass],
+      ExpressionEncoder[Long])
+    val cls = classOf[StringLongClass]
+
+    val structType = new StructType().add("a", StringType).add("b", ByteType, false)
+    val attrs = Seq('a.struct(structType), 'b.int)
+    val fromRowExpr: Expression = encoder.resolve(attrs, null).fromRowExpression
+    val expected: Expression = NewInstance(
+      classOf[Tuple2[_, _]],
+      Seq(
+        NewInstance(
+          cls,
+          Seq(
+            toExternalString(GetStructField(
+              'a.struct(structType),
+              structType(0),
+              0)),
+            GetStructField(
+              'a.struct(structType),
+              structType(1),
+              1).cast(LongType)),
+          false,
+          ObjectType(cls)),
+        'b.int.cast(LongType)),
+      false,
+      ObjectType(classOf[Tuple2[_, _]]))
+    compareExpressions(fromRowExpr, expected)
+  }
+
+  private def toExternalString(e: Expression): Expression = {
+    Invoke(e, "toString", ObjectType(classOf[String]), Nil)
+  }
+}

sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala

@@ -384,7 +384,7 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
       Seq((JavaData(1), 1L), (JavaData(2), 1L)))
   }
 
-  ignore("Java encoder self join") {
+  test("Java encoder self join") {
     implicit val kryoEncoder = Encoders.javaSerialization[JavaData]
     val ds = Seq(JavaData(1), JavaData(2)).toDS()
     assert(ds.joinWith(ds, lit(true)).collect().toSet ==
@@ -394,6 +394,11 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
         (JavaData(2), JavaData(1)),
         (JavaData(2), JavaData(2))))
   }
+
+  test("change encoder with compatible schema") {
+    val ds = Seq(2 -> 2.toByte, 3 -> 3.toByte).toDF("a", "b").as[ClassData]
+    assert(ds.collect().toSeq == Seq(ClassData("2", 2), ClassData("3", 3)))
+  }
 }