temp

Author: ilicmarkodb

python/pyspark/sql/connect/types.py

@@ -129,6 +129,10 @@ def pyspark_types_to_proto_types(data_type: DataType) -> pb2.DataType:
     ret = pb2.DataType()
     if isinstance(data_type, NullType):
         ret.null.CopyFrom(pb2.DataType.NULL())
+    elif isinstance(data_type, CharType):
+        ret.char.length = data_type.length
+    elif isinstance(data_type, VarcharType):
+        ret.varchar.length = data_type.length
     elif isinstance(data_type, StringType):
         ret.string.collation = data_type.collation
     elif isinstance(data_type, BooleanType):

python/pyspark/sql/tests/test_udf.py

@@ -1402,6 +1402,36 @@ def my_udf(input_val):
             result_type = df_result.schema["result"].dataType
             self.assertEqual(result_type, StringType("fr"))
 
+    def test_udf_with_char_varchar_return_type(self):
+        char_type = "char(10)"
+        varchar_type = "varchar(8)"
+        array_with_char_type = "array<char(5)>"
+        array_with_varchar_type = "array<varchar(12)>"
+        map_type = f"map<{char_type}, {varchar_type}>"
+        struct_type = f"struct<f1: {char_type}, f2: {varchar_type}>"
+
+        return_types = [
+            char_type,
+            varchar_type,
+            array_with_char_type,
+            array_with_varchar_type,
+            map_type,
+            struct_type,
+            f"struct<f1: {array_with_char_type}, f2: {array_with_varchar_type}, f3: {map_type}>",
+            f"map<{array_with_char_type}, {array_with_varchar_type}>",
+            f"array<{struct_type}>"
+        ]
+
+        for return_type in return_types:
+            with self.assertRaisesRegex(
+                AnalysisException,
+                "(Please use a different output data type for your UDF or DataFrame|"
+                "Invalid return type with Arrow-optimized Python UDF)"
+            ):
+                @udf(return_type)
+                def my_udf():
+                    return return_type
+
 
 class UDFTests(BaseUDFTestsMixin, ReusedSQLTestCase):
     @classmethod