snowpark-connect 0.23.0__py3-none-any.whl → 0.25.0__py3-none-any.whl

snowflake/snowpark_connect/includes/python/pyspark/sql/tests/test_udtf.py

@@ -1,1456 +0,0 @@
-#
-# Licensed to the Apache Software Foundation (ASF) under one or more
-# contributor license agreements.  See the NOTICE file distributed with
-# this work for additional information regarding copyright ownership.
-# The ASF licenses this file to You under the Apache License, Version 2.0
-# (the "License"); you may not use this file except in compliance with
-# the License.  You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-import os
-import tempfile
-import unittest
-from typing import Iterator
-
-from py4j.protocol import Py4JJavaError
-
-from pyspark.errors import (
-    PySparkAttributeError,
-    PythonException,
-    PySparkTypeError,
-    AnalysisException,
-    PySparkRuntimeError,
-)
-from pyspark.rdd import PythonEvalType
-from pyspark.sql.functions import lit, udf, udtf
-from pyspark.sql.types import IntegerType, MapType, Row, StringType, StructType
-from pyspark.testing import assertDataFrameEqual
-from pyspark.testing.sqlutils import (
-    have_pandas,
-    have_pyarrow,
-    pandas_requirement_message,
-    pyarrow_requirement_message,
-    ReusedSQLTestCase,
-)
-
-
-class BaseUDTFTestsMixin:
-    def test_simple_udtf(self):
-        class TestUDTF:
-            def eval(self):
-                yield "hello", "world"
-
-        func = udtf(TestUDTF, returnType="c1: string, c2: string")
-        rows = func().collect()
-        self.assertEqual(rows, [Row(c1="hello", c2="world")])
-
-    def test_udtf_yield_single_row_col(self):
-        class TestUDTF:
-            def eval(self, a: int):
-                yield a,
-
-        func = udtf(TestUDTF, returnType="a: int")
-        rows = func(lit(1)).collect()
-        self.assertEqual(rows, [Row(a=1)])
-
-    def test_udtf_yield_multi_cols(self):
-        class TestUDTF:
-            def eval(self, a: int):
-                yield a, a + 1
-
-        func = udtf(TestUDTF, returnType="a: int, b: int")
-        rows = func(lit(1)).collect()
-        self.assertEqual(rows, [Row(a=1, b=2)])
-
-    def test_udtf_yield_multi_rows(self):
-        class TestUDTF:
-            def eval(self, a: int):
-                yield a,
-                yield a + 1,
-
-        func = udtf(TestUDTF, returnType="a: int")
-        rows = func(lit(1)).collect()
-        self.assertEqual(rows, [Row(a=1), Row(a=2)])
-
-    def test_udtf_yield_multi_row_col(self):
-        class TestUDTF:
-            def eval(self, a: int, b: int):
-                yield a, b, a + b
-                yield a, b, a - b
-                yield a, b, b - a
-
-        func = udtf(TestUDTF, returnType="a: int, b: int, c: int")
-        rows = func(lit(1), lit(2)).collect()
-        self.assertEqual(rows, [Row(a=1, b=2, c=3), Row(a=1, b=2, c=-1), Row(a=1, b=2, c=1)])
-
-    def test_udtf_decorator(self):
-        @udtf(returnType="a: int, b: int")
-        class TestUDTF:
-            def eval(self, a: int):
-                yield a, a + 1
-
-        rows = TestUDTF(lit(1)).collect()
-        self.assertEqual(rows, [Row(a=1, b=2)])
-
-    def test_udtf_registration(self):
-        class TestUDTF:
-            def eval(self, a: int, b: int):
-                yield a, b, a + b
-                yield a, b, a - b
-                yield a, b, b - a
-
-        func = udtf(TestUDTF, returnType="a: int, b: int, c: int")
-        self.spark.udtf.register("testUDTF", func)
-        df = self.spark.sql("SELECT * FROM testUDTF(1, 2)")
-        self.assertEqual(
-            df.collect(), [Row(a=1, b=2, c=3), Row(a=1, b=2, c=-1), Row(a=1, b=2, c=1)]
-        )
-
-    def test_udtf_with_lateral_join(self):
-        class TestUDTF:
-            def eval(self, a: int, b: int) -> Iterator:
-                yield a, b, a + b
-                yield a, b, a - b
-
-        func = udtf(TestUDTF, returnType="a: int, b: int, c: int")
-        self.spark.udtf.register("testUDTF", func)
-        df = self.spark.sql(
-            "SELECT f.* FROM values (0, 1), (1, 2) t(a, b), LATERAL testUDTF(a, b) f"
-        )
-        expected = self.spark.createDataFrame(
-            [(0, 1, 1), (0, 1, -1), (1, 2, 3), (1, 2, -1)], schema=["a", "b", "c"]
-        )
-        self.assertEqual(df.collect(), expected.collect())
-
-    def test_udtf_eval_with_return_stmt(self):
-        class TestUDTF:
-            def eval(self, a: int, b: int):
-                return [(a, a + 1), (b, b + 1)]
-
-        func = udtf(TestUDTF, returnType="a: int, b: int")
-        rows = func(lit(1), lit(2)).collect()
-        self.assertEqual(rows, [Row(a=1, b=2), Row(a=2, b=3)])
-
-    def test_udtf_eval_returning_non_tuple(self):
-        @udtf(returnType="a: int")
-        class TestUDTF:
-            def eval(self, a: int):
-                yield a
-
-        with self.assertRaisesRegex(PythonException, "UDTF_INVALID_OUTPUT_ROW_TYPE"):
-            TestUDTF(lit(1)).collect()
-
-        @udtf(returnType="a: int")
-        class TestUDTF:
-            def eval(self, a: int):
-                return (a,)
-
-        with self.assertRaisesRegex(PythonException, "UDTF_INVALID_OUTPUT_ROW_TYPE"):
-            TestUDTF(lit(1)).collect()
-
-    def test_udtf_with_invalid_return_value(self):
-        @udtf(returnType="x: int")
-        class TestUDTF:
-            def eval(self, a):
-                return a
-
-        with self.assertRaisesRegex(PythonException, "UDTF_RETURN_NOT_ITERABLE"):
-            TestUDTF(lit(1)).collect()
-
-    def test_udtf_with_zero_arg_and_invalid_return_value(self):
-        @udtf(returnType="x: int")
-        class TestUDTF:
-            def eval(self):
-                return 1
-
-        with self.assertRaisesRegex(PythonException, "UDTF_RETURN_NOT_ITERABLE"):
-            TestUDTF().collect()
-
-    def test_udtf_with_invalid_return_value_in_terminate(self):
-        @udtf(returnType="x: int")
-        class TestUDTF:
-            def eval(self, a):
-                ...
-
-            def terminate(self):
-                return 1
-
-        with self.assertRaisesRegex(PythonException, "UDTF_RETURN_NOT_ITERABLE"):
-            TestUDTF(lit(1)).collect()
-
-    def test_udtf_eval_with_no_return(self):
-        @udtf(returnType="a: int")
-        class TestUDTF:
-            def eval(self, a: int):
-                ...
-
-        self.assertEqual(TestUDTF(lit(1)).collect(), [])
-
-        @udtf(returnType="a: int")
-        class TestUDTF:
-            def eval(self, a: int):
-                return
-
-        self.assertEqual(TestUDTF(lit(1)).collect(), [])
-
-    def test_udtf_with_conditional_return(self):
-        class TestUDTF:
-            def eval(self, a: int):
-                if a > 5:
-                    yield a,
-
-        func = udtf(TestUDTF, returnType="a: int")
-        self.spark.udtf.register("test_udtf", func)
-        self.assertEqual(
-            self.spark.sql("SELECT * FROM range(0, 8) JOIN LATERAL test_udtf(id)").collect(),
-            [Row(id=6, a=6), Row(id=7, a=7)],
-        )
-
-    def test_udtf_with_empty_yield(self):
-        @udtf(returnType="a: int")
-        class TestUDTF:
-            def eval(self, a: int):
-                yield
-
-        assertDataFrameEqual(TestUDTF(lit(1)), [Row(a=None)])
-
-    def test_udtf_with_none_output(self):
-        @udtf(returnType="a: int")
-        class TestUDTF:
-            def eval(self, a: int):
-                yield a,
-                yield None,
-
-        self.assertEqual(TestUDTF(lit(1)).collect(), [Row(a=1), Row(a=None)])
-        df = self.spark.createDataFrame([(0, 1), (1, 2)], schema=["a", "b"])
-        self.assertEqual(TestUDTF(lit(1)).join(df, "a", "inner").collect(), [Row(a=1, b=2)])
-        assertDataFrameEqual(
-            TestUDTF(lit(1)).join(df, "a", "left"), [Row(a=None, b=None), Row(a=1, b=2)]
-        )
-
-    def test_udtf_with_none_input(self):
-        @udtf(returnType="a: int")
-        class TestUDTF:
-            def eval(self, a: int):
-                yield a,
-
-        self.assertEqual(TestUDTF(lit(None)).collect(), [Row(a=None)])
-        self.spark.udtf.register("testUDTF", TestUDTF)
-        df = self.spark.sql("SELECT * FROM testUDTF(null)")
-        self.assertEqual(df.collect(), [Row(a=None)])
-
-    def test_udtf_with_wrong_num_input(self):
-        @udtf(returnType="a: int, b: int")
-        class TestUDTF:
-            def eval(self, a: int):
-                yield a, a + 1
-
-        with self.assertRaisesRegex(
-            PythonException, r"eval\(\) missing 1 required positional argument: 'a'"
-        ):
-            TestUDTF().collect()
-
-        with self.assertRaisesRegex(
-            PythonException, r"eval\(\) takes 2 positional arguments but 3 were given"
-        ):
-            TestUDTF(lit(1), lit(2)).collect()
-
-    def test_udtf_init_with_additional_args(self):
-        @udtf(returnType="x int")
-        class TestUDTF:
-            def __init__(self, a: int):
-                ...
-
-            def eval(self, a: int):
-                yield a,
-
-        with self.assertRaisesRegex(
-            PythonException, r"__init__\(\) missing 1 required positional argument: 'a'"
-        ):
-            TestUDTF(lit(1)).show()
-
-    def test_udtf_terminate_with_additional_args(self):
-        @udtf(returnType="x int")
-        class TestUDTF:
-            def eval(self, a: int):
-                yield a,
-
-            def terminate(self, a: int):
-                ...
-
-        with self.assertRaisesRegex(
-            PythonException, r"terminate\(\) missing 1 required positional argument: 'a'"
-        ):
-            TestUDTF(lit(1)).show()
-
-    def test_udtf_with_wrong_num_output(self):
-        err_msg = (
-            r"\[UDTF_RETURN_SCHEMA_MISMATCH\] The number of columns in the "
-            "result does not match the specified schema."
-        )
-
-        # Output less columns than specified return schema
-        @udtf(returnType="a: int, b: int")
-        class TestUDTF:
-            def eval(self, a: int):
-                yield a,
-
-        with self.assertRaisesRegex(PythonException, err_msg):
-            TestUDTF(lit(1)).collect()
-
-        # Output more columns than specified return schema
-        @udtf(returnType="a: int")
-        class TestUDTF:
-            def eval(self, a: int):
-                yield a, a + 1
-
-        with self.assertRaisesRegex(PythonException, err_msg):
-            TestUDTF(lit(1)).collect()
-
-    def test_udtf_with_empty_output_schema_and_non_empty_output(self):
-        @udtf(returnType=StructType())
-        class TestUDTF:
-            def eval(self):
-                yield 1,
-
-        with self.assertRaisesRegex(PythonException, "UDTF_RETURN_SCHEMA_MISMATCH"):
-            TestUDTF().collect()
-
-    def test_udtf_with_non_empty_output_schema_and_empty_output(self):
-        @udtf(returnType="a: int")
-        class TestUDTF:
-            def eval(self):
-                yield tuple()
-
-        with self.assertRaisesRegex(PythonException, "UDTF_RETURN_SCHEMA_MISMATCH"):
-            TestUDTF().collect()
-
-    def test_udtf_init(self):
-        @udtf(returnType="a: int, b: int, c: string")
-        class TestUDTF:
-            def __init__(self):
-                self.key = "test"
-
-            def eval(self, a: int):
-                yield a, a + 1, self.key
-
-        rows = TestUDTF(lit(1)).collect()
-        self.assertEqual(rows, [Row(a=1, b=2, c="test")])
-
-    def test_udtf_terminate(self):
-        @udtf(returnType="key: string, value: float")
-        class TestUDTF:
-            def __init__(self):
-                self._count = 0
-                self._sum = 0
-
-            def eval(self, x: int):
-                self._count += 1
-                self._sum += x
-                yield "input", float(x)
-
-            def terminate(self):
-                yield "count", float(self._count)
-                yield "avg", self._sum / self._count
-
-        self.assertEqual(
-            TestUDTF(lit(1)).collect(),
-            [Row(key="input", value=1), Row(key="count", value=1.0), Row(key="avg", value=1.0)],
-        )
-
-        self.spark.udtf.register("test_udtf", TestUDTF)
-        df = self.spark.sql(
-            "SELECT id, key, value FROM range(0, 10, 1, 2), "
-            "LATERAL test_udtf(id) WHERE key != 'input'"
-        )
-        self.assertEqual(
-            df.collect(),
-            [
-                Row(id=4, key="count", value=5.0),
-                Row(id=4, key="avg", value=2.0),
-                Row(id=9, key="count", value=5.0),
-                Row(id=9, key="avg", value=7.0),
-            ],
-        )
-
-    def test_init_with_exception(self):
-        @udtf(returnType="x: int")
-        class TestUDTF:
-            def __init__(self):
-                raise Exception("error")
-
-            def eval(self):
-                yield 1,
-
-        with self.assertRaisesRegex(
-            PythonException,
-            r"\[UDTF_EXEC_ERROR\] User defined table function encountered an error "
-            r"in the '__init__' method: error",
-        ):
-            TestUDTF().show()
-
-    def test_eval_with_exception(self):
-        @udtf(returnType="x: int")
-        class TestUDTF:
-            def eval(self):
-                raise Exception("error")
-
-        with self.assertRaisesRegex(
-            PythonException,
-            r"\[UDTF_EXEC_ERROR\] User defined table function encountered an error "
-            r"in the 'eval' method: error",
-        ):
-            TestUDTF().show()
-
-    def test_terminate_with_exceptions(self):
-        @udtf(returnType="a: int, b: int")
-        class TestUDTF:
-            def eval(self, a: int):
-                yield a, a + 1
-
-            def terminate(self):
-                raise ValueError("terminate error")
-
-        with self.assertRaisesRegex(
-            PythonException,
-            r"\[UDTF_EXEC_ERROR\] User defined table function encountered an error "
-            r"in the 'terminate' method: terminate error",
-        ):
-            TestUDTF(lit(1)).collect()
-
-    def test_udtf_terminate_with_wrong_num_output(self):
-        err_msg = (
-            r"\[UDTF_RETURN_SCHEMA_MISMATCH\] The number of columns in the result "
-            "does not match the specified schema."
-        )
-
-        @udtf(returnType="a: int, b: int")
-        class TestUDTF:
-            def eval(self, a: int):
-                yield a, a + 1
-
-            def terminate(self):
-                yield 1, 2, 3
-
-        with self.assertRaisesRegex(PythonException, err_msg):
-            TestUDTF(lit(1)).show()
-
-        @udtf(returnType="a: int, b: int")
-        class TestUDTF:
-            def eval(self, a: int):
-                yield a, a + 1
-
-            def terminate(self):
-                yield 1,
-
-        with self.assertRaisesRegex(PythonException, err_msg):
-            TestUDTF(lit(1)).show()
-
-    def test_udtf_determinism(self):
-        class TestUDTF:
-            def eval(self, a: int):
-                yield a,
-
-        func = udtf(TestUDTF, returnType="x: int")
-        # The UDTF is marked as non-deterministic by default.
-        self.assertFalse(func.deterministic)
-        func = func.asDeterministic()
-        self.assertTrue(func.deterministic)
-
-    def test_nondeterministic_udtf(self):
-        import random
-
-        class RandomUDTF:
-            def eval(self, a: int):
-                yield a + int(random.random()),
-
-        random_udtf = udtf(RandomUDTF, returnType="x: int")
-        assertDataFrameEqual(random_udtf(lit(1)), [Row(x=1)])
-        self.spark.udtf.register("random_udtf", random_udtf)
-        assertDataFrameEqual(self.spark.sql("select * from random_udtf(1)"), [Row(x=1)])
-
-    def test_udtf_with_nondeterministic_input(self):
-        from pyspark.sql.functions import rand
-
-        @udtf(returnType="x: int")
-        class TestUDTF:
-            def eval(self, a: int):
-                yield 1 if a > 100 else 0,
-
-        assertDataFrameEqual(TestUDTF(rand(0) * 100), [Row(x=0)])
-
-    def test_udtf_with_invalid_return_type(self):
-        @udtf(returnType="int")
-        class TestUDTF:
-            def eval(self, a: int):
-                yield a + 1,
-
-        with self.assertRaises(PySparkTypeError) as e:
-            TestUDTF(lit(1)).collect()
-
-        self.check_error(
-            exception=e.exception,
-            error_class="UDTF_RETURN_TYPE_MISMATCH",
-            message_parameters={"name": "TestUDTF", "return_type": "IntegerType()"},
-        )
-
-        @udtf(returnType=MapType(StringType(), IntegerType()))
-        class TestUDTF:
-            def eval(self, a: int):
-                yield a + 1,
-
-        with self.assertRaises(PySparkTypeError) as e:
-            TestUDTF(lit(1)).collect()
-
-        self.check_error(
-            exception=e.exception,
-            error_class="UDTF_RETURN_TYPE_MISMATCH",
-            message_parameters={
-                "name": "TestUDTF",
-                "return_type": "MapType(StringType(), IntegerType(), True)",
-            },
-        )
-
-    def test_udtf_with_struct_input_type(self):
-        @udtf(returnType="x: string")
-        class TestUDTF:
-            def eval(self, person):
-                yield f"{person.name}: {person.age}",
-
-        self.spark.udtf.register("test_udtf", TestUDTF)
-        self.assertEqual(
-            self.spark.sql(
-                "select * from test_udtf(named_struct('name', 'Alice', 'age', 1))"
-            ).collect(),
-            [Row(x="Alice: 1")],
-        )
-
-    def test_udtf_with_array_input_type(self):
-        @udtf(returnType="x: string")
-        class TestUDTF:
-            def eval(self, args):
-                yield str(args),
-
-        self.spark.udtf.register("test_udtf", TestUDTF)
-        self.assertEqual(
-            self.spark.sql("select * from test_udtf(array(1, 2, 3))").collect(),
-            [Row(x="[1, 2, 3]")],
-        )
-
-    def test_udtf_with_map_input_type(self):
-        @udtf(returnType="x: string")
-        class TestUDTF:
-            def eval(self, m):
-                yield str(m),
-
-        self.spark.udtf.register("test_udtf", TestUDTF)
-        self.assertEqual(
-            self.spark.sql("select * from test_udtf(map('key', 'value'))").collect(),
-            [Row(x="{'key': 'value'}")],
-        )
-
-    def test_udtf_with_struct_output_types(self):
-        @udtf(returnType="x: struct<a:int,b:int>")
-        class TestUDTF:
-            def eval(self, x: int):
-                yield {"a": x, "b": x + 1},
-
-        self.assertEqual(TestUDTF(lit(1)).collect(), [Row(x=Row(a=1, b=2))])
-
-    def test_udtf_with_array_output_types(self):
-        @udtf(returnType="x: array<int>")
-        class TestUDTF:
-            def eval(self, x: int):
-                yield [x, x + 1, x + 2],
-
-        self.assertEqual(TestUDTF(lit(1)).collect(), [Row(x=[1, 2, 3])])
-
-    def test_udtf_with_map_output_types(self):
-        @udtf(returnType="x: map<int,string>")
-        class TestUDTF:
-            def eval(self, x: int):
-                yield {x: str(x)},
-
-        self.assertEqual(TestUDTF(lit(1)).collect(), [Row(x={1: "1"})])
-
-    def test_udtf_with_empty_output_types(self):
-        @udtf(returnType=StructType())
-        class TestUDTF:
-            def eval(self):
-                yield tuple()
-
-        assertDataFrameEqual(TestUDTF(), [Row()])
-
-    def _check_result_or_exception(
-        self, func_handler, ret_type, expected, *, err_type=PythonException
-    ):
-        func = udtf(func_handler, returnType=ret_type)
-        if not isinstance(expected, str):
-            assertDataFrameEqual(func(), expected)
-        else:
-            with self.assertRaisesRegex(err_type, expected):
-                func().collect()
-
-    def test_numeric_output_type_casting(self):
-        class TestUDTF:
-            def eval(self):
-                yield 1,
-
-        for i, (ret_type, expected) in enumerate(
-            [
-                ("x: boolean", [Row(x=None)]),
-                ("x: tinyint", [Row(x=1)]),
-                ("x: smallint", [Row(x=1)]),
-                ("x: int", [Row(x=1)]),
-                ("x: bigint", [Row(x=1)]),
-                ("x: string", [Row(x="1")]),  # int to string is ok, but string to int is None
-                (
-                    "x: date",
-                    "AttributeError",
-                ),  # AttributeError: 'int' object has no attribute 'toordinal'
-                (
-                    "x: timestamp",
-                    "AttributeError",
-                ),  # AttributeError: 'int' object has no attribute 'tzinfo'
-                ("x: byte", [Row(x=1)]),
-                ("x: binary", [Row(x=None)]),
-                ("x: float", [Row(x=None)]),
-                ("x: double", [Row(x=None)]),
-                ("x: decimal(10, 0)", [Row(x=None)]),
-                ("x: array<int>", [Row(x=None)]),
-                ("x: map<string,int>", [Row(x=None)]),
-                ("x: struct<a:int>", "UNEXPECTED_TUPLE_WITH_STRUCT"),
-            ]
-        ):
-            with self.subTest(ret_type=ret_type):
-                self._check_result_or_exception(TestUDTF, ret_type, expected)
-
-    def test_numeric_string_output_type_casting(self):
-        class TestUDTF:
-            def eval(self):
-                yield "1",
-
-        for ret_type, expected in [
-            ("x: boolean", [Row(x=None)]),
-            ("x: tinyint", [Row(x=None)]),
-            ("x: smallint", [Row(x=None)]),
-            ("x: int", [Row(x=None)]),
-            ("x: bigint", [Row(x=None)]),
-            ("x: string", [Row(x="1")]),
-            ("x: date", "AttributeError"),
-            ("x: timestamp", "AttributeError"),
-            ("x: byte", [Row(x=None)]),
-            ("x: binary", [Row(x=bytearray(b"1"))]),
-            ("x: float", [Row(x=None)]),
-            ("x: double", [Row(x=None)]),
-            ("x: decimal(10, 0)", [Row(x=None)]),
-            ("x: array<int>", [Row(x=None)]),
-            ("x: map<string,int>", [Row(x=None)]),
-            ("x: struct<a:int>", "UNEXPECTED_TUPLE_WITH_STRUCT"),
-        ]:
-            with self.subTest(ret_type=ret_type):
-                self._check_result_or_exception(TestUDTF, ret_type, expected)
-
-    def test_string_output_type_casting(self):
-        class TestUDTF:
-            def eval(self):
-                yield "hello",
-
-        for ret_type, expected in [
-            ("x: boolean", [Row(x=None)]),
-            ("x: tinyint", [Row(x=None)]),
-            ("x: smallint", [Row(x=None)]),
-            ("x: int", [Row(x=None)]),
-            ("x: bigint", [Row(x=None)]),
-            ("x: string", [Row(x="hello")]),
-            ("x: date", "AttributeError"),
-            ("x: timestamp", "AttributeError"),
-            ("x: byte", [Row(x=None)]),
-            ("x: binary", [Row(x=bytearray(b"hello"))]),
-            ("x: float", [Row(x=None)]),
-            ("x: double", [Row(x=None)]),
-            ("x: decimal(10, 0)", [Row(x=None)]),
-            ("x: array<int>", [Row(x=None)]),
-            ("x: map<string,int>", [Row(x=None)]),
-            ("x: struct<a:int>", "UNEXPECTED_TUPLE_WITH_STRUCT"),
-        ]:
-            with self.subTest(ret_type=ret_type):
-                self._check_result_or_exception(TestUDTF, ret_type, expected)
-
-    def test_array_output_type_casting(self):
-        class TestUDTF:
-            def eval(self):
-                yield [0, 1.1, 2],
-
-        for ret_type, expected in [
-            ("x: boolean", [Row(x=None)]),
-            ("x: tinyint", [Row(x=None)]),
-            ("x: smallint", [Row(x=None)]),
-            ("x: int", [Row(x=None)]),
-            ("x: bigint", [Row(x=None)]),
-            ("x: string", [Row(x="[0, 1.1, 2]")]),
-            ("x: date", "AttributeError"),
-            ("x: timestamp", "AttributeError"),
-            ("x: byte", [Row(x=None)]),
-            ("x: binary", [Row(x=None)]),
-            ("x: float", [Row(x=None)]),
-            ("x: double", [Row(x=None)]),
-            ("x: decimal(10, 0)", [Row(x=None)]),
-            ("x: array<int>", [Row(x=[0, None, 2])]),
-            ("x: array<double>", [Row(x=[None, 1.1, None])]),
-            ("x: array<string>", [Row(x=["0", "1.1", "2"])]),
-            ("x: array<boolean>", [Row(x=[None, None, None])]),
-            ("x: array<array<int>>", [Row(x=[None, None, None])]),
-            ("x: map<string,int>", [Row(x=None)]),
-            ("x: struct<a:int,b:int,c:int>", [Row(x=Row(a=0, b=None, c=2))]),
-        ]:
-            with self.subTest(ret_type=ret_type):
-                self._check_result_or_exception(TestUDTF, ret_type, expected)
-
-    def test_map_output_type_casting(self):
-        class TestUDTF:
-            def eval(self):
-                yield {"a": 0, "b": 1.1, "c": 2},
-
-        for ret_type, expected in [
-            ("x: boolean", [Row(x=None)]),
-            ("x: tinyint", [Row(x=None)]),
-            ("x: smallint", [Row(x=None)]),
-            ("x: int", [Row(x=None)]),
-            ("x: bigint", [Row(x=None)]),
-            ("x: string", [Row(x="{a=0, b=1.1, c=2}")]),
-            ("x: date", "AttributeError"),
-            ("x: timestamp", "AttributeError"),
-            ("x: byte", [Row(x=None)]),
-            ("x: binary", [Row(x=None)]),
-            ("x: float", [Row(x=None)]),
-            ("x: double", [Row(x=None)]),
-            ("x: decimal(10, 0)", [Row(x=None)]),
-            ("x: array<string>", [Row(x=None)]),
-            ("x: map<string,string>", [Row(x={"a": "0", "b": "1.1", "c": "2"})]),
-            ("x: map<string,boolean>", [Row(x={"a": None, "b": None, "c": None})]),
-            ("x: map<string,int>", [Row(x={"a": 0, "b": None, "c": 2})]),
-            ("x: map<string,float>", [Row(x={"a": None, "b": 1.1, "c": None})]),
-            ("x: map<string,map<string,int>>", [Row(x={"a": None, "b": None, "c": None})]),
-            ("x: struct<a:int>", [Row(x=Row(a=0))]),
-        ]:
-            with self.subTest(ret_type=ret_type):
-                self._check_result_or_exception(TestUDTF, ret_type, expected)
-
-    def test_struct_output_type_casting_dict(self):
-        class TestUDTF:
-            def eval(self):
-                yield {"a": 0, "b": 1.1, "c": 2},
-
-        for ret_type, expected in [
-            ("x: boolean", [Row(x=None)]),
-            ("x: tinyint", [Row(x=None)]),
-            ("x: smallint", [Row(x=None)]),
-            ("x: int", [Row(x=None)]),
-            ("x: bigint", [Row(x=None)]),
-            ("x: string", [Row(x="{a=0, b=1.1, c=2}")]),
-            ("x: date", "AttributeError"),
-            ("x: timestamp", "AttributeError"),
-            ("x: byte", [Row(x=None)]),
-            ("x: binary", [Row(x=None)]),
-            ("x: float", [Row(x=None)]),
-            ("x: double", [Row(x=None)]),
-            ("x: decimal(10, 0)", [Row(x=None)]),
-            ("x: array<string>", [Row(x=None)]),
-            ("x: map<string,string>", [Row(x={"a": "0", "b": "1.1", "c": "2"})]),
-            ("x: struct<a:string,b:string,c:string>", [Row(Row(a="0", b="1.1", c="2"))]),
-            ("x: struct<a:int,b:int,c:int>", [Row(Row(a=0, b=None, c=2))]),
-            ("x: struct<a:float,b:float,c:float>", [Row(Row(a=None, b=1.1, c=None))]),
-        ]:
-            with self.subTest(ret_type=ret_type):
-                self._check_result_or_exception(TestUDTF, ret_type, expected)
-
-    def test_struct_output_type_casting_row(self):
-        self.check_struct_output_type_casting_row(Py4JJavaError)
-
-    def check_struct_output_type_casting_row(self, error_type):
-        class TestUDTF:
-            def eval(self):
-                yield Row(a=0, b=1.1, c=2),
-
-        err = ("PickleException", error_type)
-
-        for ret_type, expected in [
-            ("x: boolean", err),
-            ("x: tinyint", err),
-            ("x: smallint", err),
-            ("x: int", err),
-            ("x: bigint", err),
-            ("x: string", err),
-            ("x: date", "ValueError"),
-            ("x: timestamp", "ValueError"),
-            ("x: byte", err),
-            ("x: binary", err),
-            ("x: float", err),
-            ("x: double", err),
-            ("x: decimal(10, 0)", err),
-            ("x: array<string>", err),
-            ("x: map<string,string>", err),
-            ("x: struct<a:string,b:string,c:string>", [Row(Row(a="0", b="1.1", c="2"))]),
-            ("x: struct<a:int,b:int,c:int>", [Row(Row(a=0, b=None, c=2))]),
-            ("x: struct<a:float,b:float,c:float>", [Row(Row(a=None, b=1.1, c=None))]),
-        ]:
-            with self.subTest(ret_type=ret_type):
-                if isinstance(expected, tuple):
-                    self._check_result_or_exception(
-                        TestUDTF, ret_type, expected[0], err_type=expected[1]
-                    )
-                else:
-                    self._check_result_or_exception(TestUDTF, ret_type, expected)
-
-    def test_inconsistent_output_types(self):
-        class TestUDTF:
-            def eval(self):
-                yield 1,
-                yield [1, 2],
-
-        for ret_type, expected in [
-            ("x: int", [Row(x=1), Row(x=None)]),
-            ("x: array<int>", [Row(x=None), Row(x=[1, 2])]),
-        ]:
-            with self.subTest(ret_type=ret_type):
-                assertDataFrameEqual(udtf(TestUDTF, returnType=ret_type)(), expected)
-
-    @unittest.skipIf(not have_pandas, pandas_requirement_message)
-    def test_udtf_with_pandas_input_type(self):
-        import pandas as pd
-
-        @udtf(returnType="corr: double")
-        class TestUDTF:
-            def eval(self, s1: pd.Series, s2: pd.Series):
-                yield s1.corr(s2)
-
-        self.spark.udtf.register("test_udtf", TestUDTF)
-        # TODO(SPARK-43968): check during compile time instead of runtime
-        with self.assertRaisesRegex(
-            PythonException, "AttributeError: 'int' object has no attribute 'corr'"
-        ):
-            self.spark.sql(
-                "select * from values (1, 2), (2, 3) t(a, b), lateral test_udtf(a, b)"
-            ).collect()
-
-    def test_udtf_register_error(self):
-        @udf
-        def upper(s: str):
-            return s.upper()
-
-        with self.assertRaises(PySparkTypeError) as e:
-            self.spark.udtf.register("test_udf", upper)
-
-        self.check_error(
-            exception=e.exception,
-            error_class="INVALID_UDTF_EVAL_TYPE",
-            message_parameters={
-                "name": "test_udf",
-                "eval_type": "SQL_TABLE_UDF, SQL_ARROW_TABLE_UDF",
-            },
-        )
-
-    def test_udtf_pickle_error(self):
-        with tempfile.TemporaryDirectory() as d:
-            file = os.path.join(d, "file.txt")
-            file_obj = open(file, "w")
-
-            @udtf(returnType="x: int")
-            class TestUDTF:
-                def eval(self):
-                    file_obj
-                    yield 1,
-
-            with self.assertRaisesRegex(PySparkRuntimeError, "UDTF_SERIALIZATION_ERROR"):
-                TestUDTF().collect()
-
-    def test_udtf_access_spark_session(self):
-        df = self.spark.range(10)
-
-        @udtf(returnType="x: int")
-        class TestUDTF:
-            def eval(self):
-                df.collect()
-                yield 1,
-
-        with self.assertRaisesRegex(PySparkRuntimeError, "UDTF_SERIALIZATION_ERROR"):
-            TestUDTF().collect()
-
-    def test_udtf_no_eval(self):
-        with self.assertRaises(PySparkAttributeError) as e:
-
-            @udtf(returnType="a: int, b: int")
-            class TestUDTF:
-                def run(self, a: int):
-                    yield a, a + 1
-
-        self.check_error(
-            exception=e.exception,
-            error_class="INVALID_UDTF_NO_EVAL",
-            message_parameters={"name": "TestUDTF"},
-        )
-
-    def test_udtf_with_no_handler_class(self):
-        with self.assertRaises(PySparkTypeError) as e:
-
-            @udtf(returnType="a: int")
-            def test_udtf(a: int):
-                yield a,
-
-        self.check_error(
-            exception=e.exception,
-            error_class="INVALID_UDTF_HANDLER_TYPE",
-            message_parameters={"type": "function"},
-        )
-
-        with self.assertRaises(PySparkTypeError) as e:
-            udtf(1, returnType="a: int")
-
-        self.check_error(
-            exception=e.exception,
-            error_class="INVALID_UDTF_HANDLER_TYPE",
-            message_parameters={"type": "int"},
-        )
-
-    def test_udtf_with_table_argument_query(self):
-        class TestUDTF:
-            def eval(self, row: Row):
-                if row["id"] > 5:
-                    yield row["id"],
-
-        func = udtf(TestUDTF, returnType="a: int")
-        self.spark.udtf.register("test_udtf", func)
-        self.assertEqual(
-            self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)))").collect(),
-            [Row(a=6), Row(a=7)],
-        )
-
-    def test_udtf_with_int_and_table_argument_query(self):
-        class TestUDTF:
-            def eval(self, i: int, row: Row):
-                if row["id"] > i:
-                    yield row["id"],
-
-        func = udtf(TestUDTF, returnType="a: int")
-        self.spark.udtf.register("test_udtf", func)
-        self.assertEqual(
-            self.spark.sql(
-                "SELECT * FROM test_udtf(5, TABLE (SELECT id FROM range(0, 8)))"
-            ).collect(),
-            [Row(a=6), Row(a=7)],
-        )
-
-    def test_udtf_with_table_argument_identifier(self):
-        class TestUDTF:
-            def eval(self, row: Row):
-                if row["id"] > 5:
-                    yield row["id"],
-
-        func = udtf(TestUDTF, returnType="a: int")
-        self.spark.udtf.register("test_udtf", func)
-
-        with self.tempView("v"):
-            self.spark.sql("CREATE OR REPLACE TEMPORARY VIEW v as SELECT id FROM range(0, 8)")
-            self.assertEqual(
-                self.spark.sql("SELECT * FROM test_udtf(TABLE (v))").collect(),
-                [Row(a=6), Row(a=7)],
-            )
-
-    def test_udtf_with_int_and_table_argument_identifier(self):
-        class TestUDTF:
-            def eval(self, i: int, row: Row):
-                if row["id"] > i:
-                    yield row["id"],
-
-        func = udtf(TestUDTF, returnType="a: int")
-        self.spark.udtf.register("test_udtf", func)
-
-        with self.tempView("v"):
-            self.spark.sql("CREATE OR REPLACE TEMPORARY VIEW v as SELECT id FROM range(0, 8)")
-            self.assertEqual(
-                self.spark.sql("SELECT * FROM test_udtf(5, TABLE (v))").collect(),
-                [Row(a=6), Row(a=7)],
-            )
-
-    def test_udtf_with_table_argument_unknown_identifier(self):
-        class TestUDTF:
-            def eval(self, row: Row):
-                if row["id"] > 5:
-                    yield row["id"],
-
-        func = udtf(TestUDTF, returnType="a: int")
-        self.spark.udtf.register("test_udtf", func)
-
-        with self.assertRaisesRegex(AnalysisException, "TABLE_OR_VIEW_NOT_FOUND"):
-            self.spark.sql("SELECT * FROM test_udtf(TABLE (v))").collect()
-
-    def test_udtf_with_table_argument_malformed_query(self):
-        class TestUDTF:
-            def eval(self, row: Row):
-                if row["id"] > 5:
-                    yield row["id"],
-
-        func = udtf(TestUDTF, returnType="a: int")
-        self.spark.udtf.register("test_udtf", func)
-
-        with self.assertRaisesRegex(AnalysisException, "TABLE_OR_VIEW_NOT_FOUND"):
-            self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT * FROM v))").collect()
-
-    def test_udtf_with_table_argument_cte_inside(self):
-        class TestUDTF:
-            def eval(self, row: Row):
-                if row["id"] > 5:
-                    yield row["id"],
-
-        func = udtf(TestUDTF, returnType="a: int")
-        self.spark.udtf.register("test_udtf", func)
-        self.assertEqual(
-            self.spark.sql(
-                """
-                SELECT * FROM test_udtf(TABLE (
-                  WITH t AS (
-                    SELECT id FROM range(0, 8)
-                  )
-                  SELECT * FROM t
-                ))
-                """
-            ).collect(),
-            [Row(a=6), Row(a=7)],
-        )
-
-    def test_udtf_with_table_argument_cte_outside(self):
-        class TestUDTF:
-            def eval(self, row: Row):
-                if row["id"] > 5:
-                    yield row["id"],
-
-        func = udtf(TestUDTF, returnType="a: int")
-        self.spark.udtf.register("test_udtf", func)
-        self.assertEqual(
-            self.spark.sql(
-                """
-                WITH t AS (
-                  SELECT id FROM range(0, 8)
-                )
-                SELECT * FROM test_udtf(TABLE (SELECT id FROM t))
-                """
-            ).collect(),
-            [Row(a=6), Row(a=7)],
-        )
-
-        self.assertEqual(
-            self.spark.sql(
-                """
-                WITH t AS (
-                  SELECT id FROM range(0, 8)
-                )
-                SELECT * FROM test_udtf(TABLE (t))
-                """
-            ).collect(),
-            [Row(a=6), Row(a=7)],
-        )
-
-    # TODO(SPARK-44233): Fix the subquery resolution.
-    @unittest.skip("Fails to resolve the subquery.")
-    def test_udtf_with_table_argument_lateral_join(self):
-        class TestUDTF:
-            def eval(self, row: Row):
-                if row["id"] > 5:
-                    yield row["id"],
-
-        func = udtf(TestUDTF, returnType="a: int")
-        self.spark.udtf.register("test_udtf", func)
-        self.assertEqual(
-            self.spark.sql(
-                """
-                SELECT * FROM
-                  range(0, 8) AS t,
-                  LATERAL test_udtf(TABLE (t))
-                """
-            ).collect(),
-            [Row(a=6), Row(a=7)],
-        )
-
-    def test_udtf_with_table_argument_multiple(self):
-        class TestUDTF:
-            def eval(self, a: Row, b: Row):
-                yield a[0], b[0]
-
-        func = udtf(TestUDTF, returnType="a: int, b: int")
-        self.spark.udtf.register("test_udtf", func)
-
-        query = """
-          SELECT * FROM test_udtf(
-            TABLE (SELECT id FROM range(0, 2)),
-            TABLE (SELECT id FROM range(0, 3)))
-        """
-
-        with self.sql_conf({"spark.sql.tvf.allowMultipleTableArguments.enabled": False}):
-            with self.assertRaisesRegex(
-                AnalysisException, "TABLE_VALUED_FUNCTION_TOO_MANY_TABLE_ARGUMENTS"
-            ):
-                self.spark.sql(query).collect()
-
-        with self.sql_conf({"spark.sql.tvf.allowMultipleTableArguments.enabled": True}):
-            self.assertEqual(
-                self.spark.sql(query).collect(),
-                [
-                    Row(a=0, b=0),
-                    Row(a=1, b=0),
-                    Row(a=0, b=1),
-                    Row(a=1, b=1),
-                    Row(a=0, b=2),
-                    Row(a=1, b=2),
-                ],
-            )
-
-    def test_docstring(self):
-        class TestUDTF:
-            """A UDTF for test."""
-
-            def __init__(self):
-                """Initialize the UDTF"""
-                ...
-
-            def eval(self, x: int):
-                """Evaluate the input row."""
-                yield x + 1,
-
-            def terminate(self):
-                """Terminate the UDTF."""
-                ...
-
-        cls = udtf(TestUDTF, returnType="y: int").func
-        self.assertIn("A UDTF for test", cls.__doc__)
-        self.assertIn("Initialize the UDTF", cls.__init__.__doc__)
-        self.assertIn("Evaluate the input row", cls.eval.__doc__)
-        self.assertIn("Terminate the UDTF", cls.terminate.__doc__)
-
-
-class UDTFTests(BaseUDTFTestsMixin, ReusedSQLTestCase):
-    @classmethod
-    def setUpClass(cls):
-        super(UDTFTests, cls).setUpClass()
-        cls.spark.conf.set("spark.sql.execution.pythonUDTF.arrow.enabled", "false")
-
-    @classmethod
-    def tearDownClass(cls):
-        try:
-            cls.spark.conf.unset("spark.sql.execution.pythonUDTF.arrow.enabled")
-        finally:
-            super(UDTFTests, cls).tearDownClass()
-
-
-@unittest.skipIf(
-    not have_pandas or not have_pyarrow, pandas_requirement_message or pyarrow_requirement_message
-)
-class UDTFArrowTestsMixin(BaseUDTFTestsMixin):
-    def test_eval_type(self):
-        def upper(x: str):
-            return upper(x)
-
-        class TestUDTF:
-            def eval(self, x: str):
-                return upper(x)
-
-        self.assertEqual(
-            udtf(TestUDTF, returnType="x: string", useArrow=False).evalType,
-            PythonEvalType.SQL_TABLE_UDF,
-        )
-
-        self.assertEqual(
-            udtf(TestUDTF, returnType="x: string", useArrow=True).evalType,
-            PythonEvalType.SQL_ARROW_TABLE_UDF,
-        )
-
-    def test_udtf_arrow_sql_conf(self):
-        class TestUDTF:
-            def eval(self):
-                yield 1,
-
-        # We do not use `self.sql_conf` here to test the SQL SET command
-        # instead of using PySpark's `spark.conf.set`.
-        old_value = self.spark.conf.get("spark.sql.execution.pythonUDTF.arrow.enabled")
-        self.spark.sql("SET spark.sql.execution.pythonUDTF.arrow.enabled=False")
-        self.assertEqual(udtf(TestUDTF, returnType="x: int").evalType, PythonEvalType.SQL_TABLE_UDF)
-        self.spark.sql("SET spark.sql.execution.pythonUDTF.arrow.enabled=True")
-        self.assertEqual(
-            udtf(TestUDTF, returnType="x: int").evalType, PythonEvalType.SQL_ARROW_TABLE_UDF
-        )
-        self.spark.conf.set("spark.sql.execution.pythonUDTF.arrow.enabled", old_value)
-
-    def test_udtf_eval_returning_non_tuple(self):
-        @udtf(returnType="a: int")
-        class TestUDTF:
-            def eval(self, a: int):
-                yield a
-
-        # When arrow is enabled, it can handle non-tuple return value.
-        assertDataFrameEqual(TestUDTF(lit(1)), [Row(a=1)])
-
-        @udtf(returnType="a: int")
-        class TestUDTF:
-            def eval(self, a: int):
-                return [a]
-
-        assertDataFrameEqual(TestUDTF(lit(1)), [Row(a=1)])
-
-    def test_numeric_output_type_casting(self):
-        class TestUDTF:
-            def eval(self):
-                yield 1,
-
-        err = "UDTF_ARROW_TYPE_CAST_ERROR"
-
-        for ret_type, expected in [
-            ("x: boolean", [Row(x=True)]),
-            ("x: tinyint", [Row(x=1)]),
-            ("x: smallint", [Row(x=1)]),
-            ("x: int", [Row(x=1)]),
-            ("x: bigint", [Row(x=1)]),
-            ("x: string", [Row(x="1")]),  # require arrow.cast
-            ("x: date", err),
-            ("x: byte", [Row(x=1)]),
-            ("x: binary", [Row(x=bytearray(b"\x01"))]),
-            ("x: float", [Row(x=1.0)]),
-            ("x: double", [Row(x=1.0)]),
-            ("x: decimal(10, 0)", err),
-            ("x: array<int>", err),
-            ("x: map<string,int>", err),
-            ("x: struct<a:int>", err),
-        ]:
-            with self.subTest(ret_type=ret_type):
-                self._check_result_or_exception(TestUDTF, ret_type, expected)
-
-    def test_numeric_string_output_type_casting(self):
-        class TestUDTF:
-            def eval(self):
-                yield "1",
-
-        err = "UDTF_ARROW_TYPE_CAST_ERROR"
-
-        for ret_type, expected in [
-            ("x: boolean", [Row(x=True)]),
-            ("x: tinyint", [Row(x=1)]),
-            ("x: smallint", [Row(x=1)]),
-            ("x: int", [Row(x=1)]),
-            ("x: bigint", [Row(x=1)]),
-            ("x: string", [Row(x="1")]),
-            ("x: date", err),
-            ("x: timestamp", err),
-            ("x: byte", [Row(x=1)]),
-            ("x: binary", [Row(x=bytearray(b"1"))]),
-            ("x: float", [Row(x=1.0)]),
-            ("x: double", [Row(x=1.0)]),
-            ("x: decimal(10, 0)", [Row(x=1)]),
-            ("x: array<string>", [Row(x=["1"])]),
-            ("x: array<int>", [Row(x=[1])]),
-            ("x: map<string,int>", err),
-            ("x: struct<a:int>", err),
-        ]:
-            with self.subTest(ret_type=ret_type):
-                self._check_result_or_exception(TestUDTF, ret_type, expected)
-
-    def test_string_output_type_casting(self):
-        class TestUDTF:
-            def eval(self):
-                yield "hello",
-
-        err = "UDTF_ARROW_TYPE_CAST_ERROR"
-
-        for ret_type, expected in [
-            ("x: boolean", err),
-            ("x: tinyint", err),
-            ("x: smallint", err),
-            ("x: int", err),
-            ("x: bigint", err),
-            ("x: string", [Row(x="hello")]),
-            ("x: date", err),
-            ("x: timestamp", err),
-            ("x: byte", err),
-            ("x: binary", [Row(x=bytearray(b"hello"))]),
-            ("x: float", err),
-            ("x: double", err),
-            ("x: decimal(10, 0)", err),
-            ("x: array<string>", [Row(x=["h", "e", "l", "l", "o"])]),
-            ("x: array<int>", err),
-            ("x: map<string,int>", err),
-            ("x: struct<a:int>", err),
-        ]:
-            with self.subTest(ret_type=ret_type):
-                self._check_result_or_exception(TestUDTF, ret_type, expected)
-
-    def test_array_output_type_casting(self):
-        class TestUDTF:
-            def eval(self):
-                yield [0, 1.1, 2],
-
-        err = "UDTF_ARROW_TYPE_CAST_ERROR"
-
-        for ret_type, expected in [
-            ("x: boolean", err),
-            ("x: tinyint", err),
-            ("x: smallint", err),
-            ("x: int", err),
-            ("x: bigint", err),
-            ("x: string", err),
-            ("x: date", err),
-            ("x: timestamp", err),
-            ("x: byte", err),
-            ("x: binary", err),
-            ("x: float", err),
-            ("x: double", err),
-            ("x: decimal(10, 0)", err),
-            ("x: array<string>", [Row(x=["0", "1.1", "2"])]),
-            ("x: array<boolean>", [Row(x=[False, True, True])]),
-            ("x: array<int>", [Row(x=[0, 1, 2])]),
-            ("x: array<float>", [Row(x=[0, 1.1, 2])]),
-            ("x: array<array<int>>", err),
-            ("x: map<string,int>", err),
-            ("x: struct<a:int>", err),
-            ("x: struct<a:int,b:int,c:int>", err),
-        ]:
-            with self.subTest(ret_type=ret_type):
-                self._check_result_or_exception(TestUDTF, ret_type, expected)
-
-    def test_map_output_type_casting(self):
-        class TestUDTF:
-            def eval(self):
-                yield {"a": 0, "b": 1.1, "c": 2},
-
-        err = "UDTF_ARROW_TYPE_CAST_ERROR"
-
-        for ret_type, expected in [
-            ("x: boolean", err),
-            ("x: tinyint", err),
-            ("x: smallint", err),
-            ("x: int", err),
-            ("x: bigint", err),
-            ("x: string", err),
-            ("x: date", err),
-            ("x: timestamp", err),
-            ("x: byte", err),
-            ("x: binary", err),
-            ("x: float", err),
-            ("x: double", err),
-            ("x: decimal(10, 0)", err),
-            ("x: array<string>", [Row(x=["a", "b", "c"])]),
-            ("x: map<string,string>", err),
-            ("x: map<string,boolean>", err),
-            ("x: map<string,int>", [Row(x={"a": 0, "b": 1, "c": 2})]),
-            ("x: map<string,float>", [Row(x={"a": 0, "b": 1.1, "c": 2})]),
-            ("x: map<string,map<string,int>>", err),
-            ("x: struct<a:int>", [Row(x=Row(a=0))]),
-        ]:
-            with self.subTest(ret_type=ret_type):
-                self._check_result_or_exception(TestUDTF, ret_type, expected)
-
-    def test_struct_output_type_casting_dict(self):
-        class TestUDTF:
-            def eval(self):
-                yield {"a": 0, "b": 1.1, "c": 2},
-
-        err = "UDTF_ARROW_TYPE_CAST_ERROR"
-
-        for ret_type, expected in [
-            ("x: boolean", err),
-            ("x: tinyint", err),
-            ("x: smallint", err),
-            ("x: int", err),
-            ("x: bigint", err),
-            ("x: string", err),
-            ("x: date", err),
-            ("x: timestamp", err),
-            ("x: byte", err),
-            ("x: binary", err),
-            ("x: float", err),
-            ("x: double", err),
-            ("x: decimal(10, 0)", err),
-            ("x: array<string>", [Row(x=["a", "b", "c"])]),
-            ("x: map<string,string>", err),
-            ("x: struct<a:string,b:string,c:string>", [Row(Row(a="0", b="1.1", c="2"))]),
-            ("x: struct<a:int,b:int,c:int>", [Row(Row(a=0, b=1, c=2))]),
-            ("x: struct<a:float,b:float,c:float>", [Row(Row(a=0, b=1.1, c=2))]),
-            ("x: struct<a:struct<>,b:struct<>,c:struct<>>", err),
-        ]:
-            with self.subTest(ret_type=ret_type):
-                self._check_result_or_exception(TestUDTF, ret_type, expected)
-
-    def test_struct_output_type_casting_row(self):
-        class TestUDTF:
-            def eval(self):
-                yield Row(a=0, b=1.1, c=2),
-
-        err = "UDTF_ARROW_TYPE_CAST_ERROR"
-
-        for ret_type, expected in [
-            ("x: boolean", err),
-            ("x: tinyint", err),
-            ("x: smallint", err),
-            ("x: int", err),
-            ("x: bigint", err),
-            ("x: string", err),
-            ("x: date", err),
-            ("x: timestamp", err),
-            ("x: byte", err),
-            ("x: binary", err),
-            ("x: float", err),
-            ("x: double", err),
-            ("x: decimal(10, 0)", err),
-            ("x: array<string>", [Row(x=["0", "1.1", "2"])]),
-            ("x: map<string,string>", err),
-            ("x: struct<a:string,b:string,c:string>", [Row(Row(a="0", b="1.1", c="2"))]),
-            ("x: struct<a:int,b:int,c:int>", [Row(Row(a=0, b=1, c=2))]),
-            ("x: struct<a:float,b:float,c:float>", [Row(Row(a=0, b=1.1, c=2))]),
-            ("x: struct<a:struct<>,b:struct<>,c:struct<>>", err),
-        ]:
-            with self.subTest(ret_type=ret_type):
-                self._check_result_or_exception(TestUDTF, ret_type, expected)
-
-    def test_inconsistent_output_types(self):
-        class TestUDTF:
-            def eval(self):
-                yield 1,
-                yield [1, 2],
-
-        for ret_type in [
-            "x: int",
-            "x: array<int>",
-        ]:
-            with self.subTest(ret_type=ret_type):
-                with self.assertRaisesRegex(PythonException, "UDTF_ARROW_TYPE_CAST_ERROR"):
-                    udtf(TestUDTF, returnType=ret_type)().collect()
-
-
-class UDTFArrowTests(UDTFArrowTestsMixin, ReusedSQLTestCase):
-    @classmethod
-    def setUpClass(cls):
-        super(UDTFArrowTests, cls).setUpClass()
-        cls.spark.conf.set("spark.sql.execution.pythonUDTF.arrow.enabled", "true")
-
-    @classmethod
-    def tearDownClass(cls):
-        try:
-            cls.spark.conf.unset("spark.sql.execution.pythonUDTF.arrow.enabled")
-        finally:
-            super(UDTFArrowTests, cls).tearDownClass()
-
-
-if __name__ == "__main__":
-    from pyspark.sql.tests.test_udtf import *  # noqa: F401
-
-    try:
-        import xmlrunner  # type: ignore
-
-        testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
-    except ImportError:
-        testRunner = None
-    unittest.main(testRunner=testRunner, verbosity=2)