Fixes and tests

Author: heyihong

python/pyspark/sql/connect/expressions.py

@@ -377,6 +377,52 @@ 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 literal.array.HasField("element_type"):
+                return ArrayType(
+                    proto_schema_to_pyspark_data_type(literal.array.element_type), True
+                )
+            element_type = None
+            if len(literal.array.elements) > 0:
+                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 at the moment.
+        # 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 +472,21 @@ 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"):
+            elementType = None
+            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])
-            else:
+                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])
+
+            if elementType is None:
                 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/connect/proto/expressions_pb2.py

@@ -477,7 +477,8 @@ class Expression(google.protobuf.message.Message):
             @property
             def element_type(self) -> pyspark.sql.connect.proto.types_pb2.DataType:
                 """(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 inferable, since spark 4.1+ supports
+                inferring the element type from the elements.
                 """
             @property
             def elements(
@@ -492,25 +493,14 @@ class Expression(google.protobuf.message.Message):
                 elements: collections.abc.Iterable[global___Expression.Literal] | None = ...,
             ) -> None: ...
             def HasField(
-                self,
-                field_name: typing_extensions.Literal[
-                    "_element_type", b"_element_type", "element_type", b"element_type"
-                ],
+                self, field_name: typing_extensions.Literal["element_type", b"element_type"]
             ) -> builtins.bool: ...
             def ClearField(
                 self,
                 field_name: typing_extensions.Literal[
-                    "_element_type",
-                    b"_element_type",
-                    "element_type",
-                    b"element_type",
-                    "elements",
-                    b"elements",
+                    "element_type", b"element_type", "elements", b"elements"
                 ],
             ) -> None: ...
-            def WhichOneof(
-                self, oneof_group: typing_extensions.Literal["_element_type", b"_element_type"]
-            ) -> typing_extensions.Literal["element_type"] | None: ...
 
         class Map(google.protobuf.message.Message):
             DESCRIPTOR: google.protobuf.descriptor.Descriptor
@@ -522,12 +512,14 @@ class Expression(google.protobuf.message.Message):
             @property
             def key_type(self) -> pyspark.sql.connect.proto.types_pb2.DataType:
                 """(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 inferable, since spark 4.1+ supports
+                inferring the key type from the keys
                 """
             @property
             def value_type(self) -> pyspark.sql.connect.proto.types_pb2.DataType:
                 """(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 inferable, since spark 4.1+ supports
+                inferring the value type from the values.
                 """
             @property
             def keys(
@@ -552,23 +544,12 @@ class Expression(google.protobuf.message.Message):
             def HasField(
                 self,
                 field_name: typing_extensions.Literal[
-                    "_key_type",
-                    b"_key_type",
-                    "_value_type",
-                    b"_value_type",
-                    "key_type",
-                    b"key_type",
-                    "value_type",
-                    b"value_type",
+                    "key_type", b"key_type", "value_type", b"value_type"
                 ],
             ) -> builtins.bool: ...
             def ClearField(
                 self,
                 field_name: typing_extensions.Literal[
-                    "_key_type",
-                    b"_key_type",
-                    "_value_type",
-                    b"_value_type",
                     "key_type",
                     b"key_type",
                     "keys",
@@ -579,14 +560,6 @@ class Expression(google.protobuf.message.Message):
                     b"values",
                 ],
             ) -> None: ...
-            @typing.overload
-            def WhichOneof(
-                self, oneof_group: typing_extensions.Literal["_key_type", b"_key_type"]
-            ) -> typing_extensions.Literal["key_type"] | None: ...
-            @typing.overload
-            def WhichOneof(
-                self, oneof_group: typing_extensions.Literal["_value_type", b"_value_type"]
-            ) -> typing_extensions.Literal["value_type"] | None: ...
 
         class Struct(google.protobuf.message.Message):
             DESCRIPTOR: google.protobuf.descriptor.Descriptor

python/pyspark/sql/connect/proto/expressions_pb2.pyi

@@ -979,21 +979,20 @@ def test_literal_expression_with_arrays(self):
         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)

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

@@ -216,18 +216,21 @@ 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.
-      optional DataType element_type = 1;
+      // empty or the element type is not inferable, since spark 4.1+ supports
+      // inferring the element type from the elements.
+      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
-      optional DataType key_type = 1;
+      // empty or the key type is not inferable, since spark 4.1+ supports
+      // inferring the key type from the keys
+      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.
-      optional DataType value_type = 2;
+      // empty or the value type is not inferable, since spark 4.1+ supports
+      // inferring the value type from the values.
+      DataType value_type = 2;
       repeated Literal keys = 3;
       repeated Literal values = 4;
     }