python/pyspark/sql/tests/test_udtf.py - master - githubmirror/spark

python/pyspark/sql/tests/test_udtf.py
4 014 строк144 KB
Linhong Liu
[SPARK-57271][PYTHON] Propagate traceback locals to Python planner runner
29 июл 2026, 15:52
29 июл 2026, 15:5287dbc24
#
# 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.
#

from decimal import Decimal
import datetime
import os
import shutil
import tempfile
import unittest
import logging
import time
from dataclasses import dataclass
from typing import Iterator, Optional

from pyspark.errors import (
    PySparkAttributeError,
    PythonException,
    PySparkTypeError,
    AnalysisException,
    PySparkPicklingError,
    IllegalArgumentException,
)
from pyspark.util import PythonEvalType
from pyspark.sql.functions import (
    array,
    col,
    create_map,
    lit,
    named_struct,
    udf,
    udtf,
    AnalyzeArgument,
    AnalyzeResult,
    OrderingColumn,
    PartitioningColumn,
    SelectedColumn,
    SkipRestOfInputTableException,
)
from pyspark.sql.types import (
    ArrayType,
    BooleanType,
    DataType,
    IntegerType,
    LongType,
    MapType,
    NullType,
    Row,
    StringType,
    StructField,
    StructType,
    VariantVal,
)
from pyspark.logger import PySparkLogger
from pyspark.testing import assertDataFrameEqual, assertSchemaEqual
from pyspark.testing.objects import ExamplePoint, ExamplePointUDT
from pyspark.testing.sqlutils import ReusedSQLTestCase
from pyspark.testing.utils import (
    have_pandas,
    have_pyarrow,
    pandas_requirement_message,
    pyarrow_requirement_message,
)
from pyspark.util import is_remote_only


class BaseUDTFTestsMixin:
    def test_simple_udtf(self):
        class TestUDTF:
            def eval(self):
                yield "hello", "world"

        func = udtf(TestUDTF, returnType="c1: string, c2: string")
        assertDataFrameEqual(func(), [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")
        assertDataFrameEqual(func(lit(1)), [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")
        assertDataFrameEqual(func(lit(1)), [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")
        assertDataFrameEqual(func(lit(1)), [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")
        res = func(lit(1), lit(2))
        assertDataFrameEqual(res, [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

        assertDataFrameEqual(TestUDTF(lit(1)), [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)")
        assertDataFrameEqual(df, [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"
        )
        schema = StructType(
            [
                StructField("a", IntegerType(), True),
                StructField("b", IntegerType(), True),
                StructField("c", IntegerType(), True),
            ]
        )
        expected = self.spark.createDataFrame(
            [(0, 1, 1), (0, 1, -1), (1, 2, 3), (1, 2, -1)], schema=schema
        )
        assertDataFrameEqual(df, expected)

    def test_udtf_with_lateral_join_dataframe(self):
        @udtf(returnType="a: int, b: int, c: int")
        class TestUDTF:
            def eval(self, a: int, b: int) -> Iterator:
                yield a, b, a + b
                yield a, b, a - b

        self.spark.udtf.register("testUDTF", TestUDTF)

        for i, df in enumerate(
            [
                self.spark.sql("values (0, 1), (1, 2) t(a, b)").lateralJoin(
                    TestUDTF(col("a").outer(), col("b").outer())
                ),
                self.spark.sql("values (0, 1), (1, 2) t(a, b)").lateralJoin(
                    TestUDTF(a=col("a").outer(), b=col("b").outer())
                ),
                self.spark.sql("values (0, 1), (1, 2) t(a, b)").lateralJoin(
                    TestUDTF(b=col("b").outer(), a=col("a").outer())
                ),
            ]
        ):
            with self.subTest(query_no=i):
                assertDataFrameEqual(
                    df,
                    self.spark.sql(
                        "SELECT * FROM values (0, 1), (1, 2) t(a, b), LATERAL testUDTF(a, b)"
                    ),
                )

        @udtf(returnType="a: int")
        class TestUDTF:
            def eval(self):
                yield (1,)
                yield (2,)

        assertDataFrameEqual(
            self.spark.range(3, numPartitions=1).lateralJoin(TestUDTF()),
            [
                Row(id=0, a=1),
                Row(id=0, a=2),
                Row(id=1, a=1),
                Row(id=1, a=2),
                Row(id=2, a=1),
                Row(id=2, a=2),
            ],
        )

        @udtf(returnType="a: int")
        class TestUDTF:
            def eval(self, i: int):
                for n in range(i):
                    yield (n,)

        assertDataFrameEqual(
            self.spark.range(3, numPartitions=1).lateralJoin(TestUDTF(col("id").outer())),
            [
                Row(id=1, a=0),
                Row(id=2, a=0),
                Row(id=2, a=1),
            ],
        )

    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")
        res = func(lit(1), lit(2))
        assertDataFrameEqual(res, [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()

        @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()

        @udtf(returnType=StructType().add("point", ExamplePointUDT()))
        class TestUDTF:
            def eval(self, x: float, y: float):
                yield ExamplePoint(x=x * 10, y=y * 10)

        with self.assertRaisesRegex(PythonException, "UDTF_INVALID_OUTPUT_ROW_TYPE"):
            TestUDTF(lit(1.0), lit(2.0)).collect()

    def test_udtf_eval_returning_tuple_with_struct_type(self):
        @udtf(returnType="a: struct<b: int, c: int>")
        class TestUDTF:
            def eval(self, a: int):
                yield ((a, a + 1),)

        assertDataFrameEqual(TestUDTF(lit(1)), [Row(a=Row(b=1, c=2))])

        @udtf(returnType="a: struct<b: int, c: int>")
        class TestUDTF:
            def eval(self, a: int):
                yield a, a + 1

        with self.assertRaisesRegex(PythonException, "UDTF_RETURN_SCHEMA_MISMATCH"):
            TestUDTF(lit(1)).collect()

    def test_udtf_eval_returning_udt(self):
        @udtf(returnType=StructType().add("point", ExamplePointUDT()))
        class TestUDTF:
            def eval(self, x: float, y: float):
                yield (ExamplePoint(x=x * 10, y=y * 10),)

        assertDataFrameEqual(
            TestUDTF(lit(1.0), lit(2.0)), [Row(point=ExamplePoint(x=10.0, y=20.0))]
        )

    def test_udtf_eval_taking_udt(self):
        @udtf(returnType="x: double, y: double")
        class TestUDTF:
            def eval(self, point: ExamplePoint):
                yield point.x * 10, point.y * 10

        df = self.spark.createDataFrame(
            [(ExamplePoint(x=1.0, y=2.0),)], schema=StructType().add("point", ExamplePointUDT())
        )
        assertDataFrameEqual(
            df.lateralJoin(TestUDTF(col("point").outer())),
            [Row(point=ExamplePoint(x=1.0, y=2.0), x=10.0, y=20.0)],
        )

    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): ...

        assertDataFrameEqual(TestUDTF(lit(1)), [])

        @udtf(returnType="a: int")
        class TestUDTF:
            def eval(self, a: int):
                return

        assertDataFrameEqual(TestUDTF(lit(1)), [])

    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)
        assertDataFrameEqual(
            self.spark.sql("SELECT * FROM range(0, 8) JOIN LATERAL test_udtf(id)"),
            [Row(id=6, a=6), Row(id=7, a=7)],
        )

    def test_udtf_with_conditional_return_dataframe(self):
        @udtf(returnType="a: int")
        class TestUDTF:
            def eval(self, a: int):
                if a > 5:
                    yield (a,)

        self.spark.udtf.register("test_udtf", TestUDTF)

        assertDataFrameEqual(
            self.spark.range(8).lateralJoin(TestUDTF(col("id").outer())),
            self.spark.sql("SELECT * FROM range(0, 8) JOIN LATERAL test_udtf(id)"),
        )

    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,)

        assertDataFrameEqual(TestUDTF(lit(1)), [Row(a=1), Row(a=None)])
        df = self.spark.createDataFrame([(0, 1), (1, 2)], schema=["a", "b"])
        assertDataFrameEqual(TestUDTF(lit(1)).join(df, "a", "inner"), [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,)

        assertDataFrameEqual(TestUDTF(lit(None)), [Row(a=None)])
        self.spark.udtf.register("testUDTF", TestUDTF)
        df = self.spark.sql("SELECT * FROM testUDTF(null)")
        assertDataFrameEqual(df, [Row(a=None)])

    # These are expected error message substrings to be used in test cases below.
    tooManyPositionalArguments = "too many positional arguments"
    missingARequiredArgument = "missing a required argument"
    multipleValuesForArgument = "multiple values for argument"

    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, BaseUDTFTestsMixin.missingARequiredArgument):
            TestUDTF().collect()

        with self.assertRaisesRegex(PythonException, BaseUDTFTestsMixin.tooManyPositionalArguments):
            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".*constructor has more than one argument.*"):
            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):
        # 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, "UDTF_RETURN_SCHEMA_MISMATCH"):
            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, "UDTF_RETURN_SCHEMA_MISMATCH"):
            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

        res = TestUDTF(lit(1))
        assertDataFrameEqual(res, [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

        assertDataFrameEqual(
            TestUDTF(lit(1)),
            [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'"
        )
        assertDataFrameEqual(
            df,
            [
                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_udtf_cleanup_with_exception_in_eval(self):
        with tempfile.TemporaryDirectory(prefix="test_udtf_cleanup_with_exception_in_eval") as d:
            path = os.path.join(d, "file.txt")

            @udtf(returnType="x: int")
            class TestUDTF:
                def __init__(self):
                    self.path = path

                def eval(self, x: int):
                    raise Exception("eval error")

                def terminate(self):
                    with open(self.path, "a") as f:
                        f.write("terminate")

                def cleanup(self):
                    with open(self.path, "a") as f:
                        f.write("cleanup")

            with self.assertRaisesRegex(PythonException, "eval error"):
                TestUDTF(lit(1)).show()

            with open(path, "r") as f:
                data = f.read()

            # Only the cleanup method should be called, not terminate. The UDTF may be retried,
            # so cleanup can run more than once and the file may contain multiple "cleanup"
            # strings; just check that "cleanup" is present and "terminate" is not.
            self.assertIn("cleanup", data)
            self.assertNotIn("terminate", data)

    def test_udtf_cleanup_with_exception_in_terminate(self):
        with tempfile.TemporaryDirectory(
            prefix="test_udtf_cleanup_with_exception_in_terminate"
        ) as d:
            path = os.path.join(d, "file.txt")

            @udtf(returnType="x: int")
            class TestUDTF:
                def __init__(self):
                    self.path = path

                def eval(self, x: int):
                    yield (x,)

                def terminate(self):
                    raise Exception("terminate error")

                def cleanup(self):
                    with open(self.path, "a") as f:
                        f.write("cleanup")

            with self.assertRaisesRegex(PythonException, "terminate error"):
                TestUDTF(lit(1)).show()

            with open(path, "r") as f:
                data = f.read()

            # The cleanup method should be called even when terminate raises. The UDTF may be
            # retried, so cleanup can run more than once and the file may contain multiple
            # "cleanup" strings; just check that "cleanup" is present.
            self.assertIn("cleanup", data)

    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):
        @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, "UDTF_RETURN_SCHEMA_MISMATCH"):
            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, "UDTF_RETURN_SCHEMA_MISMATCH"):
            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,
            errorClass="UDTF_RETURN_TYPE_MISMATCH",
            messageParameters={"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,
            errorClass="UDTF_RETURN_TYPE_MISMATCH",
            messageParameters={
                "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)
        assertDataFrameEqual(
            self.spark.sql("select * from test_udtf(named_struct('name', 'Alice', 'age', 1))"),
            [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.assertIn(
            self.spark.sql("select * from test_udtf(array(1, 2, 3))").collect(),
            [[Row(x="[1, 2, 3]")], [Row(x="[np.int32(1), np.int32(2), np.int32(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)
        assertDataFrameEqual(
            self.spark.sql("select * from test_udtf(map('key', 'value'))"),
            [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},)

        assertDataFrameEqual(TestUDTF(lit(1)), [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],)

        assertDataFrameEqual(TestUDTF(lit(1)), [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)},)

        assertDataFrameEqual(TestUDTF(lit(1)), [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_udtf_nullable_check(self):
        for ret_type, value, expected in [
            (
                StructType([StructField("value", ArrayType(IntegerType(), False))]),
                ([None],),
                "PySparkRuntimeError",
            ),
            (
                StructType([StructField("value", ArrayType(IntegerType(), True))]),
                ([None],),
                [Row(value=[None])],
            ),
            (
                StructType([StructField("value", MapType(StringType(), IntegerType(), False))]),
                ({"a": None},),
                "PySparkRuntimeError",
            ),
            (
                StructType([StructField("value", MapType(StringType(), IntegerType(), True))]),
                ({"a": None},),
                [Row(value={"a": None})],
            ),
            (
                StructType([StructField("value", MapType(StringType(), IntegerType(), True))]),
                ({None: 1},),
                "PySparkRuntimeError",
            ),
            (
                StructType([StructField("value", MapType(StringType(), IntegerType(), False))]),
                ({None: 1},),
                "PySparkRuntimeError",
            ),
            (
                StructType(
                    [
                        StructField(
                            "value", MapType(StringType(), ArrayType(IntegerType(), False), False)
                        )
                    ]
                ),
                ({"s": [None]},),
                "PySparkRuntimeError",
            ),
            (
                StructType(
                    [
                        StructField(
                            "value",
                            MapType(
                                StructType([StructField("value", StringType(), False)]),
                                IntegerType(),
                                False,
                            ),
                        )
                    ]
                ),
                ({(None,): 1},),
                "PySparkRuntimeError",
            ),
            (
                StructType(
                    [
                        StructField(
                            "value",
                            MapType(
                                StructType([StructField("value", StringType(), False)]),
                                IntegerType(),
                                True,
                            ),
                        )
                    ]
                ),
                ({(None,): 1},),
                "PySparkRuntimeError",
            ),
            (
                StructType(
                    [StructField("value", StructType([StructField("value", StringType(), False)]))]
                ),
                ((None,),),
                "PySparkRuntimeError",
            ),
            (
                StructType(
                    [
                        StructField(
                            "value",
                            StructType(
                                [StructField("value", ArrayType(StringType(), False), False)]
                            ),
                        )
                    ]
                ),
                (([None],),),
                "PySparkRuntimeError",
            ),
        ]:

            class TestUDTF:
                def eval(self):
                    yield value

            with self.subTest(ret_type=ret_type, value=value):
                self._check_result_or_exception(TestUDTF, ret_type, expected)

    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):
        from py4j.protocol import Py4JJavaError

        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,
            errorClass="CANNOT_REGISTER_UDTF",
            messageParameters={
                "name": "test_udf",
            },
        )

        class TestUDTF: ...

        with self.assertRaises(PySparkTypeError) as e:
            self.spark.udtf.register("test_udtf", TestUDTF)

        self.check_error(
            exception=e.exception,
            errorClass="CANNOT_REGISTER_UDTF",
            messageParameters={
                "name": "test_udtf",
            },
        )

    def test_udtf_pickle_error(self):
        with tempfile.TemporaryDirectory(prefix="test_udtf_pickle_error") 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(PySparkPicklingError, "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(PySparkPicklingError, "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,
            errorClass="INVALID_UDTF_NO_EVAL",
            messageParameters={"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,
            errorClass="INVALID_UDTF_HANDLER_TYPE",
            messageParameters={"type": "function"},
        )

        with self.assertRaises(PySparkTypeError) as e:
            udtf(1, returnType="a: int")

        self.check_error(
            exception=e.exception,
            errorClass="INVALID_UDTF_HANDLER_TYPE",
            messageParameters={"type": "int"},
        )

    def test_udtf_with_table_argument_query(self):
        func = self.udtf_for_table_argument()
        self.spark.udtf.register("test_udtf", func)
        assertDataFrameEqual(
            self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)))"),
            [Row(a=6), Row(a=7)],
        )

    def test_df_asTable(self):
        func = self.udtf_for_table_argument()
        self.spark.udtf.register("test_udtf", func)
        df = self.spark.range(8)
        assertDataFrameEqual(
            func(df.asTable()),
            self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)))"),
        )

    def udtf_for_table_argument(self):
        class TestUDTF:
            def eval(self, row: Row):
                if row["id"] > 5:
                    yield (row["id"],)

        func = udtf(TestUDTF, returnType="a: int")
        return func

    def test_df_asTable_chaining_methods(self):
        class TestUDTF:
            def eval(self, row: Row):
                yield row["key"], row["value"]

        func = udtf(TestUDTF, returnType="key: int, value: string")
        df = self.spark.createDataFrame(
            [(1, "a", 3), (1, "b", 3), (2, "c", 4), (2, "d", 4)], ["key", "value", "number"]
        )
        assertDataFrameEqual(
            func(df.asTable().partitionBy("key").orderBy(df.value)),
            [
                Row(key=1, value="a"),
                Row(key=1, value="b"),
                Row(key=2, value="c"),
                Row(key=2, value="d"),
            ],
            checkRowOrder=True,
        )
        assertDataFrameEqual(
            func(df.asTable().partitionBy(["key", "number"]).orderBy(df.value)),
            [
                Row(key=1, value="a"),
                Row(key=1, value="b"),
                Row(key=2, value="c"),
                Row(key=2, value="d"),
            ],
            checkRowOrder=True,
        )
        assertDataFrameEqual(
            func(df.asTable().partitionBy("key").orderBy(df.value.desc())),
            [
                Row(key=1, value="b"),
                Row(key=1, value="a"),
                Row(key=2, value="d"),
                Row(key=2, value="c"),
            ],
            checkRowOrder=True,
        )
        assertDataFrameEqual(
            func(df.asTable().partitionBy("key").orderBy(["number", "value"])),
            [
                Row(key=1, value="a"),
                Row(key=1, value="b"),
                Row(key=2, value="c"),
                Row(key=2, value="d"),
            ],
            checkRowOrder=True,
        )
        assertDataFrameEqual(
            func(df.asTable().withSinglePartition()),
            [
                Row(key=1, value="a"),
                Row(key=1, value="b"),
                Row(key=2, value="c"),
                Row(key=2, value="d"),
            ],
        )

        assertDataFrameEqual(
            func(df.asTable().withSinglePartition().orderBy("value")),
            [
                Row(key=1, value="a"),
                Row(key=1, value="b"),
                Row(key=2, value="c"),
                Row(key=2, value="d"),
            ],
        )

        with self.assertRaisesRegex(
            IllegalArgumentException,
            r"Cannot call withSinglePartition\(\) after partitionBy\(\)"
            r" or withSinglePartition\(\) has been called",
        ):
            df.asTable().partitionBy(df.key).withSinglePartition()

        with self.assertRaisesRegex(
            IllegalArgumentException,
            r"Cannot call partitionBy\(\) after partitionBy\(\)"
            r" or withSinglePartition\(\) has been called",
        ):
            df.asTable().withSinglePartition().partitionBy(df.key)

        with self.assertRaisesRegex(
            IllegalArgumentException,
            r"Please call partitionBy\(\) or withSinglePartition\(\) before orderBy\(\)",
        ):
            df.asTable().orderBy(df.key)

        with self.assertRaisesRegex(
            IllegalArgumentException,
            r"Please call partitionBy\(\) or withSinglePartition\(\) before orderBy\(\)",
        ):
            df.asTable().partitionBy().orderBy(df.key)

        with self.assertRaisesRegex(
            IllegalArgumentException,
            r"Cannot call partitionBy\(\) after partitionBy\(\)"
            r" or withSinglePartition\(\) has been called",
        ):
            df.asTable().partitionBy(df.key).partitionBy()

    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)
        assertDataFrameEqual(
            self.spark.sql("SELECT * FROM test_udtf(5, TABLE (SELECT id FROM range(0, 8)))"),
            [Row(a=6), Row(a=7)],
        )

    def test_udtf_with_table_argument_identifier(self):
        func = self.udtf_for_table_argument()
        self.spark.udtf.register("test_udtf", func)

        with self.temp_view("v"):
            self.spark.sql("CREATE OR REPLACE TEMPORARY VIEW v as SELECT id FROM range(0, 8)")
            assertDataFrameEqual(
                self.spark.sql("SELECT * FROM test_udtf(TABLE (v))"),
                [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.temp_view("v"):
            self.spark.sql("CREATE OR REPLACE TEMPORARY VIEW v as SELECT id FROM range(0, 8)")
            assertDataFrameEqual(
                self.spark.sql("SELECT * FROM test_udtf(5, TABLE (v))"),
                [Row(a=6), Row(a=7)],
            )

    def test_udtf_with_table_argument_unknown_identifier(self):
        func = self.udtf_for_table_argument()
        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):
        func = self.udtf_for_table_argument()
        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):
        func = self.udtf_for_table_argument()
        self.spark.udtf.register("test_udtf", func)
        assertDataFrameEqual(
            self.spark.sql("""
                SELECT * FROM test_udtf(TABLE (
                  WITH t AS (
                    SELECT id FROM range(0, 8)
                  )
                  SELECT * FROM t
                ))
                """),
            [Row(a=6), Row(a=7)],
        )

    def test_udtf_with_table_argument_cte_outside(self):
        func = self.udtf_for_table_argument()
        self.spark.udtf.register("test_udtf", func)
        assertDataFrameEqual(
            self.spark.sql("""
                WITH t AS (
                  SELECT id FROM range(0, 8)
                )
                SELECT * FROM test_udtf(TABLE (SELECT id FROM t))
                """),
            [Row(a=6), Row(a=7)],
        )

        assertDataFrameEqual(
            self.spark.sql("""
                WITH t AS (
                  SELECT id FROM range(0, 8)
                )
                SELECT * FROM test_udtf(TABLE (t))
                """),
            [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):
        func = self.udtf_for_table_argument()
        self.spark.udtf.register("test_udtf", func)
        assertDataFrameEqual(
            self.spark.sql("""
                SELECT * FROM
                  range(0, 8) AS t,
                  LATERAL test_udtf(TABLE (t))
                """),
            [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}):
            assertDataFrameEqual(
                self.spark.sql(query),
                [
                    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"""
                ...

            @staticmethod
            def analyze(x: AnalyzeArgument) -> AnalyzeResult:
                """Analyze the argument."""
                ...

            def eval(self, x: int):
                """Evaluate the input row."""
                yield (x + 1,)

            def terminate(self):
                """Terminate the UDTF."""
                ...

        cls = udtf(TestUDTF).func
        self.assertIn("A UDTF for test", cls.__doc__)
        self.assertIn("Initialize the UDTF", cls.__init__.__doc__)
        self.assertIn("Analyze the argument", cls.analyze.__doc__)
        self.assertIn("Evaluate the input row", cls.eval.__doc__)
        self.assertIn("Terminate the UDTF", cls.terminate.__doc__)

    def test_simple_udtf_with_analyze(self):
        class TestUDTF:
            @staticmethod
            def analyze() -> AnalyzeResult:
                return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))

            def eval(self):
                yield "hello", "world"

        func = udtf(TestUDTF)
        self.spark.udtf.register("test_udtf", func)

        expected = [Row(c1="hello", c2="world")]
        assertDataFrameEqual(func(), expected)
        assertDataFrameEqual(self.spark.sql("SELECT * FROM test_udtf()"), expected)

    def test_udtf_with_analyze(self):
        class TestUDTF:
            @staticmethod
            def analyze(a: AnalyzeArgument) -> AnalyzeResult:
                assert isinstance(a, AnalyzeArgument)
                assert isinstance(a.dataType, DataType)
                assert a.value is not None
                assert a.isTable is False
                assert a.isConstantExpression is True
                return AnalyzeResult(StructType().add("a", a.dataType))

            def eval(self, a):
                yield (a,)

        func = udtf(TestUDTF)
        self.spark.udtf.register("test_udtf", func)

        for i, (df, expected_schema, expected_results) in enumerate(
            [
                (func(lit(1)), StructType().add("a", IntegerType()), [Row(a=1)]),
                # another data type
                (func(lit("x")), StructType().add("a", StringType()), [Row(a="x")]),
                # array type
                (
                    func(array(lit(1), lit(2), lit(3))),
                    StructType().add("a", ArrayType(IntegerType(), containsNull=False)),
                    [Row(a=[1, 2, 3])],
                ),
                # map type
                (
                    func(create_map(lit("x"), lit(1), lit("y"), lit(2))),
                    StructType().add(
                        "a", MapType(StringType(), IntegerType(), valueContainsNull=False)
                    ),
                    [Row(a={"x": 1, "y": 2})],
                ),
                # struct type
                (
                    func(named_struct(lit("x"), lit(1), lit("y"), lit(2))),
                    StructType().add(
                        "a",
                        StructType()
                        .add("x", IntegerType(), nullable=False)
                        .add("y", IntegerType(), nullable=False),
                    ),
                    [Row(a=Row(x=1, y=2))],
                ),
                # use SQL
                (
                    self.spark.sql("SELECT * from test_udtf(1)"),
                    StructType().add("a", IntegerType()),
                    [Row(a=1)],
                ),
            ]
        ):
            with self.subTest(query_no=i):
                assertSchemaEqual(df.schema, expected_schema)
                assertDataFrameEqual(df, expected_results)

    def test_udtf_with_analyze_decorator(self):
        @udtf
        class TestUDTF:
            @staticmethod
            def analyze() -> AnalyzeResult:
                return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))

            def eval(self):
                yield "hello", "world"

        self.spark.udtf.register("test_udtf", TestUDTF)

        expected = [Row(c1="hello", c2="world")]
        assertDataFrameEqual(TestUDTF(), expected)
        assertDataFrameEqual(self.spark.sql("SELECT * FROM test_udtf()"), expected)

    def test_udtf_with_analyze_decorator_parens(self):
        @udtf()
        class TestUDTF:
            @staticmethod
            def analyze() -> AnalyzeResult:
                return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))

            def eval(self):
                yield "hello", "world"

        self.spark.udtf.register("test_udtf", TestUDTF)

        expected = [Row(c1="hello", c2="world")]
        assertDataFrameEqual(TestUDTF(), expected)
        assertDataFrameEqual(self.spark.sql("SELECT * FROM test_udtf()"), expected)

    def test_udtf_with_analyze_multiple_arguments(self):
        class TestUDTF:
            @staticmethod
            def analyze(a: AnalyzeArgument, b: AnalyzeArgument) -> AnalyzeResult:
                return AnalyzeResult(StructType().add("a", a.dataType).add("b", b.dataType))

            def eval(self, a, b):
                yield a, b

        func = udtf(TestUDTF)
        self.spark.udtf.register("test_udtf", func)

        for i, (df, expected_schema, expected_results) in enumerate(
            [
                (
                    func(lit(1), lit("x")),
                    StructType().add("a", IntegerType()).add("b", StringType()),
                    [Row(a=1, b="x")],
                ),
                (
                    self.spark.sql("SELECT * FROM test_udtf(1, 'x')"),
                    StructType().add("a", IntegerType()).add("b", StringType()),
                    [Row(a=1, b="x")],
                ),
            ]
        ):
            with self.subTest(query_no=i):
                assertSchemaEqual(df.schema, expected_schema)
                assertDataFrameEqual(df, expected_results)

    def test_udtf_with_analyze_arbitrary_number_arguments(self):
        class TestUDTF:
            @staticmethod
            def analyze(*args: AnalyzeArgument) -> AnalyzeResult:
                return AnalyzeResult(
                    StructType([StructField(f"col{i}", a.dataType) for i, a in enumerate(args)])
                )

            def eval(self, *args):
                yield args

        func = udtf(TestUDTF)
        self.spark.udtf.register("test_udtf", func)

        for i, (df, expected_schema, expected_results) in enumerate(
            [
                (
                    func(lit(1)),
                    StructType().add("col0", IntegerType()),
                    [Row(a=1)],
                ),
                (
                    self.spark.sql("SELECT * FROM test_udtf(1, 'x')"),
                    StructType().add("col0", IntegerType()).add("col1", StringType()),
                    [Row(a=1, b="x")],
                ),
                (func(), StructType(), [Row()]),
            ]
        ):
            with self.subTest(query_no=i):
                assertSchemaEqual(df.schema, expected_schema)
                assertDataFrameEqual(df, expected_results)

    def test_udtf_with_analyze_table_argument(self):
        class TestUDTF:
            @staticmethod
            def analyze(a: AnalyzeArgument) -> AnalyzeResult:
                assert isinstance(a, AnalyzeArgument)
                assert isinstance(a.dataType, StructType)
                assert a.value is None
                assert a.isTable is True
                assert a.isConstantExpression is False
                return AnalyzeResult(StructType().add("a", a.dataType[0].dataType))

            def eval(self, a: Row):
                if a["id"] > 5:
                    yield (a["id"],)

        func = udtf(TestUDTF)
        self.spark.udtf.register("test_udtf", func)

        df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)))")
        assertSchemaEqual(df.schema, StructType().add("a", LongType()))
        assertDataFrameEqual(df, [Row(a=6), Row(a=7)])

    def test_udtf_with_analyze_table_argument_adding_columns(self):
        class TestUDTF:
            @staticmethod
            def analyze(a: AnalyzeArgument) -> AnalyzeResult:
                assert isinstance(a.dataType, StructType)
                assert a.isTable is True
                assert a.isConstantExpression is False
                return AnalyzeResult(a.dataType.add("is_even", BooleanType()))

            def eval(self, a: Row):
                yield a["id"], a["id"] % 2 == 0

        func = udtf(TestUDTF)
        self.spark.udtf.register("test_udtf", func)

        df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 4)))")
        assertSchemaEqual(
            df.schema,
            StructType().add("id", LongType(), nullable=False).add("is_even", BooleanType()),
        )
        assertDataFrameEqual(
            df,
            [
                Row(a=0, is_even=True),
                Row(a=1, is_even=False),
                Row(a=2, is_even=True),
                Row(a=3, is_even=False),
            ],
        )

    def test_udtf_with_analyze_table_argument_repeating_rows(self):
        class TestUDTF:
            @staticmethod
            def analyze(n, row) -> AnalyzeResult:
                if n.value is None or not isinstance(n.value, int) or (n.value < 1 or n.value > 10):
                    raise Exception("The first argument must be a scalar integer between 1 and 10")

                if row.isTable is False:
                    raise Exception("The second argument must be a table argument")

                assert isinstance(row.dataType, StructType)
                return AnalyzeResult(row.dataType)

            def eval(self, n: int, row: Row):
                for _ in range(n):
                    yield row

        func = udtf(TestUDTF)
        self.spark.udtf.register("test_udtf", func)

        expected_schema = StructType().add("id", LongType(), nullable=False)
        expected_results = [
            Row(a=0),
            Row(a=0),
            Row(a=1),
            Row(a=1),
            Row(a=2),
            Row(a=2),
            Row(a=3),
            Row(a=3),
        ]
        for i, df in enumerate(
            [
                self.spark.sql("SELECT * FROM test_udtf(2, TABLE (SELECT id FROM range(0, 4)))"),
                self.spark.sql(
                    "SELECT * FROM test_udtf(1 + 1, TABLE (SELECT id FROM range(0, 4)))"
                ),
            ]
        ):
            with self.subTest(query_no=i):
                assertSchemaEqual(df.schema, expected_schema)
                assertDataFrameEqual(df, expected_results)

        with self.assertRaisesRegex(
            AnalysisException, "The first argument must be a scalar integer between 1 and 10"
        ):
            self.spark.sql(
                "SELECT * FROM test_udtf(0, TABLE (SELECT id FROM range(0, 4)))"
            ).collect()

        with (
            self.sql_conf({"spark.sql.tvf.allowMultipleTableArguments.enabled": True}),
            self.assertRaisesRegex(
                AnalysisException, "The first argument must be a scalar integer between 1 and 10"
            ),
        ):
            self.spark.sql("""
                SELECT * FROM test_udtf(
                  TABLE (SELECT id FROM range(0, 1)),
                  TABLE (SELECT id FROM range(0, 4)))
                """).collect()

        with self.assertRaisesRegex(
            AnalysisException, "The second argument must be a table argument"
        ):
            self.spark.sql("SELECT * FROM test_udtf(1, 'x')").collect()

    def test_udtf_with_analyze_table_select(self):
        @udtf
        class TestUDTF:
            @staticmethod
            def analyze(*args, **kwargs) -> AnalyzeResult:
                return AnalyzeResult(
                    StructType().add("id", IntegerType()), select=[SelectedColumn("id")]
                )

            def eval(self, row: Row):
                assert "value" not in row
                yield (row["id"],)

        df = self.spark.createDataFrame([(1, "a"), (2, "b"), (3, "c")], ["id", "value"])
        assertDataFrameEqual(
            TestUDTF(df.asTable()).collect(),
            [Row(id=1), Row(id=2), Row(id=3)],
        )

    def test_udtf_with_both_return_type_and_analyze(self):
        class TestUDTF:
            @staticmethod
            def analyze() -> AnalyzeResult:
                return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))

            def eval(self):
                yield "hello", "world"

        with self.assertRaises(PySparkAttributeError) as e:
            udtf(TestUDTF, returnType="c1: string, c2: string")

        self.check_error(
            exception=e.exception,
            errorClass="INVALID_UDTF_BOTH_RETURN_TYPE_AND_ANALYZE",
            messageParameters={"name": "TestUDTF"},
        )

    def test_udtf_with_neither_return_type_nor_analyze(self):
        class TestUDTF:
            def eval(self):
                yield "hello", "world"

        with self.assertRaises(PySparkAttributeError) as e:
            udtf(TestUDTF)

        self.check_error(
            exception=e.exception,
            errorClass="INVALID_UDTF_RETURN_TYPE",
            messageParameters={"name": "TestUDTF"},
        )

    def test_udtf_with_analyze_non_staticmethod(self):
        class TestUDTF:
            def analyze(self) -> AnalyzeResult:
                return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))

            def eval(self):
                yield "hello", "world"

        with self.assertRaises(PySparkAttributeError) as e:
            udtf(TestUDTF)

        self.check_error(
            exception=e.exception,
            errorClass="INVALID_UDTF_RETURN_TYPE",
            messageParameters={"name": "TestUDTF"},
        )

        with self.assertRaises(PySparkAttributeError) as e:
            udtf(TestUDTF, returnType="c1: string, c2: string")

        self.check_error(
            exception=e.exception,
            errorClass="INVALID_UDTF_BOTH_RETURN_TYPE_AND_ANALYZE",
            messageParameters={"name": "TestUDTF"},
        )

    def test_udtf_with_scalar_analyze_returning_wrong_result(self):
        # (wrong_type, error_message_regex)
        invalid_results = [
            (StringType(), r".*AnalyzeResult.*StringType.*"),
            (AnalyzeResult(StringType()), r".*AnalyzeResult.*schema.*"),
            (
                AnalyzeResult(StructType().add("a", StringType()), withSinglePartition=True),
                r".*withSinglePartition.*",
            ),
            (
                AnalyzeResult(
                    StructType().add("a", StringType()), partitionBy=[PartitioningColumn("a")]
                ),
                r".*partitionBy.*",
            ),
        ]

        for wrong_type, error_message_regex in invalid_results:
            with self.subTest(wrong_type=wrong_type):

                class TestUDTF:
                    @staticmethod
                    def analyze() -> AnalyzeResult:
                        return wrong_type

                    def eval(self):
                        yield "hello", "world"

                func = udtf(TestUDTF)

                with self.assertRaisesRegex(AnalysisException, error_message_regex):
                    func().collect()

    def test_udtf_with_table_analyze_returning_wrong_result(self):
        invalid_results = [
            (
                AnalyzeResult(
                    StructType().add("a", StringType()), partitionBy=[OrderingColumn("a")]
                ),
                r".*partitionBy.*",
            ),
            (
                AnalyzeResult(
                    StructType().add("a", StringType()), orderBy=[PartitioningColumn("a")]
                ),
                r".*orderBy.*",
            ),
            (
                AnalyzeResult(StructType().add("a", StringType()), select=SelectedColumn("a")),
                r".*select.*",
            ),
        ]

        for wrong_type, error_message_regex in invalid_results:
            with self.subTest(wrong_type=wrong_type):

                class TestUDTF:
                    @staticmethod
                    def analyze(**kwargs) -> AnalyzeResult:
                        return wrong_type

                    def eval(self, **kwargs):
                        yield tuple(value for _, value in sorted(kwargs.items()))

                func = udtf(TestUDTF)

                with self.assertRaisesRegex(AnalysisException, error_message_regex):
                    func(a=self.spark.range(3).asTable(), b=lit("x")).collect()

    def test_udtf_with_analyze_raising_an_exception(self):
        class TestUDTF:
            @staticmethod
            def analyze() -> AnalyzeResult:
                raise Exception("Failed to analyze.")

            def eval(self):
                yield "hello", "world"

        func = udtf(TestUDTF)

        with self.assertRaisesRegex(AnalysisException, "Failed to analyze."):
            func().collect()

    def test_udtf_analyze_traceback_with_locals(self):
        with self.sql_conf({"spark.sql.execution.pyspark.udf.tracebackWithLocals.enabled": True}):

            class TestUDTF:
                @staticmethod
                def analyze() -> AnalyzeResult:
                    local_marker = 1
                    if local_marker:
                        raise ValueError("boom")
                    return AnalyzeResult(StructType().add("x", StringType()))

                def eval(self):
                    yield ("x",)

            func = udtf(TestUDTF)

            with self.assertRaisesRegex(AnalysisException, "local_marker = 1"):
                func().collect()

    def test_udtf_with_analyze_null_literal(self):
        class TestUDTF:
            @staticmethod
            def analyze(a: AnalyzeArgument) -> AnalyzeResult:
                return AnalyzeResult(StructType().add("a", a.dataType))

            def eval(self, a):
                yield (a,)

        func = udtf(TestUDTF)

        df = func(lit(None))
        assertSchemaEqual(df.schema, StructType().add("a", NullType()))
        assertDataFrameEqual(df, [Row(a=None)])

    def test_udtf_with_analyze_taking_wrong_number_of_arguments(self):
        class TestUDTF:
            @staticmethod
            def analyze(a: AnalyzeArgument, b: AnalyzeArgument) -> AnalyzeResult:
                return AnalyzeResult(StructType().add("a", a.dataType).add("b", b.dataType))

            def eval(self, a, b):
                yield a, a + 1

        func = udtf(TestUDTF)

        with self.assertRaisesRegex(AnalysisException, r"arguments"):
            func(lit(1)).collect()

        with self.assertRaisesRegex(AnalysisException, r"arguments"):
            func(lit(1), lit(2), lit(3)).collect()

    def test_udtf_with_analyze_taking_keyword_arguments(self):
        @udtf
        class TestUDTF:
            @staticmethod
            def analyze(**kwargs) -> AnalyzeResult:
                return AnalyzeResult(StructType().add("a", StringType()).add("b", StringType()))

            def eval(self, **kwargs):
                yield "hello", "world"

        self.spark.udtf.register("test_udtf", TestUDTF)

        expected = [Row(c1="hello", c2="world")]
        assertDataFrameEqual(TestUDTF(), expected)
        assertDataFrameEqual(self.spark.sql("SELECT * FROM test_udtf()"), expected)
        assertDataFrameEqual(self.spark.sql("SELECT * FROM test_udtf(a=>1)"), expected)

        with self.assertRaisesRegex(
            AnalysisException, BaseUDTFTestsMixin.tooManyPositionalArguments
        ):
            TestUDTF(lit(1)).collect()

        with self.assertRaisesRegex(
            AnalysisException, BaseUDTFTestsMixin.tooManyPositionalArguments
        ):
            self.spark.sql("SELECT * FROM test_udtf(1, 'x')").collect()

    def test_udtf_with_analyze_using_broadcast(self):
        colname = self.sc.broadcast("col1")

        @udtf
        class TestUDTF:
            @staticmethod
            def analyze(a: AnalyzeArgument) -> AnalyzeResult:
                return AnalyzeResult(StructType().add(colname.value, a.dataType))

            def eval(self, a):
                assert colname.value == "col1"
                yield (a,)

            def terminate(self):
                assert colname.value == "col1"
                yield (100,)

        self.spark.udtf.register("test_udtf", TestUDTF)

        for i, df in enumerate([TestUDTF(lit(10)), self.spark.sql("SELECT * FROM test_udtf(10)")]):
            with self.subTest(query_no=i):
                assertSchemaEqual(df.schema, StructType().add("col1", IntegerType()))
                assertDataFrameEqual(df, [Row(col1=10), Row(col1=100)])

    def test_udtf_with_analyze_using_accumulator(self):
        test_accum = self.sc.accumulator(0)

        @udtf
        class TestUDTF:
            @staticmethod
            def analyze(a: AnalyzeArgument) -> AnalyzeResult:
                test_accum.add(1)
                return AnalyzeResult(StructType().add("col1", a.dataType))

            def eval(self, a):
                test_accum.add(10)
                yield (a,)

            def terminate(self):
                test_accum.add(100)
                yield (100,)

        self.spark.udtf.register("test_udtf", TestUDTF)

        for i, df in enumerate([TestUDTF(lit(10)), self.spark.sql("SELECT * FROM test_udtf(10)")]):
            with self.subTest(query_no=i):
                assertSchemaEqual(df.schema, StructType().add("col1", IntegerType()))
                assertDataFrameEqual(df, [Row(col1=10), Row(col1=100)])

        self.assertEqual(test_accum.value, 222)

    def _add_pyfile(self, path):
        self.sc.addPyFile(path)

    def test_udtf_with_analyze_using_pyfile(self):
        with tempfile.TemporaryDirectory(prefix="test_udtf_with_analyze_using_pyfile") as d:
            pyfile_path = os.path.join(d, "my_pyfile.py")
            with open(pyfile_path, "w") as f:
                f.write("my_func = lambda: 'col1'")

            self._add_pyfile(pyfile_path)

            class TestUDTF:
                @staticmethod
                def call_my_func() -> str:
                    import my_pyfile

                    return my_pyfile.my_func()

                @staticmethod
                def analyze(a: AnalyzeArgument) -> AnalyzeResult:
                    return AnalyzeResult(StructType().add(TestUDTF.call_my_func(), a.dataType))

                def eval(self, a):
                    assert TestUDTF.call_my_func() == "col1"
                    yield (a,)

                def terminate(self):
                    assert TestUDTF.call_my_func() == "col1"
                    yield (100,)

            test_udtf = udtf(TestUDTF)
            self.spark.udtf.register("test_udtf", test_udtf)

            for i, df in enumerate(
                [test_udtf(lit(10)), self.spark.sql("SELECT * FROM test_udtf(10)")]
            ):
                with self.subTest(query_no=i):
                    assertSchemaEqual(df.schema, StructType().add("col1", IntegerType()))
                    assertDataFrameEqual(df, [Row(col1=10), Row(col1=100)])

    def test_udtf_with_analyze_using_zipped_package(self):
        with tempfile.TemporaryDirectory(prefix="test_udtf_with_analyze_using_zipped_package") as d:
            package_path = os.path.join(d, "my_zipfile")
            os.mkdir(package_path)
            pyfile_path = os.path.join(package_path, "__init__.py")
            with open(pyfile_path, "w") as f:
                f.write("my_func = lambda: 'col1'")
            shutil.make_archive(package_path, "zip", d, "my_zipfile")

            self._add_pyfile(f"{package_path}.zip")

            class TestUDTF:
                @staticmethod
                def call_my_func() -> str:
                    import my_zipfile

                    return my_zipfile.my_func()

                @staticmethod
                def analyze(a: AnalyzeArgument) -> AnalyzeResult:
                    return AnalyzeResult(StructType().add(TestUDTF.call_my_func(), a.dataType))

                def eval(self, a):
                    assert TestUDTF.call_my_func() == "col1"
                    yield (a,)

                def terminate(self):
                    assert TestUDTF.call_my_func() == "col1"
                    yield (100,)

            test_udtf = udtf(TestUDTF)
            self.spark.udtf.register("test_udtf", test_udtf)

            for i, df in enumerate(
                [test_udtf(lit(10)), self.spark.sql("SELECT * FROM test_udtf(10)")]
            ):
                with self.subTest(query_no=i):
                    assertSchemaEqual(df.schema, StructType().add("col1", IntegerType()))
                    assertDataFrameEqual(df, [Row(col1=10), Row(col1=100)])

    def _add_archive(self, path):
        self.sc.addArchive(path)

    def test_udtf_with_analyze_using_archive(self):
        from pyspark.core.files import SparkFiles

        self.check_udtf_with_analyze_using_archive(SparkFiles.getRootDirectory())

    def check_udtf_with_analyze_using_archive(self, exec_root_dir):
        with tempfile.TemporaryDirectory(prefix="test_udtf_with_analyze_using_archive") as d:
            archive_path = os.path.join(d, "my_archive")
            os.mkdir(archive_path)
            pyfile_path = os.path.join(archive_path, "my_file.txt")
            with open(pyfile_path, "w") as f:
                f.write("col1")
            shutil.make_archive(archive_path, "zip", d, "my_archive")

            self._add_archive(f"{archive_path}.zip#my_files")

            class TestUDTF:
                @staticmethod
                def read_my_archive() -> str:
                    with open(
                        os.path.join(exec_root_dir, "my_files", "my_archive", "my_file.txt"),
                        "r",
                    ) as my_file:
                        return my_file.read().strip()

                @staticmethod
                def analyze(a: AnalyzeArgument) -> AnalyzeResult:
                    return AnalyzeResult(StructType().add(TestUDTF.read_my_archive(), a.dataType))

                def eval(self, a):
                    assert TestUDTF.read_my_archive() == "col1"
                    yield (a,)

                def terminate(self):
                    assert TestUDTF.read_my_archive() == "col1"
                    yield (100,)

            test_udtf = udtf(TestUDTF)
            self.spark.udtf.register("test_udtf", test_udtf)

            for i, df in enumerate(
                [test_udtf(lit(10)), self.spark.sql("SELECT * FROM test_udtf(10)")]
            ):
                with self.subTest(query_no=i):
                    assertSchemaEqual(df.schema, StructType().add("col1", IntegerType()))
                    assertDataFrameEqual(df, [Row(col1=10), Row(col1=100)])

    def _add_file(self, path):
        self.sc.addFile(path)

    def test_udtf_with_analyze_using_file(self):
        from pyspark.core.files import SparkFiles

        self.check_udtf_with_analyze_using_file(SparkFiles.getRootDirectory())

    def check_udtf_with_analyze_using_file(self, exec_root_dir):
        with tempfile.TemporaryDirectory(prefix="test_udtf_with_analyze_using_file") as d:
            file_path = os.path.join(d, "my_file.txt")
            with open(file_path, "w") as f:
                f.write("col1")

            self._add_file(file_path)

            class TestUDTF:
                @staticmethod
                def read_my_file() -> str:
                    with open(os.path.join(exec_root_dir, "my_file.txt"), "r") as my_file:
                        return my_file.read().strip()

                @staticmethod
                def analyze(a: AnalyzeArgument) -> AnalyzeResult:
                    return AnalyzeResult(StructType().add(TestUDTF.read_my_file(), a.dataType))

                def eval(self, a):
                    assert TestUDTF.read_my_file() == "col1"
                    yield (a,)

                def terminate(self):
                    assert TestUDTF.read_my_file() == "col1"
                    yield (100,)

            test_udtf = udtf(TestUDTF)
            self.spark.udtf.register("test_udtf", test_udtf)

            for i, df in enumerate(
                [test_udtf(lit(10)), self.spark.sql("SELECT * FROM test_udtf(10)")]
            ):
                with self.subTest(query_no=i):
                    assertSchemaEqual(df.schema, StructType().add("col1", IntegerType()))
                    assertDataFrameEqual(df, [Row(col1=10), Row(col1=100)])

    def test_udtf_with_named_arguments(self):
        @udtf(returnType="a: int")
        class TestUDTF:
            def eval(self, a, b):
                yield (a,)

        self.spark.udtf.register("test_udtf", TestUDTF)

        for i, df in enumerate(
            [
                self.spark.sql("SELECT * FROM test_udtf(a => 10, b => 'x')"),
                self.spark.sql("SELECT * FROM test_udtf(b => 'x', a => 10)"),
                TestUDTF(a=lit(10), b=lit("x")),
                TestUDTF(b=lit("x"), a=lit(10)),
            ]
        ):
            with self.subTest(query_no=i):
                assertDataFrameEqual(df, [Row(a=10)])

    def test_udtf_with_named_table_arguments(self):
        @udtf(returnType="a: int")
        class TestUDTF:
            def eval(self, a, b):
                yield (a.id,)

        self.spark.udtf.register("test_udtf", TestUDTF)

        for i, df in enumerate(
            [
                self.spark.sql("SELECT * FROM test_udtf(a => TABLE(FROM range(3)), b => 'x')"),
                self.spark.sql("SELECT * FROM test_udtf(b => 'x', a => TABLE(FROM range(3)))"),
                TestUDTF(a=self.spark.range(3).asTable(), b=lit("x")),
                TestUDTF(b=lit("x"), a=self.spark.range(3).asTable()),
            ]
        ):
            with self.subTest(query_no=i):
                assertDataFrameEqual(df, [Row(a=i) for i in range(3)])

    def test_udtf_with_named_arguments_negative(self):
        @udtf(returnType="a: int")
        class TestUDTF:
            def eval(self, a, b):
                yield (a,)

        self.spark.udtf.register("test_udtf", TestUDTF)

        with self.assertRaisesRegex(
            AnalysisException,
            "DUPLICATE_ROUTINE_PARAMETER_ASSIGNMENT.DOUBLE_NAMED_ARGUMENT_REFERENCE",
        ):
            self.spark.sql("SELECT * FROM test_udtf(a => 10, a => 100)").show()

        with self.assertRaisesRegex(AnalysisException, "UNEXPECTED_POSITIONAL_ARGUMENT"):
            self.spark.sql("SELECT * FROM test_udtf(a => 10, 'x')").show()

        with self.assertRaisesRegex(PythonException, BaseUDTFTestsMixin.missingARequiredArgument):
            self.spark.sql("SELECT * FROM test_udtf(c => 'x')").show()

        with self.assertRaisesRegex(PythonException, BaseUDTFTestsMixin.multipleValuesForArgument):
            self.spark.sql("SELECT * FROM test_udtf(10, a => 100)").show()

    def test_udtf_with_kwargs(self):
        @udtf(returnType="a: int, b: string")
        class TestUDTF:
            def eval(self, **kwargs):
                yield kwargs["a"], kwargs["b"]

        self.spark.udtf.register("test_udtf", TestUDTF)

        for i, df in enumerate(
            [
                self.spark.sql("SELECT * FROM test_udtf(a => 10, b => 'x')"),
                self.spark.sql("SELECT * FROM test_udtf(b => 'x', a => 10)"),
                TestUDTF(a=lit(10), b=lit("x")),
                TestUDTF(b=lit("x"), a=lit(10)),
            ]
        ):
            with self.subTest(query_no=i):
                assertDataFrameEqual(df, [Row(a=10, b="x")])

        # negative
        with self.assertRaisesRegex(
            AnalysisException,
            "DUPLICATE_ROUTINE_PARAMETER_ASSIGNMENT.DOUBLE_NAMED_ARGUMENT_REFERENCE",
        ):
            self.spark.sql("SELECT * FROM test_udtf(a => 10, a => 100)").show()

        with self.assertRaisesRegex(AnalysisException, "UNEXPECTED_POSITIONAL_ARGUMENT"):
            self.spark.sql("SELECT * FROM test_udtf(a => 10, 'x')").show()

    def test_udtf_with_table_argument_and_kwargs(self):
        @udtf(returnType="a: int, b: string")
        class TestUDTF:
            def eval(self, **kwargs):
                yield kwargs["a"].id, kwargs["b"]

        self.spark.udtf.register("test_udtf", TestUDTF)

        for i, df in enumerate(
            [
                self.spark.sql("SELECT * FROM test_udtf(a => TABLE(FROM range(3)), b => 'x')"),
                self.spark.sql("SELECT * FROM test_udtf(b => 'x', a => TABLE(FROM range(3)))"),
                TestUDTF(a=self.spark.range(3).asTable(), b=lit("x")),
                TestUDTF(b=lit("x"), a=self.spark.range(3).asTable()),
            ]
        ):
            with self.subTest(query_no=i):
                assertDataFrameEqual(df, [Row(a=i, b="x") for i in range(3)])

    def test_udtf_with_analyze_kwargs(self):
        @udtf
        class TestUDTF:
            @staticmethod
            def analyze(**kwargs: AnalyzeArgument) -> AnalyzeResult:
                assert isinstance(kwargs["a"].dataType, IntegerType)
                assert kwargs["a"].value == 10
                assert kwargs["a"].isTable is False
                assert kwargs["a"].isConstantExpression is True
                assert isinstance(kwargs["b"].dataType, StringType)
                assert kwargs["b"].value == "x"
                assert not kwargs["b"].isTable
                return AnalyzeResult(
                    StructType(
                        [StructField(key, arg.dataType) for key, arg in sorted(kwargs.items())]
                    )
                )

            def eval(self, **kwargs):
                yield tuple(value for _, value in sorted(kwargs.items()))

        self.spark.udtf.register("test_udtf", TestUDTF)

        for i, df in enumerate(
            [
                self.spark.sql("SELECT * FROM test_udtf(a => 10, b => 'x')"),
                self.spark.sql("SELECT * FROM test_udtf(b => 'x', a => 10)"),
                TestUDTF(a=lit(10), b=lit("x")),
                TestUDTF(b=lit("x"), a=lit(10)),
            ]
        ):
            with self.subTest(query_no=i):
                assertDataFrameEqual(df, [Row(a=10, b="x")])

    def test_udtf_with_table_argument_and_analyze_kwargs(self):
        @udtf
        class TestUDTF:
            @staticmethod
            def analyze(**kwargs: AnalyzeArgument) -> AnalyzeResult:
                assert isinstance(kwargs["a"].dataType, StructType)
                assert kwargs["a"].isTable is True
                assert isinstance(kwargs["b"].dataType, StringType)
                assert kwargs["b"].value == "x"
                assert not kwargs["b"].isTable
                return AnalyzeResult(
                    StructType(
                        [StructField(key, arg.dataType) for key, arg in sorted(kwargs.items())]
                    )
                )

            def eval(self, **kwargs):
                yield tuple(value for _, value in sorted(kwargs.items()))

        self.spark.udtf.register("test_udtf", TestUDTF)

        for i, df in enumerate(
            [
                self.spark.sql("SELECT * FROM test_udtf(a => TABLE(FROM range(3)), b => 'x')"),
                self.spark.sql("SELECT * FROM test_udtf(b => 'x', a => TABLE(FROM range(3)))"),
                TestUDTF(a=self.spark.range(3).asTable(), b=lit("x")),
                TestUDTF(b=lit("x"), a=self.spark.range(3).asTable()),
            ]
        ):
            with self.subTest(query_no=i):
                assertDataFrameEqual(df, [Row(a=Row(id=i), b="x") for i in range(3)])

    def test_udtf_with_named_arguments_lateral_join(self):
        @udtf
        class TestUDTF:
            @staticmethod
            def analyze(a, b):
                return AnalyzeResult(StructType().add("a", a.dataType))

            def eval(self, a, b):
                yield (a,)

        self.spark.udtf.register("test_udtf", TestUDTF)

        # lateral join
        for i, df in enumerate(
            [
                self.spark.sql(
                    "SELECT f.* FROM "
                    "VALUES (0, 'x'), (1, 'y') t(a, b), LATERAL test_udtf(a => a, b => b) f"
                ),
                self.spark.sql(
                    "SELECT f.* FROM "
                    "VALUES (0, 'x'), (1, 'y') t(a, b), LATERAL test_udtf(b => b, a => a) f"
                ),
            ]
        ):
            with self.subTest(query_no=i):
                assertDataFrameEqual(df, [Row(a=0), Row(a=1)])

    def test_udtf_with_named_arguments_and_defaults(self):
        @udtf
        class TestUDTF:
            @staticmethod
            def analyze(a: AnalyzeArgument, b: Optional[AnalyzeArgument] = None):
                assert isinstance(a.dataType, IntegerType)
                assert a.value == 10
                assert not a.isTable
                assert a.isConstantExpression is True
                if b is not None:
                    assert isinstance(b.dataType, StringType)
                    assert b.value == "z"
                    assert not b.isTable
                schema = StructType().add("a", a.dataType)
                if b is None:
                    return AnalyzeResult(schema.add("b", IntegerType()))
                else:
                    return AnalyzeResult(schema.add("b", b.dataType))

            def eval(self, a, b=100):
                yield a, b

        self.spark.udtf.register("test_udtf", TestUDTF)

        # without "b"
        for i, df in enumerate(
            [
                self.spark.sql("SELECT * FROM test_udtf(10)"),
                self.spark.sql("SELECT * FROM test_udtf(a => 10)"),
                TestUDTF(lit(10)),
                TestUDTF(a=lit(10)),
            ]
        ):
            with self.subTest(with_b=False, query_no=i):
                assertDataFrameEqual(df, [Row(a=10, b=100)])

        # with "b"
        for i, df in enumerate(
            [
                self.spark.sql("SELECT * FROM test_udtf(10, b => 'z')"),
                self.spark.sql("SELECT * FROM test_udtf(a => 10, b => 'z')"),
                self.spark.sql("SELECT * FROM test_udtf(b => 'z', a => 10)"),
                TestUDTF(lit(10), b=lit("z")),
                TestUDTF(a=lit(10), b=lit("z")),
                TestUDTF(b=lit("z"), a=lit(10)),
            ]
        ):
            with self.subTest(with_b=True, query_no=i):
                assertDataFrameEqual(df, [Row(a=10, b="z")])

    def test_udtf_with_table_argument_and_partition_by(self):
        class TestUDTF:
            def __init__(self):
                self._sum = 0
                self._partition_col = None

            def eval(self, row: Row):
                # Make sure that the PARTITION BY expressions were projected out.
                assert len(row.asDict().items()) == 2
                assert "partition_col" in row
                assert "input" in row
                self._sum += row["input"]
                if self._partition_col is not None and self._partition_col != row["partition_col"]:
                    # Make sure that all values of the partitioning column are the same
                    # for each row consumed by this method for this instance of the class.
                    raise Exception(
                        f"self._partition_col was {self._partition_col} but the row "
                        + f"value was {row['partition_col']}"
                    )
                self._partition_col = row["partition_col"]

            def terminate(self):
                yield self._partition_col, self._sum

        # This is a basic example.
        func = udtf(TestUDTF, returnType="partition_col: int, total: int")
        self.spark.udtf.register("test_udtf", func)

        base_query = """
            WITH t AS (
              SELECT id AS partition_col, 1 AS input FROM range(1, 21)
              UNION ALL
              SELECT id AS partition_col, 2 AS input FROM range(1, 21)
            )
            SELECT partition_col, total
            FROM test_udtf({table_arg})
            ORDER BY 1, 2
        """

        for table_arg in [
            "TABLE(t) PARTITION BY partition_col - 1",
            "row => TABLE(t) PARTITION BY partition_col - 1",
        ]:
            with self.subTest(table_arg=table_arg):
                assertDataFrameEqual(
                    self.spark.sql(base_query.format(table_arg=table_arg)),
                    [Row(partition_col=x, total=3) for x in range(1, 21)],
                )

        base_query = """
            WITH t AS (
              SELECT {str_first} AS partition_col, id AS input FROM range(0, 2)
              UNION ALL
              SELECT {str_second} AS partition_col, id AS input FROM range(0, 2)
            )
            SELECT partition_col, total
            FROM test_udtf({table_arg})
            ORDER BY 1, 2
        """

        # These cases partition by constant values.
        for str_first, str_second, result_first, result_second in (
            ("123", "456", 123, 456),
            ("123", "NULL", None, 123),
        ):
            for table_arg in [
                "TABLE(t) PARTITION BY partition_col",
                "row => TABLE(t) PARTITION BY partition_col",
            ]:
                with self.subTest(str_first=str_first, str_second=str_second, table_arg=table_arg):
                    assertDataFrameEqual(
                        self.spark.sql(
                            base_query.format(
                                str_first=str_first, str_second=str_second, table_arg=table_arg
                            )
                        ),
                        [
                            Row(partition_col=result_first, total=1),
                            Row(partition_col=result_second, total=1),
                        ],
                    )

        # Combine a lateral join with a TABLE argument with PARTITION BY .
        func = udtf(TestUDTF, returnType="partition_col: int, total: int")
        self.spark.udtf.register("test_udtf", func)

        base_query = """
            WITH t AS (
              SELECT id AS partition_col, 1 AS input FROM range(1, 3)
              UNION ALL
              SELECT id AS partition_col, 2 AS input FROM range(1, 3)
            )
            SELECT v.a, v.b, f.partition_col, f.total
            FROM VALUES (0, 1) AS v(a, b),
            LATERAL test_udtf({table_arg}) f
            ORDER BY 1, 2, 3, 4
        """

        for table_arg in [
            "TABLE(t) PARTITION BY partition_col - 1",
            "row => TABLE(t) PARTITION BY partition_col - 1",
        ]:
            with self.subTest(func_call=table_arg):
                assertDataFrameEqual(
                    self.spark.sql(base_query.format(table_arg=table_arg)),
                    [
                        Row(a=0, b=1, partition_col=1, total=3),
                        Row(a=0, b=1, partition_col=2, total=3),
                    ],
                )

    def test_udtf_with_table_argument_and_partition_by_no_terminate(self):
        func = self.udtf_for_table_argument()  # a udtf with no terminate method defined
        self.spark.udtf.register("test_udtf", func)

        assertDataFrameEqual(
            self.spark.sql(
                "SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)) PARTITION BY id)"
            ),
            [Row(a=6), Row(a=7)],
        )

    def test_udtf_with_table_argument_and_partition_by_and_order_by(self):
        class TestUDTF:
            def __init__(self):
                self._last = None
                self._partition_col = None

            def eval(self, row: Row, partition_col: str):
                # Make sure that the PARTITION BY and ORDER BY expressions were projected out.
                assert len(row.asDict().items()) == 2
                assert "partition_col" in row
                assert "input" in row
                # Make sure that all values of the partitioning column are the same
                # for each row consumed by this method for this instance of the class.
                if self._partition_col is not None and self._partition_col != row[partition_col]:
                    raise Exception(
                        f"self._partition_col was {self._partition_col} but the row "
                        + f"value was {row[partition_col]}"
                    )
                self._last = row["input"]
                self._partition_col = row[partition_col]

            def terminate(self):
                yield self._partition_col, self._last

        func = udtf(TestUDTF, returnType="partition_col: int, last: int")
        self.spark.udtf.register("test_udtf", func)

        base_query = """
            WITH t AS (
              SELECT id AS partition_col, 1 AS input FROM range(1, 21)
              UNION ALL
              SELECT id AS partition_col, 2 AS input FROM range(1, 21)
            )
            SELECT partition_col, last
            FROM test_udtf(
              {table_arg},
              partition_col => 'partition_col')
            ORDER BY 1, 2
        """

        for order_by_str, result_val in (
            ("input ASC", 2),
            ("input + 1 ASC", 2),
            ("input DESC", 1),
            ("input - 1 DESC", 1),
        ):
            for table_arg in [
                f"TABLE(t) PARTITION BY partition_col - 1 ORDER BY {order_by_str}",
                f"row => TABLE(t) PARTITION BY partition_col - 1 ORDER BY {order_by_str}",
            ]:
                with self.subTest(table_arg=table_arg):
                    assertDataFrameEqual(
                        self.spark.sql(base_query.format(table_arg=table_arg)),
                        [Row(partition_col=x, last=result_val) for x in range(1, 21)],
                    )

    def test_udtf_with_table_argument_with_single_partition(self):
        class TestUDTF:
            def __init__(self):
                self._count = 0
                self._sum = 0
                self._last = None

            def eval(self, row: Row):
                # Make sure that the rows arrive in the expected order.
                if self._last is not None and self._last > row["input"]:
                    raise Exception(
                        f"self._last was {self._last} but the row value was {row['input']}"
                    )
                self._count += 1
                self._last = row["input"]
                self._sum += row["input"]

            def terminate(self):
                yield self._count, self._sum, self._last

        func = udtf(TestUDTF, returnType="count: int, total: int, last: int")
        self.spark.udtf.register("test_udtf", func)

        base_query = """
            WITH t AS (
              SELECT id AS partition_col, 1 AS input FROM range(1, 21)
              UNION ALL
              SELECT id AS partition_col, 2 AS input FROM range(1, 21)
            )
            SELECT count, total, last
            FROM test_udtf({table_arg})
            ORDER BY 1, 2
        """

        for table_arg in [
            "TABLE(t) WITH SINGLE PARTITION ORDER BY (input, partition_col)",
            "row => TABLE(t) WITH SINGLE PARTITION ORDER BY (input, partition_col)",
        ]:
            with self.subTest(table_arg=table_arg):
                assertDataFrameEqual(
                    self.spark.sql(base_query.format(table_arg=table_arg)),
                    [Row(count=40, total=60, last=2)],
                )

    def test_udtf_with_table_argument_with_single_partition_from_analyze(self):
        @udtf
        class TestUDTF:
            def __init__(self):
                self._count = 0
                self._sum = 0
                self._last = None

            @staticmethod
            def analyze(*args, **kwargs):
                return AnalyzeResult(
                    schema=StructType()
                    .add("count", IntegerType())
                    .add("total", IntegerType())
                    .add("last", IntegerType()),
                    withSinglePartition=True,
                    orderBy=[OrderingColumn("input"), OrderingColumn("partition_col")],
                )

            def eval(self, row: Row):
                # Make sure that the rows arrive in the expected order.
                if self._last is not None and self._last > row["input"]:
                    raise Exception(
                        f"self._last was {self._last} but the row value was {row['input']}"
                    )
                self._count += 1
                self._last = row["input"]
                self._sum += row["input"]

            def terminate(self):
                yield self._count, self._sum, self._last

        self.spark.udtf.register("test_udtf", TestUDTF)

        base_query = """
            WITH t AS (
              SELECT id AS partition_col, 1 AS input FROM range(1, 21)
              UNION ALL
              SELECT id AS partition_col, 2 AS input FROM range(1, 21)
            )
            SELECT count, total, last
            FROM test_udtf({table_arg})
            ORDER BY 1, 2
        """

        for table_arg in ["TABLE(t)", "row => TABLE(t)"]:
            with self.subTest(table_arg):
                assertDataFrameEqual(
                    self.spark.sql(base_query.format(table_arg=table_arg)),
                    [Row(count=40, total=60, last=2)],
                )

    def test_udtf_with_table_argument_with_partition_by_and_order_by_from_analyze(self):
        @udtf
        class TestUDTF:
            def __init__(self):
                self._partition_col = None
                self._count = 0
                self._sum = 0
                self._last = None

            @staticmethod
            def analyze(*args, **kwargs):
                return AnalyzeResult(
                    schema=StructType()
                    .add("partition_col", IntegerType())
                    .add("count", IntegerType())
                    .add("total", IntegerType())
                    .add("last", IntegerType()),
                    partitionBy=[PartitioningColumn("partition_col")],
                    orderBy=[
                        OrderingColumn(name="input", ascending=True, overrideNullsFirst=False)
                    ],
                )

            def eval(self, row: Row):
                # Make sure that the PARTITION BY and ORDER BY expressions were projected out.
                assert len(row.asDict().items()) == 2
                assert "partition_col" in row
                assert "input" in row
                # Make sure that all values of the partitioning column are the same
                # for each row consumed by this method for this instance of the class.
                if self._partition_col is not None and self._partition_col != row["partition_col"]:
                    raise Exception(
                        f"self._partition_col was {self._partition_col} but the row "
                        + f"value was {row['partition_col']}"
                    )
                # Make sure that the rows arrive in the expected order.
                if (
                    self._last is not None
                    and row["input"] is not None
                    and self._last > row["input"]
                ):
                    raise Exception(
                        f"self._last was {self._last} but the row value was {row['input']}"
                    )
                self._partition_col = row["partition_col"]
                self._count += 1
                self._last = row["input"]
                if row["input"] is not None:
                    self._sum += row["input"]

            def terminate(self):
                yield self._partition_col, self._count, self._sum, self._last

        self.spark.udtf.register("test_udtf", TestUDTF)

        base_query = """
            WITH t AS (
              SELECT id AS partition_col, 1 AS input FROM range(1, 21)
              UNION ALL
              SELECT id AS partition_col, 2 AS input FROM range(1, 21)
              UNION ALL
              SELECT 42 AS partition_col, NULL AS input
              UNION ALL
              SELECT 42 AS partition_col, 1 AS input
              UNION ALL
              SELECT 42 AS partition_col, 2 AS input
            )
            SELECT partition_col, count, total, last
            FROM test_udtf({table_arg})
            ORDER BY 1, 2
        """

        for table_arg in ["TABLE(t)", "row => TABLE(t)"]:
            with self.subTest(table_arg=table_arg):
                assertDataFrameEqual(
                    self.spark.sql(base_query.format(table_arg=table_arg)),
                    [Row(partition_col=x, count=2, total=3, last=2) for x in range(1, 21)]
                    + [Row(partition_col=42, count=3, total=3, last=None)],
                )

    def test_udtf_with_analyze_order_by_override_nulls_first(self):
        for override_nulls_first in [True, False]:

            @udtf
            class TestUDTF:
                @staticmethod
                def analyze(*args, **kwargs) -> AnalyzeResult:
                    return AnalyzeResult(
                        StructType().add("id", IntegerType()),
                        withSinglePartition=True,
                        orderBy=[OrderingColumn("id", overrideNullsFirst=override_nulls_first)],
                    )

                def eval(self, row: Row):
                    yield (row["id"],)

            df = self.spark.createDataFrame([(1,), (None,)], ["id"])
            assertDataFrameEqual(
                TestUDTF(df.asTable()).collect(),
                [Row(id=None), Row(id=1)] if override_nulls_first else [Row(id=1), Row(id=None)],
                checkRowOrder=True,
            )

    def test_udtf_with_prepare_string_from_analyze(self):
        @dataclass
        class AnalyzeResultWithBuffer(AnalyzeResult):
            buffer: str = ""

        @udtf
        class TestUDTF:
            def __init__(self, analyze_result=None):
                self._total = 0
                if analyze_result is not None:
                    self._buffer = analyze_result.buffer
                else:
                    self._buffer = ""

            @staticmethod
            def analyze(argument, _):
                if (
                    argument.value is None
                    or argument.isTable
                    or not isinstance(argument.value, str)
                    or len(argument.value) == 0
                ):
                    raise Exception("The first argument must be non-empty string")
                assert argument.dataType == StringType()
                assert not argument.isTable
                return AnalyzeResultWithBuffer(
                    schema=StructType().add("total", IntegerType()).add("buffer", StringType()),
                    withSinglePartition=True,
                    buffer=argument.value,
                )

            def eval(self, argument, row: Row):
                self._total += 1

            def terminate(self):
                yield self._total, self._buffer

        self.spark.udtf.register("test_udtf", TestUDTF)

        assertDataFrameEqual(
            self.spark.sql("""
                WITH t AS (
                  SELECT id FROM range(1, 21)
                )
                SELECT total, buffer
                FROM test_udtf("abc", TABLE(t))
                """),
            [Row(count=20, buffer="abc")],
        )

    def test_udtf_with_skip_rest_of_input_table_exception(self):
        @udtf(returnType="current: int, total: int")
        class TestUDTF:
            def __init__(self):
                self._current = 0
                self._total = 0

            def eval(self, input: Row):
                self._current = input["id"]
                self._total += 1
                if self._total >= 4:
                    raise SkipRestOfInputTableException("Stop at self._total >= 4")

            def terminate(self):
                yield self._current, self._total

        self.spark.udtf.register("test_udtf", TestUDTF)

        # Run a test case including WITH SINGLE PARTITION on the UDTF call. The
        # SkipRestOfInputTableException stops scanning rows after the fourth input row is consumed.
        assertDataFrameEqual(
            self.spark.sql("""
                WITH t AS (
                  SELECT id FROM range(1, 21)
                )
                SELECT current, total
                FROM test_udtf(TABLE(t) WITH SINGLE PARTITION ORDER BY id)
                """),
            [Row(current=4, total=4)],
        )

        # Run a test case including WITH SINGLE PARTITION on the UDTF call. The
        # SkipRestOfInputTableException stops scanning rows for each of the two partitions
        # separately.
        assertDataFrameEqual(
            self.spark.sql("""
                WITH t AS (
                  SELECT id FROM range(1, 21)
                )
                SELECT current, total
                FROM test_udtf(TABLE(t) PARTITION BY floor(id / 10) ORDER BY id)
                ORDER BY ALL
                """),
            [Row(current=4, total=4), Row(current=13, total=4), Row(current=20, total=1)],
        )

    def test_udtf_with_variant_input(self):
        @udtf(returnType="i int, s: string")
        class TestUDTF:
            def eval(self, v):
                for i in range(10):
                    yield i, v.toJson()

        self.spark.udtf.register("test_udtf", TestUDTF)
        assertDataFrameEqual(
            self.spark.sql('select i, s from test_udtf(parse_json(\'{"a":"b"}\'))'),
            [Row(i=n, s='{"a":"b"}') for n in range(10)],
        )

    def test_udtf_with_nested_variant_input(self):
        # struct<variant>
        @udtf(returnType="i int, s: string")
        class TestUDTFStruct:
            def eval(self, v):
                for i in range(10):
                    yield i, v["v"].toJson()

        self.spark.udtf.register("test_udtf_struct", TestUDTFStruct)
        res = self.spark.sql(
            "select i, s from test_udtf_struct(named_struct('v', parse_json('{\"a\":\"c\"}')))"
        )
        assertDataFrameEqual(res, [Row(i=n, s='{"a":"c"}') for n in range(10)])

        # array<variant>
        @udtf(returnType="i int, s: string")
        class TestUDTFArray:
            def eval(self, v):
                for i in range(10):
                    yield i, v[0].toJson()

        self.spark.udtf.register("test_udtf_array", TestUDTFArray)
        res = self.spark.sql('select i, s from test_udtf_array(array(parse_json(\'{"a":"d"}\')))')
        assertDataFrameEqual(res, [Row(i=n, s='{"a":"d"}') for n in range(10)])

        # map<string, variant>
        @udtf(returnType="i int, s: string")
        class TestUDTFMap:
            def eval(self, v):
                for i in range(10):
                    yield i, v["v"].toJson()

        self.spark.udtf.register("test_udtf_map", TestUDTFMap)
        res = self.spark.sql(
            "select i, s from test_udtf_map(map('v', parse_json('{\"a\":\"e\"}')))"
        )
        assertDataFrameEqual(res, [Row(i=n, s='{"a":"e"}') for n in range(10)])

    def test_udtf_with_variant_output(self):
        @udtf(returnType="i int, v: variant")
        class TestUDTF:
            # TODO(SPARK-50284): Replace when an easy Python API to construct Variants is created.
            def eval(self, n):
                for i in range(n):
                    yield i, VariantVal(bytes([2, 1, 0, 0, 2, 5, 97 + i]), bytes([1, 1, 0, 1, 97]))

        self.spark.udtf.register("test_udtf", TestUDTF)
        res = self.spark.sql("select i, to_json(v) from test_udtf(8)")
        assertDataFrameEqual(res, [Row(i=n, s=f'{{"a":"{chr(97 + n)}"}}') for n in range(8)])

    def test_udtf_with_nested_variant_output(self):
        # struct<variant>
        @udtf(returnType="i int, v: struct<v1 variant>")
        class TestUDTFStruct:
            # TODO(SPARK-50284): Replace when an easy Python API to construct Variants is created.
            def eval(self, n):
                for i in range(n):
                    yield (
                        i,
                        {
                            "v1": VariantVal(
                                bytes([2, 1, 0, 0, 2, 5, 97 + i]), bytes([1, 1, 0, 1, 97])
                            )
                        },
                    )

        self.spark.udtf.register("test_udtf_struct", TestUDTFStruct)
        res = self.spark.sql("select i, to_json(v.v1) from test_udtf_struct(8)")
        assertDataFrameEqual(res, [Row(i=n, s=f'{{"a":"{chr(97 + n)}"}}') for n in range(8)])

        # array<variant>
        @udtf(returnType="i int, v: array<variant>")
        class TestUDTFArray:
            # TODO(SPARK-50284): Replace when an easy Python API to construct Variants is created.
            def eval(self, n):
                for i in range(n):
                    yield (
                        i,
                        [VariantVal(bytes([2, 1, 0, 0, 2, 5, 98 + i]), bytes([1, 1, 0, 1, 97]))],
                    )

        self.spark.udtf.register("test_udtf_array", TestUDTFArray)
        res = self.spark.sql("select i, to_json(v[0]) from test_udtf_array(8)")
        assertDataFrameEqual(res, [Row(i=n, s=f'{{"a":"{chr(98 + n)}"}}') for n in range(8)])

        # map<string, variant>
        @udtf(returnType="i int, v: map<string, variant>")
        class TestUDTFStruct:
            # TODO(SPARK-50284): Replace when an easy Python API to construct Variants is created.
            def eval(self, n):
                for i in range(n):
                    yield (
                        i,
                        {
                            "v1": VariantVal(
                                bytes([2, 1, 0, 0, 2, 5, 99 + i]), bytes([1, 1, 0, 1, 97])
                            )
                        },
                    )

        self.spark.udtf.register("test_udtf_struct", TestUDTFStruct)
        res = self.spark.sql("select i, to_json(v['v1']) from test_udtf_struct(8)")
        assertDataFrameEqual(res, [Row(i=n, s=f'{{"a":"{chr(99 + n)}"}}') for n in range(8)])

    def test_udtf_segfault(self):
        for enabled, expected in [
            (True, "Segmentation fault"),
            (False, "Consider setting .* for the better Python traceback."),
        ]:
            with (
                self.subTest(enabled=enabled),
                self.sql_conf({"spark.sql.execution.pyspark.udf.faulthandler.enabled": enabled}),
            ):
                with self.subTest(method="eval"):

                    class TestUDTF:
                        def eval(self):
                            import ctypes

                            yield (ctypes.string_at(0),)

                    self._check_result_or_exception(
                        TestUDTF, "x: string", expected, err_type=Exception
                    )

                with self.subTest(method="analyze"):

                    class TestUDTFWithAnalyze:
                        @staticmethod
                        def analyze():
                            import ctypes

                            ctypes.string_at(0)
                            return AnalyzeResult(StructType().add("x", StringType()))

                        def eval(self):
                            yield ("x",)

                    self._check_result_or_exception(
                        TestUDTFWithAnalyze, None, expected, err_type=Exception
                    )

    def test_udtf_kill_on_timeout(self):
        with self.sql_conf(
            {
                "spark.sql.execution.pyspark.udf.idleTimeoutSeconds": "1s",
                "spark.sql.execution.pyspark.udf.killOnIdleTimeout": "true",
            }
        ):
            with self.subTest(method="eval"):

                class TestUDTF:
                    def eval(self):
                        time.sleep(2)
                        yield "x"

                self._check_result_or_exception(
                    TestUDTF,
                    "x: string",
                    "Python worker process terminated due to idle timeout \\(timeout: 1 seconds\\)",
                    err_type=Exception,
                )

            with self.subTest(method="analyze"):

                class TestUDTFWithAnalyze:
                    @staticmethod
                    def analyze():
                        time.sleep(2)
                        return AnalyzeResult(StructType().add("x", StringType()))

                    def eval(self):
                        yield ("x",)

                self._check_result_or_exception(
                    TestUDTFWithAnalyze,
                    None,
                    "Python worker process terminated due to idle timeout \\(timeout: 1 seconds\\)",
                    err_type=Exception,
                )

    def test_udtf_with_collated_string_types(self):
        @udtf(
            returnType="out1 string, out2 string collate UTF8_BINARY, "
            "out3 string collate UTF8_LCASE, out4 string collate UNICODE"
        )
        class MyUDTF:
            def eval(self, v1, v2, v3, v4):
                yield (v1 + "1", v2 + "2", v3 + "3", v4 + "4")

        schema = StructType(
            [
                StructField("col1", StringType(), True),
                StructField("col2", StringType("UTF8_BINARY"), True),
                StructField("col3", StringType("UTF8_LCASE"), True),
                StructField("col4", StringType("UNICODE"), True),
            ]
        )
        df = self.spark.createDataFrame([("hello",) * 4], schema=schema)

        result_df = df.lateralJoin(
            MyUDTF(
                col("col1").outer(), col("col2").outer(), col("col3").outer(), col("col4").outer()
            )
        ).select("out1", "out2", "out3", "out4")

        expected_row = ("hello1", "hello2", "hello3", "hello4")
        self.assertEqual(result_df.collect()[0], expected_row)

        expected_output_types = [
            StringType(),
            StringType("UTF8_BINARY"),
            StringType("UTF8_LCASE"),
            StringType("UNICODE"),
        ]
        for idx, field in enumerate(result_df.schema.fields):
            self.assertEqual(field.dataType, expected_output_types[idx])

    def test_udtf_binary_type(self):
        @udtf(returnType="type_name: string")
        class BinaryTypeUDTF:
            def eval(self, b):
                # Check the type of the binary input and return type name as string
                yield (type(b).__name__,)

        for conf_value in ["true", "false"]:
            expected_type = "bytes" if conf_value == "true" else "bytearray"
            with self.sql_conf({"spark.sql.execution.pyspark.binaryAsBytes": conf_value}):
                result = BinaryTypeUDTF(lit(b"test")).collect()
                self.assertEqual(result[0]["type_name"], expected_type)

    @unittest.skipIf(is_remote_only(), "Requires JVM access")
    def test_udtf_with_logging(self):
        @udtf(returnType="a: int, b: int")
        class TestUDTFWithLogging:
            def eval(self, x: int):
                logger = logging.getLogger("test_udtf")
                logger.warning(f"udtf with logging: {x}")
                yield x * 2, x + 10

        with self.sql_conf({"spark.sql.pyspark.worker.logging.enabled": "true"}):
            assertDataFrameEqual(
                self.spark.createDataFrame([(5,), (10,)], ["x"]).lateralJoin(
                    TestUDTFWithLogging(col("x").outer())
                ),
                [Row(x=x, a=x * 2, b=x + 10) for x in [5, 10]],
            )

            logs = self.spark.tvf.python_worker_logs()

            assertDataFrameEqual(
                logs.select("level", "msg", "context", "logger"),
                [
                    Row(
                        level="WARNING",
                        msg=f"udtf with logging: {x}",
                        context={"class_name": "TestUDTFWithLogging", "func_name": "eval"},
                        logger="test_udtf",
                    )
                    for x in [5, 10]
                ],
            )

    @unittest.skipIf(is_remote_only(), "Requires JVM access")
    def test_udtf_analyze_with_logging(self):
        @udtf
        class TestUDTFWithLogging:
            @staticmethod
            def analyze(x: AnalyzeArgument) -> AnalyzeResult:
                logger = logging.getLogger("test_udtf")
                logger.warning(f"udtf analyze: {x.dataType.json()}")
                return AnalyzeResult(StructType().add("a", IntegerType()).add("b", IntegerType()))

            def eval(self, x: int):
                yield x * 2, x + 10

        with self.sql_conf({"spark.sql.pyspark.worker.logging.enabled": "true"}):
            assertDataFrameEqual(
                self.spark.createDataFrame([(5,), (10,)], ["x"]).lateralJoin(
                    TestUDTFWithLogging(col("x").outer())
                ),
                [Row(x=x, a=x * 2, b=x + 10) for x in [5, 10]],
            )

            logs = self.spark.tvf.python_worker_logs()

            assertDataFrameEqual(
                logs.select(
                    "level",
                    "msg",
                    col("context.class_name").alias("context_class_name"),
                    col("context.func_name").alias("context_func_name"),
                    "logger",
                ).distinct(),
                [
                    Row(
                        level="WARNING",
                        msg='udtf analyze: "long"',
                        context_class_name="TestUDTFWithLogging",
                        context_func_name="analyze",
                        logger="test_udtf",
                    )
                ],
            )

    @unittest.skipIf(is_remote_only(), "Requires JVM access")
    def test_udtf_analyze_with_pyspark_logger(self):
        @udtf
        class TestUDTFWithLogging:
            @staticmethod
            def analyze(x: AnalyzeArgument) -> AnalyzeResult:
                logger = PySparkLogger.getLogger("PySparkLogger")
                logger.warning(f"udtf analyze: {x.dataType.json()}", dt=x.dataType.json())
                return AnalyzeResult(StructType().add("a", IntegerType()).add("b", IntegerType()))

            def eval(self, x: int):
                yield x * 2, x + 10

        with self.sql_conf({"spark.sql.pyspark.worker.logging.enabled": "true"}):
            assertDataFrameEqual(
                self.spark.createDataFrame([(5,), (10,)], ["x"]).lateralJoin(
                    TestUDTFWithLogging(col("x").outer())
                ),
                [Row(x=x, a=x * 2, b=x + 10) for x in [5, 10]],
            )

            logs = self.spark.tvf.python_worker_logs()

            assertDataFrameEqual(
                logs.select(
                    "level",
                    "msg",
                    col("context.class_name").alias("context_class_name"),
                    col("context.func_name").alias("context_func_name"),
                    col("context.dt").alias("context_dt"),
                    "logger",
                ).distinct(),
                [
                    Row(
                        level="WARNING",
                        msg='udtf analyze: "long"',
                        context_class_name="TestUDTFWithLogging",
                        context_func_name="analyze",
                        context_dt='"long"',
                        logger="PySparkLogger",
                    )
                ],
            )


class UDTFTests(BaseUDTFTestsMixin, ReusedSQLTestCase):
    @classmethod
    def setUpClass(cls):
        super().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().tearDownClass()


@unittest.skipIf(
    not have_pandas or not have_pyarrow, pandas_requirement_message or pyarrow_requirement_message
)
class LegacyUDTFArrowTestsMixin(BaseUDTFTestsMixin):
    def test_udtf_binary_type(self):
        @udtf(returnType="type_name: string")
        class BinaryTypeUDTF:
            def eval(self, b):
                # Check the type of the binary input and return type name as string
                yield (type(b).__name__,)

        # For Arrow Python UDTF with legacy conversion BinaryType is always mapped to bytes
        for conf_value in ["true", "false"]:
            with self.sql_conf({"spark.sql.execution.pyspark.binaryAsBytes": conf_value}):
                result = BinaryTypeUDTF(lit(b"test")).collect()
                self.assertEqual(result[0]["type_name"], "bytes")

    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,)

        old_value = self.spark.conf.get("spark.sql.execution.pythonUDTF.arrow.enabled")
        try:
            self.spark.conf.set("spark.sql.execution.pythonUDTF.arrow.enabled", False)
            self.assertEqual(
                udtf(TestUDTF, returnType="x: int").evalType, PythonEvalType.SQL_TABLE_UDF
            )
            self.spark.conf.set("spark.sql.execution.pythonUDTF.arrow.enabled", True)
            self.assertEqual(
                udtf(TestUDTF, returnType="x: int").evalType, PythonEvalType.SQL_ARROW_TABLE_UDF
            )
        finally:
            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)])

        @udtf(returnType="a: int")
        class TestUDTF:
            def eval(self, a: int):
                return [a]

        assertDataFrameEqual(TestUDTF(lit(1)), [Row(a=1)])

        @udtf(returnType=StructType().add("udt", ExamplePointUDT()))
        class TestUDTF:
            def eval(self, x: float, y: float):
                yield ExamplePoint(x=x * 10, y=y * 10)

        assertDataFrameEqual(TestUDTF(lit(1.0), lit(2.0)), [Row(udt=ExamplePoint(x=10.0, y=20.0))])

    def test_udtf_use_large_var_types(self):
        for use_large_var_types in [True, False]:
            with self.subTest(use_large_var_types=use_large_var_types):
                with self.sql_conf(
                    {"spark.sql.execution.arrow.useLargeVarTypes": use_large_var_types}
                ):

                    @udtf(returnType="a: string")
                    class TestUDTF:
                        def eval(self, a: int):
                            yield str(a)

                    assertDataFrameEqual(TestUDTF(lit(1)), [Row(a="1")])

    def test_numeric_output_type_casting(self):
        class TestUDTF:
            def eval(self):
                yield (1,)

        err = "Exception thrown when converting pandas.Series"

        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 = "Exception thrown when converting pandas.Series"

        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 = "Exception thrown when converting pandas.Series"

        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 = "Exception thrown when converting pandas.Series"

        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 = "Exception thrown when converting pandas.Series"

        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 = "Exception thrown when converting pandas.Series"

        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 = "Exception thrown when converting pandas.Series"

        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, "Exception thrown when converting pandas.Series"
                ):
                    udtf(TestUDTF, returnType=ret_type)().collect()


class LegacyUDTFArrowTests(LegacyUDTFArrowTestsMixin, ReusedSQLTestCase):
    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        cls.spark.conf.set("spark.sql.execution.pythonUDTF.arrow.enabled", "true")
        cls.spark.conf.set(
            "spark.sql.legacy.execution.pythonUDTF.pandas.conversion.enabled", "true"
        )

    @classmethod
    def tearDownClass(cls):
        try:
            cls.spark.conf.unset("spark.sql.execution.pythonUDTF.arrow.enabled")
            cls.spark.conf.unset("spark.sql.legacy.execution.pythonUDTF.pandas.conversion.enabled")
        finally:
            super().tearDownClass()


class UDTFArrowTestsMixin(LegacyUDTFArrowTestsMixin):
    def test_udtf_binary_type(self):
        # For Arrow Python UDTF with non-legacy conversionBinaryType is mapped to
        # bytes or bytearray consistently with non-Arrow Python UDTF behavior.
        BaseUDTFTestsMixin.test_udtf_binary_type(self)

    def test_numeric_output_type_casting(self):
        class TestUDTF:
            def eval(self):
                yield (1,)

        err = "UDTF_ARROW_DATA_CONVERSION_ERROR"

        for ret_type, expected in [
            ("x: boolean", err),
            ("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", [Row(x=datetime.date(1970, 1, 2))]),
            ("x: byte", [Row(x=1)]),
            ("x: binary", err),
            ("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_DATA_CONVERSION_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="1")]),
            ("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>", 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_string_output_type_casting(self):
        class TestUDTF:
            def eval(self):
                yield ("hello",)

        err = "UDTF_ARROW_DATA_CONVERSION_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", err),
            ("x: float", err),
            ("x: double", err),
            ("x: decimal(10, 0)", err),
            ("x: array<string>", 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_array_output_type_casting(self):
        class TestUDTF:
            def eval(self):
                yield ([0, 1.1, 2],)

        err = "UDTF_ARROW_DATA_CONVERSION_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="[0, 1.1, 2]")]),
            ("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>", err),
            ("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_DATA_CONVERSION_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=str({"a": 0, "b": 1.1, "c": 2}))]),
            ("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>", err),
            ("x: map<string,string>", [Row(x={"a": "0", "b": "1.1", "c": "2"})]),
            ("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_DATA_CONVERSION_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="{'a': 0, 'b': 1.1, 'c': 2}")]),
            ("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>", err),
            ("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=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_DATA_CONVERSION_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="Row(a=0, b=1.1, c=2)")]),
            ("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>", 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=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_DATA_CONVERSION_ERROR"):
                    udtf(TestUDTF, returnType=ret_type)().collect()

    def test_decimal_round(self):
        with self.sql_conf(
            {"spark.sql.legacy.execution.pythonUDTF.pandas.conversion.enabled": False}
        ):

            @udtf(returnType="a: DOUBLE, d: DECIMAL(38, 18)")
            class Float2Decimal:
                def eval(self, v: float):
                    yield v, Decimal(v)

            rounded = Float2Decimal(lit(1.234)).first().d
            self.assertEqual(rounded, Decimal("1.233999999999999986"))


class UDTFArrowTests(UDTFArrowTestsMixin, ReusedSQLTestCase):
    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        cls.spark.conf.set("spark.sql.execution.pythonUDTF.arrow.enabled", "true")
        cls.spark.conf.set(
            "spark.sql.legacy.execution.pythonUDTF.pandas.conversion.enabled", "false"
        )

    @classmethod
    def tearDownClass(cls):
        try:
            cls.spark.conf.unset("spark.sql.execution.pythonUDTF.arrow.enabled")
            cls.spark.conf.unset("spark.sql.legacy.execution.pythonUDTF.pandas.conversion.enabled")
        finally:
            super().tearDownClass()


if __name__ == "__main__":
    from pyspark.testing import main

    main()