#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import unittest

from pyspark.errors import PySparkValueError
from pyspark.sql import functions as sf
from pyspark.testing import assertDataFrameEqual
from pyspark.testing.sqlutils import ReusedSQLTestCase


class TVFTestsMixin:
    def test_explode(self):
        actual = self.spark.tvf.explode(sf.array(sf.lit(1), sf.lit(2)))
        expected = self.spark.sql("""SELECT * FROM explode(array(1, 2))""")
        assertDataFrameEqual(actual=actual, expected=expected)

        actual = self.spark.tvf.explode(
            sf.create_map(sf.lit("a"), sf.lit(1), sf.lit("b"), sf.lit(2))
        )
        expected = self.spark.sql("""SELECT * FROM explode(map('a', 1, 'b', 2))""")
        assertDataFrameEqual(actual=actual, expected=expected)

        # empty
        actual = self.spark.tvf.explode(sf.array())
        expected = self.spark.sql("""SELECT * FROM explode(array())""")
        assertDataFrameEqual(actual=actual, expected=expected)

        actual = self.spark.tvf.explode(sf.create_map())
        expected = self.spark.sql("""SELECT * FROM explode(map())""")
        assertDataFrameEqual(actual=actual, expected=expected)

        # null
        actual = self.spark.tvf.explode(sf.lit(None).astype("array<int>"))
        expected = self.spark.sql("""SELECT * FROM explode(null :: array<int>)""")
        assertDataFrameEqual(actual=actual, expected=expected)

        actual = self.spark.tvf.explode(sf.lit(None).astype("map<string, int>"))
        expected = self.spark.sql("""SELECT * FROM explode(null :: map<string, int>)""")
        assertDataFrameEqual(actual=actual, expected=expected)

    def test_explode_with_lateral_join(self):
        with self.tempView("t1", "t2"):
            t1 = self.spark.sql("VALUES (0, 1), (1, 2) AS t1(c1, c2)")
            t1.createOrReplaceTempView("t1")
            t3 = self.spark.sql(
                "VALUES (0, ARRAY(0, 1)), (1, ARRAY(2)), (2, ARRAY()), (null, ARRAY(4)) "
                "AS t3(c1, c2)"
            )
            t3.createOrReplaceTempView("t3")

            assertDataFrameEqual(
                t1.lateralJoin(
                    self.spark.tvf.explode(sf.array(sf.col("c1").outer(), sf.col("c2").outer()))
                    .toDF("c3")
                    .alias("t2")
                ),
                self.spark.sql("""SELECT * FROM t1, LATERAL EXPLODE(ARRAY(c1, c2)) t2(c3)"""),
            )
            assertDataFrameEqual(
                t3.lateralJoin(self.spark.tvf.explode(sf.col("c2").outer()).toDF("v").alias("t2")),
                self.spark.sql("""SELECT * FROM t3, LATERAL EXPLODE(c2) t2(v)"""),
            )
            assertDataFrameEqual(
                self.spark.tvf.explode(sf.array(sf.lit(1), sf.lit(2)))
                .toDF("v")
                .lateralJoin(
                    self.spark.range(1).select((sf.col("v").outer() + sf.lit(1)).alias("v2"))
                ),
                self.spark.sql(
                    """SELECT * FROM EXPLODE(ARRAY(1, 2)) t(v), LATERAL (SELECT v + 1 AS v2)"""
                ),
            )

    def test_explode_outer(self):
        actual = self.spark.tvf.explode_outer(sf.array(sf.lit(1), sf.lit(2)))
        expected = self.spark.sql("""SELECT * FROM explode_outer(array(1, 2))""")
        assertDataFrameEqual(actual=actual, expected=expected)

        actual = self.spark.tvf.explode_outer(
            sf.create_map(sf.lit("a"), sf.lit(1), sf.lit("b"), sf.lit(2))
        )
        expected = self.spark.sql("""SELECT * FROM explode_outer(map('a', 1, 'b', 2))""")
        assertDataFrameEqual(actual=actual, expected=expected)

        # empty
        actual = self.spark.tvf.explode_outer(sf.array())
        expected = self.spark.sql("""SELECT * FROM explode_outer(array())""")
        assertDataFrameEqual(actual=actual, expected=expected)

        actual = self.spark.tvf.explode_outer(sf.create_map())
        expected = self.spark.sql("""SELECT * FROM explode_outer(map())""")
        assertDataFrameEqual(actual=actual, expected=expected)

        # null
        actual = self.spark.tvf.explode_outer(sf.lit(None).astype("array<int>"))
        expected = self.spark.sql("""SELECT * FROM explode_outer(null :: array<int>)""")
        assertDataFrameEqual(actual=actual, expected=expected)

        actual = self.spark.tvf.explode_outer(sf.lit(None).astype("map<string, int>"))
        expected = self.spark.sql("""SELECT * FROM explode_outer(null :: map<string, int>)""")
        assertDataFrameEqual(actual=actual, expected=expected)

    def test_explode_outer_with_lateral_join(self):
        with self.tempView("t1", "t2"):
            t1 = self.spark.sql("VALUES (0, 1), (1, 2) AS t1(c1, c2)")
            t1.createOrReplaceTempView("t1")
            t3 = self.spark.sql(
                "VALUES (0, ARRAY(0, 1)), (1, ARRAY(2)), (2, ARRAY()), (null, ARRAY(4)) "
                "AS t3(c1, c2)"
            )
            t3.createOrReplaceTempView("t3")

            assertDataFrameEqual(
                t1.lateralJoin(
                    self.spark.tvf.explode_outer(
                        sf.array(sf.col("c1").outer(), sf.col("c2").outer())
                    )
                    .toDF("c3")
                    .alias("t2")
                ),
                self.spark.sql("""SELECT * FROM t1, LATERAL EXPLODE_OUTER(ARRAY(c1, c2)) t2(c3)"""),
            )
            assertDataFrameEqual(
                t3.lateralJoin(
                    self.spark.tvf.explode_outer(sf.col("c2").outer()).toDF("v").alias("t2")
                ),
                self.spark.sql("""SELECT * FROM t3, LATERAL EXPLODE_OUTER(c2) t2(v)"""),
            )
            assertDataFrameEqual(
                self.spark.tvf.explode_outer(sf.array(sf.lit(1), sf.lit(2)))
                .toDF("v")
                .lateralJoin(
                    self.spark.range(1).select((sf.col("v").outer() + sf.lit(1)).alias("v2"))
                ),
                self.spark.sql(
                    """
                    SELECT * FROM EXPLODE_OUTER(ARRAY(1, 2)) t(v), LATERAL (SELECT v + 1 AS v2)
                    """
                ),
            )

    def test_inline(self):
        actual = self.spark.tvf.inline(
            sf.array(sf.struct(sf.lit(1), sf.lit("a")), sf.struct(sf.lit(2), sf.lit("b")))
        )
        expected = self.spark.sql("""SELECT * FROM inline(array(struct(1, 'a'), struct(2, 'b')))""")
        assertDataFrameEqual(actual=actual, expected=expected)

        actual = self.spark.tvf.inline(sf.array().astype("array<struct<a:int,b:int>>"))
        expected = self.spark.sql("""SELECT * FROM inline(array() :: array<struct<a:int,b:int>>)""")
        assertDataFrameEqual(actual=actual, expected=expected)

        actual = self.spark.tvf.inline(
            sf.array(
                sf.named_struct(sf.lit("a"), sf.lit(1), sf.lit("b"), sf.lit(2)),
                sf.lit(None),
                sf.named_struct(sf.lit("a"), sf.lit(3), sf.lit("b"), sf.lit(4)),
            )
        )
        expected = self.spark.sql(
            """
            SELECT * FROM
              inline(array(named_struct('a', 1, 'b', 2), null, named_struct('a', 3, 'b', 4)))
            """
        )
        assertDataFrameEqual(actual=actual, expected=expected)

    def test_inline_with_lateral_join(self):
        with self.tempView("array_struct"):
            array_struct = self.spark.sql(
                """
                VALUES
                (1, ARRAY(STRUCT(1, 'a'), STRUCT(2, 'b'))),
                (2, ARRAY()),
                (3, ARRAY(STRUCT(3, 'c'))) AS array_struct(id, arr)
                """
            )
            array_struct.createOrReplaceTempView("array_struct")

            assertDataFrameEqual(
                array_struct.lateralJoin(self.spark.tvf.inline(sf.col("arr").outer())),
                self.spark.sql("""SELECT * FROM array_struct JOIN LATERAL INLINE(arr)"""),
            )
            assertDataFrameEqual(
                array_struct.lateralJoin(
                    self.spark.tvf.inline(sf.col("arr").outer()).toDF("k", "v").alias("t"),
                    sf.col("id") == sf.col("k"),
                    "left",
                ),
                self.spark.sql(
                    """
                    SELECT * FROM array_struct LEFT JOIN LATERAL INLINE(arr) t(k, v) ON id = k
                    """
                ),
            )

    def test_inline_outer(self):
        actual = self.spark.tvf.inline_outer(
            sf.array(sf.struct(sf.lit(1), sf.lit("a")), sf.struct(sf.lit(2), sf.lit("b")))
        )
        expected = self.spark.sql(
            """SELECT * FROM inline_outer(array(struct(1, 'a'), struct(2, 'b')))"""
        )
        assertDataFrameEqual(actual=actual, expected=expected)

        actual = self.spark.tvf.inline_outer(sf.array().astype("array<struct<a:int,b:int>>"))
        expected = self.spark.sql(
            """SELECT * FROM inline_outer(array() :: array<struct<a:int,b:int>>)"""
        )
        assertDataFrameEqual(actual=actual, expected=expected)

        actual = self.spark.tvf.inline_outer(
            sf.array(
                sf.named_struct(sf.lit("a"), sf.lit(1), sf.lit("b"), sf.lit(2)),
                sf.lit(None),
                sf.named_struct(sf.lit("a"), sf.lit(3), sf.lit("b"), sf.lit(4)),
            )
        )
        expected = self.spark.sql(
            """
            SELECT * FROM
              inline_outer(array(named_struct('a', 1, 'b', 2), null, named_struct('a', 3, 'b', 4)))
            """
        )
        assertDataFrameEqual(actual=actual, expected=expected)

    def test_inline_outer_with_lateral_join(self):
        with self.tempView("array_struct"):
            array_struct = self.spark.sql(
                """
                VALUES
                (1, ARRAY(STRUCT(1, 'a'), STRUCT(2, 'b'))),
                (2, ARRAY()),
                (3, ARRAY(STRUCT(3, 'c'))) AS array_struct(id, arr)
                """
            )
            array_struct.createOrReplaceTempView("array_struct")

            assertDataFrameEqual(
                array_struct.lateralJoin(self.spark.tvf.inline_outer(sf.col("arr").outer())),
                self.spark.sql("""SELECT * FROM array_struct JOIN LATERAL INLINE_OUTER(arr)"""),
            )
            assertDataFrameEqual(
                array_struct.lateralJoin(
                    self.spark.tvf.inline_outer(sf.col("arr").outer()).toDF("k", "v").alias("t"),
                    sf.col("id") == sf.col("k"),
                    "left",
                ),
                self.spark.sql(
                    """
                    SELECT * FROM array_struct LEFT JOIN LATERAL INLINE_OUTER(arr) t(k, v) ON id = k
                    """
                ),
            )

    def test_json_tuple(self):
        actual = self.spark.tvf.json_tuple(sf.lit('{"a":1, "b":2}'), sf.lit("a"), sf.lit("b"))
        expected = self.spark.sql("""SELECT json_tuple('{"a":1, "b":2}', 'a', 'b')""")
        assertDataFrameEqual(actual=actual, expected=expected)

        with self.assertRaises(PySparkValueError) as pe:
            self.spark.tvf.json_tuple(sf.lit('{"a":1, "b":2}'))

        self.check_error(
            exception=pe.exception,
            errorClass="CANNOT_BE_EMPTY",
            messageParameters={"item": "field"},
        )

    def test_json_tuple_with_lateral_join(self):
        with self.tempView("json_table"):
            json_table = self.spark.sql(
                """
                VALUES
                ('1', '{"f1": "1", "f2": "2", "f3": 3, "f5": 5.23}'),
                ('2', '{"f1": "1", "f3": "3", "f2": 2, "f4": 4.01}'),
                ('3', '{"f1": 3, "f4": "4", "f3": "3", "f2": 2, "f5": 5.01}'),
                ('4', cast(null as string)),
                ('5', '{"f1": null, "f5": ""}'),
                ('6', '[invalid JSON string]') AS json_table(key, jstring)
                """
            )
            json_table.createOrReplaceTempView("json_table")

            assertDataFrameEqual(
                json_table.alias("t1")
                .lateralJoin(
                    self.spark.tvf.json_tuple(
                        sf.col("jstring").outer(),
                        sf.lit("f1"),
                        sf.lit("f2"),
                        sf.lit("f3"),
                        sf.lit("f4"),
                        sf.lit("f5"),
                    ).alias("t2")
                )
                .select("t1.key", "t2.*"),
                self.spark.sql(
                    """
                    SELECT t1.key, t2.* FROM json_table t1,
                    LATERAL json_tuple(t1.jstring, 'f1', 'f2', 'f3', 'f4', 'f5') t2
                    """
                ),
            )
            assertDataFrameEqual(
                json_table.alias("t1")
                .lateralJoin(
                    self.spark.tvf.json_tuple(
                        sf.col("jstring").outer(),
                        sf.lit("f1"),
                        sf.lit("f2"),
                        sf.lit("f3"),
                        sf.lit("f4"),
                        sf.lit("f5"),
                    ).alias("t2")
                )
                .where(sf.col("t2.c0").isNotNull())
                .select("t1.key", "t2.*"),
                self.spark.sql(
                    """
                    SELECT t1.key, t2.* FROM json_table t1,
                    LATERAL json_tuple(t1.jstring, 'f1', 'f2', 'f3', 'f4', 'f5') t2
                    WHERE t2.c0 IS NOT NULL
                    """
                ),
            )

    def test_posexplode(self):
        actual = self.spark.tvf.posexplode(sf.array(sf.lit(1), sf.lit(2)))
        expected = self.spark.sql("""SELECT * FROM posexplode(array(1, 2))""")
        assertDataFrameEqual(actual=actual, expected=expected)

        actual = self.spark.tvf.posexplode(
            sf.create_map(sf.lit("a"), sf.lit(1), sf.lit("b"), sf.lit(2))
        )
        expected = self.spark.sql("""SELECT * FROM posexplode(map('a', 1, 'b', 2))""")
        assertDataFrameEqual(actual=actual, expected=expected)

        # empty
        actual = self.spark.tvf.posexplode(sf.array())
        expected = self.spark.sql("""SELECT * FROM posexplode(array())""")
        assertDataFrameEqual(actual=actual, expected=expected)

        actual = self.spark.tvf.posexplode(sf.create_map())
        expected = self.spark.sql("""SELECT * FROM posexplode(map())""")
        assertDataFrameEqual(actual=actual, expected=expected)

        # null
        actual = self.spark.tvf.posexplode(sf.lit(None).astype("array<int>"))
        expected = self.spark.sql("""SELECT * FROM posexplode(null :: array<int>)""")
        assertDataFrameEqual(actual=actual, expected=expected)

        actual = self.spark.tvf.posexplode(sf.lit(None).astype("map<string, int>"))
        expected = self.spark.sql("""SELECT * FROM posexplode(null :: map<string, int>)""")
        assertDataFrameEqual(actual=actual, expected=expected)

    def test_posexplode_with_lateral_join(self):
        with self.tempView("t1", "t2"):
            t1 = self.spark.sql("VALUES (0, 1), (1, 2) AS t1(c1, c2)")
            t1.createOrReplaceTempView("t1")
            t3 = self.spark.sql(
                "VALUES (0, ARRAY(0, 1)), (1, ARRAY(2)), (2, ARRAY()), (null, ARRAY(4)) "
                "AS t3(c1, c2)"
            )
            t3.createOrReplaceTempView("t3")

            assertDataFrameEqual(
                t1.lateralJoin(
                    self.spark.tvf.posexplode(sf.array(sf.col("c1").outer(), sf.col("c2").outer()))
                ),
                self.spark.sql("""SELECT * FROM t1, LATERAL POSEXPLODE(ARRAY(c1, c2))"""),
            )
            assertDataFrameEqual(
                t3.lateralJoin(self.spark.tvf.posexplode(sf.col("c2").outer())),
                self.spark.sql("""SELECT * FROM t3, LATERAL POSEXPLODE(c2)"""),
            )
            assertDataFrameEqual(
                self.spark.tvf.posexplode(sf.array(sf.lit(1), sf.lit(2)))
                .toDF("p", "v")
                .lateralJoin(
                    self.spark.range(1).select((sf.col("v").outer() + sf.lit(1)).alias("v2"))
                ),
                self.spark.sql(
                    """
                    SELECT * FROM POSEXPLODE(ARRAY(1, 2)) t(p, v), LATERAL (SELECT v + 1 AS v2)
                    """
                ),
            )

    def test_posexplode_outer(self):
        actual = self.spark.tvf.posexplode_outer(sf.array(sf.lit(1), sf.lit(2)))
        expected = self.spark.sql("""SELECT * FROM posexplode_outer(array(1, 2))""")
        assertDataFrameEqual(actual=actual, expected=expected)

        actual = self.spark.tvf.posexplode_outer(
            sf.create_map(sf.lit("a"), sf.lit(1), sf.lit("b"), sf.lit(2))
        )
        expected = self.spark.sql("""SELECT * FROM posexplode_outer(map('a', 1, 'b', 2))""")
        assertDataFrameEqual(actual=actual, expected=expected)

        # empty
        actual = self.spark.tvf.posexplode_outer(sf.array())
        expected = self.spark.sql("""SELECT * FROM posexplode_outer(array())""")
        assertDataFrameEqual(actual=actual, expected=expected)

        actual = self.spark.tvf.posexplode_outer(sf.create_map())
        expected = self.spark.sql("""SELECT * FROM posexplode_outer(map())""")
        assertDataFrameEqual(actual=actual, expected=expected)

        # null
        actual = self.spark.tvf.posexplode_outer(sf.lit(None).astype("array<int>"))
        expected = self.spark.sql("""SELECT * FROM posexplode_outer(null :: array<int>)""")
        assertDataFrameEqual(actual=actual, expected=expected)

        actual = self.spark.tvf.posexplode_outer(sf.lit(None).astype("map<string, int>"))
        expected = self.spark.sql("""SELECT * FROM posexplode_outer(null :: map<string, int>)""")
        assertDataFrameEqual(actual=actual, expected=expected)

    def test_posexplode_outer_with_lateral_join(self):
        with self.tempView("t1", "t2"):
            t1 = self.spark.sql("VALUES (0, 1), (1, 2) AS t1(c1, c2)")
            t1.createOrReplaceTempView("t1")
            t3 = self.spark.sql(
                "VALUES (0, ARRAY(0, 1)), (1, ARRAY(2)), (2, ARRAY()), (null, ARRAY(4)) "
                "AS t3(c1, c2)"
            )
            t3.createOrReplaceTempView("t3")

            assertDataFrameEqual(
                t1.lateralJoin(
                    self.spark.tvf.posexplode_outer(
                        sf.array(sf.col("c1").outer(), sf.col("c2").outer())
                    )
                ),
                self.spark.sql("""SELECT * FROM t1, LATERAL POSEXPLODE_OUTER(ARRAY(c1, c2))"""),
            )
            assertDataFrameEqual(
                t3.lateralJoin(self.spark.tvf.posexplode_outer(sf.col("c2").outer())),
                self.spark.sql("""SELECT * FROM t3, LATERAL POSEXPLODE_OUTER(c2)"""),
            )
            assertDataFrameEqual(
                self.spark.tvf.posexplode_outer(sf.array(sf.lit(1), sf.lit(2)))
                .toDF("p", "v")
                .lateralJoin(
                    self.spark.range(1).select((sf.col("v").outer() + sf.lit(1)).alias("v2"))
                ),
                self.spark.sql(
                    """
                    SELECT * FROM POSEXPLODE_OUTER(ARRAY(1, 2)) t(p, v),
                        LATERAL (SELECT v + 1 AS v2)
                    """
                ),
            )

    def test_stack(self):
        actual = self.spark.tvf.stack(sf.lit(2), sf.lit(1), sf.lit(2), sf.lit(3))
        expected = self.spark.sql("""SELECT * FROM stack(2, 1, 2, 3)""")
        assertDataFrameEqual(actual=actual, expected=expected)

    def test_stack_with_lateral_join(self):
        with self.tempView("t1", "t3"):
            t1 = self.spark.sql("VALUES (0, 1), (1, 2) AS t1(c1, c2)")
            t1.createOrReplaceTempView("t1")
            t3 = self.spark.sql(
                "VALUES (0, ARRAY(0, 1)), (1, ARRAY(2)), (2, ARRAY()), (null, ARRAY(4)) "
                "AS t3(c1, c2)"
            )
            t3.createOrReplaceTempView("t3")

            assertDataFrameEqual(
                t1.lateralJoin(
                    self.spark.tvf.stack(
                        sf.lit(2),
                        sf.lit("Key"),
                        sf.col("c1").outer(),
                        sf.lit("Value"),
                        sf.col("c2").outer(),
                    ).alias("t")
                ).select("t.*"),
                self.spark.sql(
                    """SELECT t.* FROM t1, LATERAL stack(2, 'Key', c1, 'Value', c2) t"""
                ),
            )
            assertDataFrameEqual(
                t1.lateralJoin(
                    self.spark.tvf.stack(sf.lit(1), sf.col("c1").outer(), sf.col("c2").outer())
                    .toDF("x", "y")
                    .alias("t")
                ).select("t.*"),
                self.spark.sql("""SELECT t.* FROM t1 JOIN LATERAL stack(1, c1, c2) t(x, y)"""),
            )
            assertDataFrameEqual(
                t1.join(t3, sf.col("t1.c1") == sf.col("t3.c1"))
                .lateralJoin(
                    self.spark.tvf.stack(
                        sf.lit(1), sf.col("t1.c2").outer(), sf.col("t3.c2").outer()
                    ).alias("t")
                )
                .select("t.*"),
                self.spark.sql(
                    """
                    SELECT t.* FROM t1 JOIN t3 ON t1.c1 = t3.c1
                        JOIN LATERAL stack(1, t1.c2, t3.c2) t
                    """
                ),
            )

    def test_collations(self):
        actual = self.spark.tvf.collations()
        expected = self.spark.sql("""SELECT * FROM collations()""")
        assertDataFrameEqual(actual=actual, expected=expected)

    def test_sql_keywords(self):
        actual = self.spark.tvf.sql_keywords()
        expected = self.spark.sql("""SELECT * FROM sql_keywords()""")
        assertDataFrameEqual(actual=actual, expected=expected)

    def test_variant_explode(self):
        actual = self.spark.tvf.variant_explode(sf.parse_json(sf.lit('["hello", "world"]')))
        expected = self.spark.sql(
            """SELECT * FROM variant_explode(parse_json('["hello", "world"]'))"""
        )
        assertDataFrameEqual(actual=actual, expected=expected)

        actual = self.spark.tvf.variant_explode(sf.parse_json(sf.lit('{"a": true, "b": 3.14}')))
        expected = self.spark.sql(
            """SELECT * FROM variant_explode(parse_json('{"a": true, "b": 3.14}'))"""
        )
        assertDataFrameEqual(actual=actual, expected=expected)

        # empty
        actual = self.spark.tvf.variant_explode(sf.parse_json(sf.lit("[]")))
        expected = self.spark.sql("""SELECT * FROM variant_explode(parse_json('[]'))""")
        assertDataFrameEqual(actual=actual, expected=expected)

        actual = self.spark.tvf.variant_explode(sf.parse_json(sf.lit("{}")))
        expected = self.spark.sql("""SELECT * FROM variant_explode(parse_json('{}'))""")
        assertDataFrameEqual(actual=actual, expected=expected)

        # null
        actual = self.spark.tvf.variant_explode(sf.lit(None).astype("variant"))
        expected = self.spark.sql("""SELECT * FROM variant_explode(null :: variant)""")
        assertDataFrameEqual(actual=actual, expected=expected)

        # not a variant object/array
        actual = self.spark.tvf.variant_explode(sf.parse_json(sf.lit("1")))
        expected = self.spark.sql("""SELECT * FROM variant_explode(parse_json('1'))""")
        assertDataFrameEqual(actual=actual, expected=expected)

    def test_variant_explode_with_lateral_join(self):
        with self.tempView("variant_table"):
            variant_table = self.spark.sql(
                """
                SELECT id, parse_json(v) AS v FROM VALUES
                    (0, '["hello", "world"]'), (1, '{"a": true, "b": 3.14}'),
                    (2, '[]'), (3, '{}'),
                    (4, NULL), (5, '1')
                    AS t(id, v)
                """
            )
            variant_table.createOrReplaceTempView("variant_table")

            assertDataFrameEqual(
                variant_table.alias("t1")
                .lateralJoin(self.spark.tvf.variant_explode(sf.col("v").outer()).alias("t"))
                .select("t1.id", "t.*"),
                self.spark.sql(
                    """
                    SELECT t1.id, t.* FROM variant_table AS t1,
                        LATERAL variant_explode(v) AS t
                    """
                ),
            )

    def test_variant_explode_outer(self):
        actual = self.spark.tvf.variant_explode_outer(sf.parse_json(sf.lit('["hello", "world"]')))
        expected = self.spark.sql(
            """SELECT * FROM variant_explode_outer(parse_json('["hello", "world"]'))"""
        )
        assertDataFrameEqual(actual=actual, expected=expected)

        actual = self.spark.tvf.variant_explode_outer(
            sf.parse_json(sf.lit('{"a": true, "b": 3.14}'))
        )
        expected = self.spark.sql(
            """SELECT * FROM variant_explode_outer(parse_json('{"a": true, "b": 3.14}'))"""
        )
        assertDataFrameEqual(actual=actual, expected=expected)

        # empty
        actual = self.spark.tvf.variant_explode_outer(sf.parse_json(sf.lit("[]")))
        expected = self.spark.sql("""SELECT * FROM variant_explode_outer(parse_json('[]'))""")
        assertDataFrameEqual(actual=actual, expected=expected)

        actual = self.spark.tvf.variant_explode_outer(sf.parse_json(sf.lit("{}")))
        expected = self.spark.sql("""SELECT * FROM variant_explode_outer(parse_json('{}'))""")
        assertDataFrameEqual(actual=actual, expected=expected)

        # null
        actual = self.spark.tvf.variant_explode_outer(sf.lit(None).astype("variant"))
        expected = self.spark.sql("""SELECT * FROM variant_explode_outer(null :: variant)""")
        assertDataFrameEqual(actual=actual, expected=expected)

        # not a variant object/array
        actual = self.spark.tvf.variant_explode_outer(sf.parse_json(sf.lit("1")))
        expected = self.spark.sql("""SELECT * FROM variant_explode_outer(parse_json('1'))""")
        assertDataFrameEqual(actual=actual, expected=expected)

    def test_variant_explode_outer_with_lateral_join(self):
        with self.tempView("variant_table"):
            variant_table = self.spark.sql(
                """
                SELECT id, parse_json(v) AS v FROM VALUES
                    (0, '["hello", "world"]'), (1, '{"a": true, "b": 3.14}'),
                    (2, '[]'), (3, '{}'),
                    (4, NULL), (5, '1')
                    AS t(id, v)
                """
            )
            variant_table.createOrReplaceTempView("variant_table")

            assertDataFrameEqual(
                variant_table.alias("t1")
                .lateralJoin(self.spark.tvf.variant_explode_outer(sf.col("v").outer()).alias("t"))
                .select("t1.id", "t.*"),
                self.spark.sql(
                    """
                    SELECT t1.id, t.* FROM variant_table AS t1,
                        LATERAL variant_explode_outer(v) AS t
                    """
                ),
            )


class TVFTests(TVFTestsMixin, ReusedSQLTestCase):
    pass


if __name__ == "__main__":
    from pyspark.sql.tests.test_tvf import *  # noqa: F401

    try:
        import xmlrunner  # type: ignore

        testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
    except ImportError:
        testRunner = None
    unittest.main(testRunner=testRunner, verbosity=2)