[SPARK-52449] Make datatypes for Expression.Literal.Map/Expression.Literal.Array optional

Author: heyihong

python/pyspark/sql/connect/expressions.py

@@ -29,6 +29,7 @@
     Optional,
 )
 
+import os
 import json
 import decimal
 import datetime
@@ -84,6 +85,9 @@
     from pyspark.sql.connect.window import WindowSpec
     from pyspark.sql.connect.plan import LogicalPlan
 
+_optional_data_types_for_map_and_array_literals_enabled = (
+    os.getenv("CONNECT_OPTIONAL_DATATYPE_FOR_MAP_AND_ARRAY_LITERALS_ENABLED", "false") == "true"
+)
 
 class Expression:
     """
@@ -426,11 +430,26 @@ def _to_value(
             assert dataType is None or isinstance(dataType, DayTimeIntervalType)
             return DayTimeIntervalType().fromInternal(literal.day_time_interval)
         elif literal.HasField("array"):
-            elementType = proto_schema_to_pyspark_data_type(literal.array.element_type)
+            elements = literal.array.elements
+            result = []
             if dataType is not None:
                 assert isinstance(dataType, ArrayType)
-                assert elementType == dataType.elementType
-            return [LiteralExpression._to_value(v, elementType) for v in literal.array.elements]
+                elementType = dataType.elementType
+            elif 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:
+                raise PySparkTypeError(
+                    errorClass="CANNOT_INFER_ARRAY_ELEMENT_TYPE",
+                    messageParameters={},
+                )
+
+            for element in elements:
+                result.append(LiteralExpression._to_value(element, elementType))
+            return result
 
         raise PySparkTypeError(
             errorClass="UNSUPPORTED_LITERAL",
@@ -475,11 +494,11 @@ 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):
-            element_type = self._dataType.elementType
-            expr.literal.array.element_type.CopyFrom(pyspark_types_to_proto_types(element_type))
+            if not _optional_data_types_for_map_and_array_literals_enabled or 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, element_type).to_plan(session).literal
+                    LiteralExpression(v, self._dataType.elementType).to_plan(session).literal
                 )
         else:
             raise PySparkTypeError(

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

@@ -475,7 +475,10 @@ class Expression(google.protobuf.message.Message):
             ELEMENT_TYPE_FIELD_NUMBER: builtins.int
             ELEMENTS_FIELD_NUMBER: builtins.int
             @property
-            def element_type(self) -> pyspark.sql.connect.proto.types_pb2.DataType: ...
+            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.
+                """
             @property
             def elements(
                 self,
@@ -489,14 +492,25 @@ 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"]
+                self,
+                field_name: typing_extensions.Literal[
+                    "_element_type", b"_element_type", "element_type", b"element_type"
+                ],
             ) -> builtins.bool: ...
             def ClearField(
                 self,
                 field_name: typing_extensions.Literal[
-                    "element_type", b"element_type", "elements", b"elements"
+                    "_element_type",
+                    b"_element_type",
+                    "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
@@ -506,9 +520,15 @@ class Expression(google.protobuf.message.Message):
             KEYS_FIELD_NUMBER: builtins.int
             VALUES_FIELD_NUMBER: builtins.int
             @property
-            def key_type(self) -> pyspark.sql.connect.proto.types_pb2.DataType: ...
+            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
+                """
             @property
-            def value_type(self) -> pyspark.sql.connect.proto.types_pb2.DataType: ...
+            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.
+                """
             @property
             def keys(
                 self,
@@ -532,12 +552,23 @@ 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",
@@ -548,6 +579,14 @@ 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

@@ -19,6 +19,7 @@
 import datetime
 import decimal
 import math
+from contextlib import contextmanager
 
 from pyspark.testing.connectutils import (
     PlanOnlyTestFixture,
@@ -1010,54 +1011,66 @@ def test_literal_expression_with_arrays(self):
         self.assertEqual(len(l4.array.elements[1].array.elements), 2)
         self.assertEqual(len(l4.array.elements[2].array.elements), 0)
 
-    def test_literal_to_any_conversion(self):
-        for value in [
-            b"binary\0\0asas",
-            True,
-            False,
-            0,
-            12,
-            -1,
-            0.0,
-            1.234567,
-            decimal.Decimal(0.0),
-            decimal.Decimal(1.234567),
-            "sss",
-            datetime.date(2022, 12, 13),
-            datetime.datetime.now(),
-            datetime.timedelta(1, 2, 3),
-            [1, 2, 3, 4, 5, 6],
-            [-1.0, 2.0, 3.0],
-            ["x", "y", "z"],
-            [[1.0, 2.0, 3.0], [4.0, 5.0], [6.0]],
-        ]:
-            lit = LiteralExpression._from_value(value)
-            proto_lit = lit.to_plan(None).literal
-            value2 = LiteralExpression._to_value(proto_lit)
-            self.assertEqual(value, value2)
-
-        with self.assertRaises(AssertionError):
-            lit = LiteralExpression._from_value(1.234567)
-            proto_lit = lit.to_plan(None).literal
-            LiteralExpression._to_value(proto_lit, StringType())
-
-        with self.assertRaises(AssertionError):
-            lit = LiteralExpression._from_value("1.234567")
-            proto_lit = lit.to_plan(None).literal
-            LiteralExpression._to_value(proto_lit, DoubleType())
-
-        with self.assertRaises(AssertionError):
-            # build a array<string> proto literal, but with incorrect elements
-            proto_lit = proto.Expression().literal
-            proto_lit.array.element_type.CopyFrom(pyspark_types_to_proto_types(StringType()))
-            proto_lit.array.elements.append(
-                LiteralExpression("string", StringType()).to_plan(None).literal
-            )
-            proto_lit.array.elements.append(
-                LiteralExpression(1.234, DoubleType()).to_plan(None).literal
-            )
+    @contextmanager
+    def _optional_data_types_enabled(self, enabled: bool):
+        from pyspark.sql.connect.expressions import _optional_data_types_for_map_and_array_literals_enabled
+        previous_value = _optional_data_types_for_map_and_array_literals_enabled
+        try:
+            _optional_data_types_for_map_and_array_literals_enabled = enabled
+            yield
+        finally:
+            _optional_data_types_for_map_and_array_literals_enabled = previous_value
 
-            LiteralExpression._to_value(proto_lit, DoubleType)
+    def test_literal_to_any_conversion(self):
+        for optional_data_types_enabled in [True, False]:
+            with self._optional_data_types_enabled(optional_data_types_enabled):
+                for value in [
+                    b"binary\0\0asas",
+                    True,
+                    False,
+                    0,
+                    12,
+                    -1,
+                    0.0,
+                    1.234567,
+                    decimal.Decimal(0.0),
+                    decimal.Decimal(1.234567),
+                    "sss",
+                    datetime.date(2022, 12, 13),
+                    datetime.datetime.now(),
+                    datetime.timedelta(1, 2, 3),
+                    [1, 2, 3, 4, 5, 6],
+                    [-1.0, 2.0, 3.0],
+                    ["x", "y", "z"],
+                    [[1.0, 2.0, 3.0], [4.0, 5.0], [6.0]],
+                ]:
+                    lit = LiteralExpression._from_value(value)
+                    proto_lit = lit.to_plan(None).literal
+                    value2 = LiteralExpression._to_value(proto_lit)
+                    self.assertEqual(value, value2)
+
+                with self.assertRaises(AssertionError):
+                    lit = LiteralExpression._from_value(1.234567)
+                    proto_lit = lit.to_plan(None).literal
+                    LiteralExpression._to_value(proto_lit, StringType())
+
+                with self.assertRaises(AssertionError):
+                    lit = LiteralExpression._from_value("1.234567")
+                    proto_lit = lit.to_plan(None).literal
+                    LiteralExpression._to_value(proto_lit, DoubleType())
+
+                with self.assertRaises(AssertionError):
+                    # build a array<string> proto literal, but with incorrect elements
+                    proto_lit = proto.Expression().literal
+                    proto_lit.array.element_type.CopyFrom(pyspark_types_to_proto_types(StringType()))
+                    proto_lit.array.elements.append(
+                        LiteralExpression("string", StringType()).to_plan(None).literal
+                    )
+                    proto_lit.array.elements.append(
+                        LiteralExpression(1.234, DoubleType()).to_plan(None).literal
+                    )
+
+                    LiteralExpression._to_value(proto_lit, DoubleType)
 
 
 if __name__ == "__main__":

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

@@ -215,13 +215,19 @@ message Expression {
     }
 
     message Array {
-      DataType element_type = 1;
+      // (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;
       repeated Literal elements = 2;
     }
 
     message Map {
-      DataType key_type = 1;
-      DataType value_type = 2;
+      // (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;
+      // (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;
       repeated Literal keys = 3;
       repeated Literal values = 4;
     }