Fixes and tests

Author: heyihong

python/pyspark/sql/connect/expressions.py

@@ -377,6 +377,54 @@ def _infer_type(cls, value: Any) -> DataType:
     def _from_value(cls, value: Any) -> "LiteralExpression":
         return LiteralExpression(value=value, dataType=LiteralExpression._infer_type(value))
 
+    @classmethod
+    def _infer_type_from_literal(cls, literal: "proto.Expression.Literal") -> Optional[DataType]:
+        if literal.HasField("null"):
+            return NullType()
+        elif literal.HasField("binary"):
+            return BinaryType()
+        elif literal.HasField("boolean"):
+            return BooleanType()
+        elif literal.HasField("byte"):
+            return ByteType()
+        elif literal.HasField("short"):
+            return ShortType()
+        elif literal.HasField("integer"):
+            return IntegerType()
+        elif literal.HasField("long"):
+            return LongType()
+        elif literal.HasField("float"):
+            return FloatType()
+        elif literal.HasField("double"):
+            return DoubleType()
+        elif literal.HasField("date"):
+            return DateType()
+        elif literal.HasField("timestamp"):
+            return TimestampType()
+        elif literal.HasField("timestamp_ntz"):
+            return TimestampNTZType()
+        elif literal.HasField("array"):
+            if len(literal.array.elements) == 0:
+                if literal.array.HasField("element_type"):
+                    return ArrayType(
+                        proto_schema_to_pyspark_data_type(literal.array.element_type), True
+                    )
+                raise PySparkTypeError(
+                    errorClass="CANNOT_INFER_ARRAY_ELEMENT_TYPE",
+                    messageParameters={},
+                )
+            element_type = LiteralExpression._infer_type_from_literal(literal.array.elements[0])
+            if element_type is None:
+                raise PySparkTypeError(
+                    errorClass="CANNOT_INFER_ARRAY_ELEMENT_TYPE",
+                    messageParameters={},
+                )
+            return ArrayType(element_type, True)
+        # Not all data types support inferring the data type from the literal.
+        # e.g. the type of DayTimeInterval contains extra information like start_field and
+        # end_field and cannot be inferred from the literal.
+        return None
+
     @classmethod
     def _to_value(
         cls, literal: "proto.Expression.Literal", dataType: Optional[DataType] = None
@@ -426,26 +474,19 @@ def _to_value(
             assert dataType is None or isinstance(dataType, DayTimeIntervalType)
             return DayTimeIntervalType().fromInternal(literal.day_time_interval)
         elif literal.HasField("array"):
-            elements = literal.array.elements
-            result = []
-            if dataType is not None:
-                assert isinstance(dataType, ArrayType)
-                elementType = dataType.elementType
-            elif literal.array.HasField("element_type"):
+            if literal.array.HasField("element_type"):
                 elementType = proto_schema_to_pyspark_data_type(literal.array.element_type)
-            elif len(elements) > 0:
-                result.append(LiteralExpression._to_value(elements[0], None))
-                elements = elements[1:]
-                elementType = LiteralExpression._infer_type(result[0])
+                if dataType is not None:
+                    assert isinstance(dataType, ArrayType)
+                    assert elementType == dataType.elementType
+            elif len(literal.array.elements) > 0:
+                elementType = LiteralExpression._infer_type_from_literal(literal.array.elements[0])
             else:
                 raise PySparkTypeError(
                     errorClass="CANNOT_INFER_ARRAY_ELEMENT_TYPE",
                     messageParameters={},
                 )
-
-            for element in elements:
-                result.append(LiteralExpression._to_value(element, elementType))
-            return result
+            return [LiteralExpression._to_value(v, elementType) for v in literal.array.elements]
 
         raise PySparkTypeError(
             errorClass="UNSUPPORTED_LITERAL",
@@ -490,14 +531,18 @@ def to_plan(self, session: "SparkConnectClient") -> "proto.Expression":
         elif isinstance(self._dataType, DayTimeIntervalType):
             expr.literal.day_time_interval = int(self._value)
         elif isinstance(self._dataType, ArrayType):
-            if len(self._value) == 0:
-                expr.literal.array.element_type.CopyFrom(
-                    pyspark_types_to_proto_types(self._dataType.elementType)
-                )
             for v in self._value:
                 expr.literal.array.elements.append(
                     LiteralExpression(v, self._dataType.elementType).to_plan(session).literal
                 )
+            if (
+                len(self._value) == 0
+                or LiteralExpression._infer_type_from_literal(expr.literal.array.elements[0])
+                is None
+            ):
+                expr.literal.array.element_type.CopyFrom(
+                    pyspark_types_to_proto_types(self._dataType.elementType)
+                )
         else:
             raise PySparkTypeError(
                 errorClass="UNSUPPORTED_DATA_TYPE",

python/pyspark/sql/tests/connect/test_connect_plan.py

@@ -972,28 +972,27 @@ def test_column_expressions(self):
 
     def test_literal_expression_with_arrays(self):
         l0 = LiteralExpression._from_value(["x", "y", "z"]).to_plan(None).literal
-        self.assertTrue(l0.array.element_type.HasField("string"))
+        self.assertFalse(l0.array.element_type.HasField("string"))
         self.assertEqual(len(l0.array.elements), 3)
         self.assertEqual(l0.array.elements[0].string, "x")
         self.assertEqual(l0.array.elements[1].string, "y")
         self.assertEqual(l0.array.elements[2].string, "z")
 
         l1 = LiteralExpression._from_value([3, -3]).to_plan(None).literal
-        self.assertTrue(l1.array.element_type.HasField("integer"))
+        self.assertFalse(l1.array.element_type.HasField("integer"))
         self.assertEqual(len(l1.array.elements), 2)
         self.assertEqual(l1.array.elements[0].integer, 3)
         self.assertEqual(l1.array.elements[1].integer, -3)
 
         l2 = LiteralExpression._from_value([float("nan"), -3.0, 0.0]).to_plan(None).literal
-        self.assertTrue(l2.array.element_type.HasField("double"))
+        self.assertFalse(l2.array.element_type.HasField("double"))
         self.assertEqual(len(l2.array.elements), 3)
         self.assertTrue(math.isnan(l2.array.elements[0].double))
         self.assertEqual(l2.array.elements[1].double, -3.0)
         self.assertEqual(l2.array.elements[2].double, 0.0)
 
         l3 = LiteralExpression._from_value([[3, 4], [5, 6, 7]]).to_plan(None).literal
-        self.assertTrue(l3.array.element_type.HasField("array"))
-        self.assertTrue(l3.array.element_type.array.element_type.HasField("integer"))
+        self.assertFalse(l3.array.element_type.HasField("array"))
         self.assertEqual(len(l3.array.elements), 2)
         self.assertEqual(len(l3.array.elements[0].array.elements), 2)
         self.assertEqual(len(l3.array.elements[1].array.elements), 3)
@@ -1003,8 +1002,7 @@ def test_literal_expression_with_arrays(self):
             .to_plan(None)
             .literal
         )
-        self.assertTrue(l4.array.element_type.HasField("array"))
-        self.assertTrue(l4.array.element_type.array.element_type.HasField("double"))
+        self.assertFalse(l4.array.element_type.HasField("array"))
         self.assertEqual(len(l4.array.elements), 3)
         self.assertEqual(len(l4.array.elements[0].array.elements), 2)
         self.assertEqual(len(l4.array.elements[1].array.elements), 2)
@@ -1033,6 +1031,8 @@ def test_literal_to_any_conversion(self):
         ]:
             lit = LiteralExpression._from_value(value)
             proto_lit = lit.to_plan(None).literal
+            if proto_lit.HasField("array"):
+                self.assertFalse(proto_lit.array.HasField("element_type"))
             value2 = LiteralExpression._to_value(proto_lit)
             self.assertEqual(value, value2)
 

sql/connect/common/src/main/protobuf/spark/connect/expressions.proto

@@ -216,17 +216,20 @@ message Expression {
 
     message Array {
       // (Optional) The element type of the array. Only need to set this when the elements are
-      // empty, since spark 4.1+ supports inferring the element type from the elements.
+      // empty or the element type is not inferrable, since spark 4.1+ supports
+      // inferring the element type from the elements.
       optional DataType element_type = 1;
       repeated Literal elements = 2;
     }
 
     message Map {
       // (Optional) The key type of the map. Only need to set this when the keys are
-      // empty, since spark 4.1+ supports inferring the key type from the keys
+      // empty or the key type is not inferrable, since spark 4.1+ supports
+      // inferring the key type from the keys
       optional DataType key_type = 1;
       // (Optional) The value type of the map. Only need to set this when the values are
-      // empty, since spark 4.1+ supports inferring the value type from the values.
+      // empty or the value type is not inferrable, since spark 4.1+ supports
+      // inferring the value type from the values.
       optional DataType value_type = 2;
       repeated Literal keys = 3;
       repeated Literal values = 4;

sql/connect/common/src/main/scala/org/apache/spark/sql/connect/common/LiteralValueProtoConverter.scala

@@ -64,7 +64,7 @@ object LiteralValueProtoConverter {
     def arrayBuilder(array: Array[_]) = {
       val ab = builder.getArrayBuilder
       array.foreach(x => ab.addElements(toLiteralProto(x)))
-      if (ab.getElementsCount == 0) {
+      if (ab.getElementsCount == 0 || getInferredDataType(ab.getElementsList.get(0)).isEmpty) {
         ab.setElementType(toConnectProtoType(toDataType(array.getClass.getComponentType)))
       }
       ab
@@ -129,7 +129,7 @@ object LiteralValueProtoConverter {
           throw new IllegalArgumentException(s"literal $other not supported (yet).")
       }
 
-      if (ab.getElementsCount == 0) {
+      if (ab.getElementsCount == 0 || getInferredDataType(ab.getElementsList.get(0)).isEmpty) {
         ab.setElementType(toConnectProtoType(elementType))
       }
 
@@ -138,8 +138,6 @@ object LiteralValueProtoConverter {
 
     def mapBuilder(scalaValue: Any, keyType: DataType, valueType: DataType) = {
       val mb = builder.getMapBuilder
-        .setKeyType(toConnectProtoType(keyType))
-        .setValueType(toConnectProtoType(valueType))
 
       scalaValue match {
         case map: scala.collection.Map[_, _] =>
@@ -151,11 +149,11 @@ object LiteralValueProtoConverter {
           throw new IllegalArgumentException(s"literal $other not supported (yet).")
       }
 
-      if (mb.getKeysCount == 0) {
+      if (mb.getKeysCount == 0 || getInferredDataType(mb.getKeysList.get(0)).isEmpty) {
         mb.setKeyType(toConnectProtoType(keyType))
       }
 
-      if (mb.getValuesCount == 0) {
+      if (mb.getValuesCount == 0 || getInferredDataType(mb.getValuesList.get(0)).isEmpty) {
         mb.setValueType(toConnectProtoType(valueType))
       }
 
@@ -370,7 +368,11 @@ object LiteralValueProtoConverter {
     }
   }
 
-  private def getOuterDataType(literal: proto.Expression.Literal): proto.DataType = {
+  private def getInferredDataType(literal: proto.Expression.Literal): Option[proto.DataType] = {
+    if (literal.hasNull) {
+      return Some(literal.getNull)
+    }
+
     val builder = proto.DataType.newBuilder()
     literal.getLiteralTypeCase match {
       case proto.Expression.Literal.LiteralTypeCase.BINARY =>
@@ -389,22 +391,14 @@ object LiteralValueProtoConverter {
         builder.setFloat(proto.DataType.Float.newBuilder.build())
       case proto.Expression.Literal.LiteralTypeCase.DOUBLE =>
         builder.setDouble(proto.DataType.Double.newBuilder.build())
-      case proto.Expression.Literal.LiteralTypeCase.STRING =>
-        builder.setString(proto.DataType.String.newBuilder.build())
       case proto.Expression.Literal.LiteralTypeCase.DATE =>
         builder.setDate(proto.DataType.Date.newBuilder.build())
       case proto.Expression.Literal.LiteralTypeCase.TIMESTAMP =>
         builder.setTimestamp(proto.DataType.Timestamp.newBuilder.build())
       case proto.Expression.Literal.LiteralTypeCase.TIMESTAMP_NTZ =>
         builder.setTimestampNtz(proto.DataType.TimestampNTZ.newBuilder.build())
-      case proto.Expression.Literal.LiteralTypeCase.YEAR_MONTH_INTERVAL =>
-        builder.setYearMonthInterval(proto.DataType.YearMonthInterval.newBuilder.build())
-      case proto.Expression.Literal.LiteralTypeCase.DECIMAL =>
-        builder.setDecimal(proto.DataType.Decimal.newBuilder.build())
       case proto.Expression.Literal.LiteralTypeCase.CALENDAR_INTERVAL =>
         builder.setCalendarInterval(proto.DataType.CalendarInterval.newBuilder.build())
-      case proto.Expression.Literal.LiteralTypeCase.DAY_TIME_INTERVAL =>
-        builder.setDayTimeInterval(proto.DataType.DayTimeInterval.newBuilder.build())
       case proto.Expression.Literal.LiteralTypeCase.ARRAY =>
         // Element type will be inferred from the elements in the array.
         builder.setArray(proto.DataType.Array.newBuilder.build())
@@ -413,9 +407,19 @@ object LiteralValueProtoConverter {
         builder.setMap(proto.DataType.Map.newBuilder.build())
       case proto.Expression.Literal.LiteralTypeCase.STRUCT =>
         builder.setStruct(literal.getStruct.getStructType.getStruct)
-      case _ => throw InvalidPlanInput(s"Unsupported Literal Type: ${literal.getLiteralTypeCase}")
+      case _ =>
+        // Not all data types support inferring the data type from the literal.
+        // e.g. the type of DayTimeInterval contains extra information like start_field and
+        // end_field and cannot be inferred from the literal.
+        return None
+    }
+    Some(builder.build())
+  }
+
+  private def getInferredDataTypeOrThrow(literal: proto.Expression.Literal): proto.DataType = {
+    getInferredDataType(literal).getOrElse {
+      throw InvalidPlanInput(s"Unsupported Literal Type: ${literal.getLiteralTypeCase}")
     }
-    builder.build()
   }
 
   def toCatalystArray(
@@ -437,7 +441,7 @@ object LiteralValueProtoConverter {
         protoArrayType(array.getElementType)
       } else if (iter.hasNext) {
         val firstElement = iter.next()
-        val outerElementType = getOuterDataType(firstElement)
+        val outerElementType = getInferredDataTypeOrThrow(firstElement)
         val (elem, inferredElementType) =
           getConverter(outerElementType, inferDataType = true)(firstElement) match {
             case LiteralValueWithDataType(elem, dataType) => (elem, dataType)
@@ -479,15 +483,23 @@ object LiteralValueProtoConverter {
         protoMapType(map.getKeyType, map.getValueType)
       } else if (iter.hasNext) {
         val (key, value) = iter.next()
-        val outerKeyType = getOuterDataType(key)
+        val (outerKeyType, inferKeyType) = if (map.hasKeyType) {
+          (map.getKeyType, false)
+        } else {
+          (getInferredDataTypeOrThrow(key), true)
+        }
         val (catalystKey, inferredKeyType) =
-          getConverter(outerKeyType, inferDataType = true)(key) match {
+          getConverter(outerKeyType, inferDataType = inferKeyType)(key) match {
             case LiteralValueWithDataType(catalystKey, dataType) => (catalystKey, dataType)
             case catalystKey => (catalystKey, outerKeyType)
           }
-        val outerValueType = getOuterDataType(value)
+        val (outerValueType, inferValueType) = if (map.hasValueType) {
+          (map.getValueType, false)
+        } else {
+          (getInferredDataTypeOrThrow(value), true)
+        }
         val (catalystValue, inferredValueType) =
-          getConverter(outerValueType, inferDataType = true)(value) match {
+          getConverter(outerValueType, inferDataType = inferValueType)(value) match {
             case LiteralValueWithDataType(catalystValue, dataType) => (catalystValue, dataType)
             case catalystValue => (catalystValue, outerValueType)
           }

sql/connect/common/src/test/resources/query-tests/queries/function_lit.json

@@ -358,10 +358,6 @@
     }, {
       "literal": {
         "array": {
-          "elementType": {
-            "integer": {
-            }
-          },
           "elements": [{
             "integer": 8
           }, {