Fix

Author: maropu

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/InternalRow.scala

@@ -17,7 +17,6 @@
 
 package org.apache.spark.sql.catalyst
 
-import org.apache.spark.SparkException
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.util.{ArrayData, MapData}
 import org.apache.spark.sql.types._
@@ -172,11 +171,8 @@ object InternalRow {
     case DoubleType => (input, v) => input.setDouble(ordinal, v.asInstanceOf[Double])
     case DecimalType.Fixed(precision, _) =>
       (input, v) => input.setDecimal(ordinal, v.asInstanceOf[Decimal], precision)
-    case CalendarIntervalType | BinaryType | _: ArrayType | StringType | _: StructType |
-         _: MapType | _: ObjectType =>
-      (input, v) => input.update(ordinal, v)
     case udt: UserDefinedType[_] => getWriter(ordinal, udt.sqlType)
     case NullType => (input, _) => input.setNullAt(ordinal)
-    case _ => throw new SparkException(s"Unsupported data type $dt")
+    case _ => (input, v) => input.update(ordinal, v)
   }
 }

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

@@ -50,8 +50,11 @@ class InterpretedMutableProjection(expressions: Seq[Expression]) extends Mutable
 
   override def target(row: InternalRow): MutableProjection = {
     // If `mutableRow` is `UnsafeRow`, `MutableProjection` accepts fixed-length types only
-    assert(!row.isInstanceOf[UnsafeRow] ||
-      validExprs.forall { case (e, _) => UnsafeRow.isFixedLength(e.dataType) })
+    require(!row.isInstanceOf[UnsafeRow] ||
+      validExprs.forall { case (e, _) => UnsafeRow.isFixedLength(e.dataType) },
+      "MutableProjection cannot use UnsafeRow for output data types: " +
+        validExprs.map(_._1.dataType).filterNot(UnsafeRow.isFixedLength)
+          .map(_.catalogString).mkString(", "))
     mutableRow = row
     this
   }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MutableProjectionSuite.scala

@@ -20,36 +20,41 @@ package org.apache.spark.sql.catalyst.expressions
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.Row
 import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow}
+import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.types.CalendarInterval
 
 class MutableProjectionSuite extends SparkFunSuite with ExpressionEvalHelper {
 
-  private def createMutableProjection(dataTypes: Array[DataType]): MutableProjection = {
+  val fixedLengthTypes = Array[DataType](
+    BooleanType, ByteType, ShortType, IntegerType, LongType, FloatType, DoubleType,
+    DateType, TimestampType)
+
+  val variableLengthTypes = Array(
+    StringType, DecimalType.defaultConcreteType, CalendarIntervalType, BinaryType,
+    ArrayType(StringType), MapType(IntegerType, StringType),
+    StructType.fromDDL("a INT, b STRING"), ObjectType(classOf[java.lang.Integer]))
+
+  def createMutableProjection(dataTypes: Array[DataType]): MutableProjection = {
     MutableProjection.create(dataTypes.zipWithIndex.map(x => BoundReference(x._2, x._1, true)))
   }
 
   testBothCodegenAndInterpreted("fixed-length types") {
-    val fixedLengthTypes = Array[DataType](
-      BooleanType, ByteType, ShortType, IntegerType, LongType, FloatType, DoubleType,
-      DateType, TimestampType)
+    val inputRow = InternalRow.fromSeq(Seq(true, 3.toByte, 15.toShort, -83, 129L, 1.0f, 5.0, 1, 2L))
     val proj = createMutableProjection(fixedLengthTypes)
-    val inputRow = InternalRow.fromSeq(
-      Seq(false, 3.toByte, 15.toShort, -83, 129L, 1.0f, 5.0, 100, 200L))
     assert(proj(inputRow) === inputRow)
+  }
 
-    // Use UnsafeRow as buffer
+  testBothCodegenAndInterpreted("unsafe buffer") {
+    val inputRow = InternalRow.fromSeq(Seq(false, 1.toByte, 9.toShort, -18, 53L, 3.2f, 7.8, 4, 9L))
     val numBytes = UnsafeRow.calculateBitSetWidthInBytes(fixedLengthTypes.length)
     val unsafeBuffer = UnsafeRow.createFromByteArray(numBytes, fixedLengthTypes.length)
+    val proj = createMutableProjection(fixedLengthTypes)
     val projUnsafeRow = proj.target(unsafeBuffer)(inputRow)
-    assert(FromUnsafeProjection(fixedLengthTypes)(projUnsafeRow) === inputRow)
+    assert(FromUnsafeProjection.apply(fixedLengthTypes)(projUnsafeRow) === inputRow)
   }
 
   testBothCodegenAndInterpreted("variable-length types") {
-    val variableLengthTypes = Array(
-      StringType, DecimalType.defaultConcreteType, CalendarIntervalType, BinaryType,
-      ArrayType(StringType), MapType(IntegerType, StringType),
-      StructType.fromDDL("a INT, b STRING"), ObjectType(classOf[java.lang.Integer]))
     val proj = createMutableProjection(variableLengthTypes)
     val scalaValues = Seq("abc", BigDecimal(10), CalendarInterval.fromString("interval 1 day"),
       Array[Byte](1, 2), Array("123", "456"), Map(1 -> "a", 2 -> "b"), Row(1, "a"),
@@ -63,4 +68,14 @@ class MutableProjectionSuite extends SparkFunSuite with ExpressionEvalHelper {
       assert(toScala(projRow.get(index, dataType)) === toScala(inputRow.get(index, dataType)))
     }
   }
+
+  test("unsupported types for unsafe buffer") {
+    withSQLConf(SQLConf.CODEGEN_FACTORY_MODE.key -> CodegenObjectFactoryMode.NO_CODEGEN.toString) {
+      val proj = createMutableProjection(Array(StringType))
+      val errMsg = intercept[IllegalArgumentException] {
+        proj.target(new UnsafeRow(1))
+      }.getMessage
+      assert(errMsg.contains("MutableProjection cannot use UnsafeRow for output data types:"))
+    }
+  }
 }