Fixes and tests

Author: heyihong

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/scala/org/apache/spark/sql/connect/common/LiteralValueProtoConverter.scala

@@ -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[_, _] =>

sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/LiteralExpressionProtoConverterSuite.scala

@@ -58,7 +58,35 @@ class LiteralExpressionProtoConverterSuite extends AnyFunSuite { // scalastyle:i
       }
   }
 
-  test("invalid array literal - empty array") {
+  test("element type of array literal is set for empty array") {
+    val literalProto =
+      LiteralValueProtoConverter.toLiteralProto(Array[Int](), ArrayType(IntegerType))
+    assert(literalProto.getArray.hasElementType)
+  }
+
+  test("element type of array literal is not set for non-empty array") {
+    val literalProto =
+      LiteralValueProtoConverter.toLiteralProto(Array(1, 2, 3), ArrayType(IntegerType))
+    assert(!literalProto.getArray.hasElementType)
+  }
+
+  test("key and value type of map literal is set for empty map") {
+    val literalProto = LiteralValueProtoConverter.toLiteralProto(
+      Map[Int, Int](),
+      MapType(IntegerType, IntegerType))
+    assert(literalProto.getMap.hasKeyType)
+    assert(literalProto.getMap.hasValueType)
+  }
+
+  test("key and value type of map literal is not set for non-empty map") {
+    val literalProto = LiteralValueProtoConverter.toLiteralProto(
+      Map(1 -> 2, 3 -> 4, 5 -> 6),
+      MapType(IntegerType, IntegerType))
+    assert(!literalProto.getMap.hasKeyType)
+    assert(!literalProto.getMap.hasValueType)
+  }
+
+  test("invalid array literal") {
     val literalProto = proto.Expression.Literal
       .newBuilder()
       .setArray(proto.Expression.Literal.Array.newBuilder())