[SPARK-55462][PYTHON] Support VariantType in convert_numpy

python/pyspark/sql/conversion.py

@@ -18,7 +18,7 @@
 import array
 import datetime
 import decimal
-from typing import TYPE_CHECKING, Any, Callable, List, Optional, Sequence, Union, overload
+from typing import TYPE_CHECKING, Any, Callable, List, Optional, Sequence, Union, cast, overload
 
 import pyspark
 from pyspark.errors import PySparkValueError
@@ -1357,6 +1357,7 @@ def _prefer_convert_numpy(
             TimestampType,
             TimestampNTZType,
             UserDefinedType,
+            VariantType,
         )
         if df_for_struct and isinstance(spark_type, StructType):
             return all(isinstance(f.dataType, supported_types) for f in spark_type.fields)
@@ -1489,13 +1490,20 @@ def convert_numpy(
                 if v is not None
                 else None
             )
+        elif isinstance(spark_type, VariantType):
+            series = arr.to_pandas(date_as_object=True)
+            series = series.apply(
+                cast(
+                    Callable,
+                    lambda v: VariantVal(v["value"], v["metadata"]) if v is not None else None,
+                )
+            )
         # elif isinstance(
         #     spark_type,
         #     (
         #         ArrayType,
         #         MapType,
         #         StructType,
-        #         VariantType,
         #         GeographyType,
         #         GeometryType,
         #     ),

python/pyspark/sql/tests/test_conversion.py

@@ -39,6 +39,8 @@
     StructField,
     StructType,
     UserDefinedType,
+    VariantType,
+    VariantVal,
 )
 from pyspark.testing.objects import ExamplePoint, ExamplePointUDT, PythonOnlyPoint, PythonOnlyUDT
 from pyspark.testing.utils import have_pyarrow, pyarrow_requirement_message
@@ -416,6 +418,40 @@ def test_udt_chunked_array(self):
         self.assertEqual(result.iloc[0], ExamplePoint(1.0, 2.0))
         self.assertEqual(result.iloc[1], ExamplePoint(3.0, 4.0))
 
+    def test_variant_convert_numpy(self):
+        import pyarrow as pa
+
+        variant_type = pa.struct(
+            [
+                pa.field("value", pa.binary(), nullable=False),
+                pa.field("metadata", pa.binary(), nullable=False, metadata={b"variant": b"true"}),
+            ]
+        )
+
+        # basic conversion with nulls
+        arr = pa.array(
+            [
+                {"value": b"\x01", "metadata": b"\x02"},
+                None,
+                {"value": b"\x03", "metadata": b"\x04"},
+            ],
+            type=variant_type,
+        )
+        result = ArrowArrayToPandasConversion.convert_numpy(arr, VariantType(), ser_name="v")
+        self.assertIsInstance(result.iloc[0], VariantVal)
+        self.assertEqual(result.iloc[0].value, b"\x01")
+        self.assertEqual(result.iloc[0].metadata, b"\x02")
+        self.assertIsNone(result.iloc[1])
+        self.assertEqual(result.iloc[2].value, b"\x03")
+        self.assertEqual(result.iloc[2].metadata, b"\x04")
+        self.assertEqual(result.name, "v")
+
+        # empty
+        result = ArrowArrayToPandasConversion.convert_numpy(
+            pa.array([], type=variant_type), VariantType()
+        )
+        self.assertEqual(len(result), 0)
+
 
 if __name__ == "__main__":
     from pyspark.testing import main