#
# 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 pyspark.errors import AnalysisException, QueryContextType, SparkRuntimeException
from pyspark.sql import functions as sf
from pyspark.testing import assertDataFrameEqual
from pyspark.testing.sqlutils import ReusedSQLTestCase


class SubqueryTestsMixin:
    @property
    def df1(self):
        return self.spark.createDataFrame(
            [
                (1, 2.0),
                (1, 2.0),
                (2, 1.0),
                (2, 2.0),
                (3, 3.0),
                (None, None),
                (None, 5.0),
                (6, None),
            ],
            ["a", "b"],
        )

    @property
    def df2(self):
        return self.spark.createDataFrame(
            [(2, 3.0), (2, 3.0), (3, 2.0), (4, 1.0), (None, None), (None, 5.0), (6, None)],
            ["c", "d"],
        )

    def test_noop_outer(self):
        assertDataFrameEqual(
            self.spark.range(1).select(sf.col("id").outer()),
            self.spark.range(1).select(sf.col("id")),
        )

        with self.assertRaises(AnalysisException) as pe:
            self.spark.range(1).select(sf.col("outer_col").outer()).collect()

        self.check_error(
            exception=pe.exception,
            errorClass="UNRESOLVED_COLUMN.WITH_SUGGESTION",
            messageParameters={"objectName": "`outer_col`", "proposal": "`id`"},
            query_context_type=QueryContextType.DataFrame,
            fragment="col",
        )

    def test_simple_uncorrelated_scalar_subquery(self):
        assertDataFrameEqual(
            self.spark.range(1).select(self.spark.range(1).select(sf.lit(1)).scalar().alias("b")),
            self.spark.sql("""select (select 1 as b) as b"""),
        )

        assertDataFrameEqual(
            self.spark.range(1).select(
                self.spark.range(1)
                .select(self.spark.range(1).select(sf.lit(1)).scalar() + 1)
                .scalar()
                + 1
            ),
            self.spark.sql("""select (select (select 1) + 1) + 1"""),
        )

        # string type
        assertDataFrameEqual(
            self.spark.range(1).select(self.spark.range(1).select(sf.lit("s")).scalar().alias("b")),
            self.spark.sql("""select (select 's' as s) as b"""),
        )

        # 0 rows
        assertDataFrameEqual(
            self.spark.range(1).select(
                self.spark.range(1).select(sf.lit("s")).limit(0).scalar().alias("b")
            ),
            self.spark.sql("""select (select 's' as s limit 0) as b"""),
        )

    def test_uncorrelated_scalar_subquery_with_view(self):
        with self.temp_view("subqueryData"):
            df = self.spark.createDataFrame(
                [(1, "one"), (2, "two"), (3, "three")], ["key", "value"]
            )
            df.createOrReplaceTempView("subqueryData")

            assertDataFrameEqual(
                self.spark.range(1).select(
                    self.spark.table("subqueryData")
                    .select("key")
                    .where(sf.col("key") > 2)
                    .orderBy("key")
                    .limit(1)
                    .scalar()
                    + 1
                ),
                self.spark.sql("""
                    select (select key from subqueryData where key > 2 order by key limit 1) + 1
                    """),
            )

            assertDataFrameEqual(
                self.spark.range(1).select(
                    (-self.spark.table("subqueryData").select(sf.max("key")).scalar()).alias(
                        "negative_max_key"
                    )
                ),
                self.spark.sql(
                    """select -(select max(key) from subqueryData) as negative_max_key"""
                ),
            )

            assertDataFrameEqual(
                self.spark.range(1).select(
                    self.spark.table("subqueryData").select("value").limit(0).scalar()
                ),
                self.spark.sql("""select (select value from subqueryData limit 0)"""),
            )

            assertDataFrameEqual(
                self.spark.range(1).select(
                    self.spark.table("subqueryData")
                    .where(
                        sf.col("key")
                        == self.spark.table("subqueryData").select(sf.max("key")).scalar() - 1
                    )
                    .select(sf.min("value"))
                    .scalar()
                ),
                self.spark.sql("""
                    select (
                        select min(value) from subqueryData
                        where key = (select max(key) from subqueryData) - 1
                    )
                    """),
            )

    def test_scalar_subquery_against_local_relations(self):
        with self.temp_view("t1", "t2"):
            self.spark.createDataFrame([(1, 1), (2, 2)], ["c1", "c2"]).createOrReplaceTempView("t1")
            self.spark.createDataFrame([(1, 1), (2, 2)], ["c1", "c2"]).createOrReplaceTempView("t2")

            assertDataFrameEqual(
                self.spark.table("t1").select(
                    self.spark.range(1).select(sf.lit(1).alias("col")).scalar()
                ),
                self.spark.sql("""SELECT (select 1 as col) from t1"""),
            )

            assertDataFrameEqual(
                self.spark.table("t1").select(self.spark.table("t2").select(sf.max("c1")).scalar()),
                self.spark.sql("""SELECT (select max(c1) from t2) from t1"""),
            )

            assertDataFrameEqual(
                self.spark.table("t1").select(
                    sf.lit(1) + self.spark.range(1).select(sf.lit(1).alias("col")).scalar()
                ),
                self.spark.sql("""SELECT 1 + (select 1 as col) from t1"""),
            )

            assertDataFrameEqual(
                self.spark.table("t1").select(
                    "c1", self.spark.table("t2").select(sf.max("c1")).scalar() + sf.col("c2")
                ),
                self.spark.sql("""SELECT c1, (select max(c1) from t2) + c2 from t1"""),
            )

            assertDataFrameEqual(
                self.spark.table("t1").select(
                    "c1",
                    (
                        self.spark.table("t2")
                        .where(sf.col("t1.c2").outer() == sf.col("t2.c2"))
                        .select(sf.max("c1"))
                        .scalar()
                    ),
                ),
                self.spark.sql(
                    """SELECT c1, (select max(c1) from t2 where t1.c2 = t2.c2) from t1"""
                ),
            )

    def test_correlated_scalar_subquery(self):
        with self.temp_view("l", "r"):
            self.df1.createOrReplaceTempView("l")
            self.df2.createOrReplaceTempView("r")

            with self.subTest("in where"):
                for cond in [
                    sf.col("a").outer() == sf.col("c"),
                    (sf.col("a") == sf.col("c")).outer(),
                    sf.expr("a = c").outer(),
                ]:
                    with self.subTest(cond=cond):
                        assertDataFrameEqual(
                            self.spark.table("l").where(
                                sf.col("b")
                                < self.spark.table("r").where(cond).select(sf.max("d")).scalar()
                            ),
                            self.spark.sql(
                                """select * from l where b < (select max(d) from r where a = c)"""
                            ),
                        )

            with self.subTest("in select"):
                df1 = self.spark.table("l").alias("t1")
                df2 = self.spark.table("l").alias("t2")

                for cond in [
                    sf.col("t1.a") == sf.col("t2.a").outer(),
                    (sf.col("t1.a") == sf.col("t2.a")).outer(),
                    sf.expr("t1.a = t2.a").outer(),
                ]:
                    with self.subTest(cond=cond):
                        assertDataFrameEqual(
                            df1.select(
                                "a",
                                df2.where(cond).select(sf.sum("b")).scalar().alias("sum_b"),
                            ),
                            self.spark.sql("""
                                select
                                    a, (select sum(b) from l t2 where t2.a = t1.a) sum_b
                                from l t1
                                """),
                        )

            with self.subTest("without .outer()"):
                assertDataFrameEqual(
                    self.spark.table("l").select(
                        "a",
                        (
                            self.spark.table("r")
                            .where(sf.col("b") == sf.col("a").outer())
                            .select(sf.sum("d"))
                            .scalar()
                            .alias("sum_d")
                        ),
                    ),
                    self.spark.sql(
                        """select a, (select sum(d) from r where b = l.a) sum_d from l"""
                    ),
                )

            with self.subTest("in select (null safe)"):
                df1 = self.spark.table("l").alias("t1")
                df2 = self.spark.table("l").alias("t2")

                assertDataFrameEqual(
                    df1.select(
                        "a",
                        (
                            df2.where(sf.col("t2.a").eqNullSafe(sf.col("t1.a").outer()))
                            .select(sf.sum("b"))
                            .scalar()
                            .alias("sum_b")
                        ),
                    ),
                    self.spark.sql("""
                        select a, (select sum(b) from l l2 where l2.a <=> l1.a) sum_b from l l1
                        """),
                )

            with self.subTest("in aggregate"):
                assertDataFrameEqual(
                    self.spark.table("l")
                    .groupBy(
                        "a",
                        (
                            self.spark.table("r")
                            .where(sf.col("a").outer() == sf.col("c"))
                            .select(sf.sum("d"))
                            .scalar()
                            .alias("sum_d")
                        ),
                    )
                    .agg({}),
                    self.spark.sql("""
                        select a, (select sum(d) from r where a = c) sum_d from l l1 group by 1, 2
                        """),
                )

            with self.subTest("non-aggregated"):
                df1 = self.spark.table("l").alias("t1")
                df2 = self.spark.table("l").alias("t2")

                with self.assertRaises(SparkRuntimeException) as pe:
                    df1.select(
                        "a",
                        df2.where(sf.col("t1.a") == sf.col("t2.a").outer()).select("b").scalar(),
                    ).collect()

                self.check_error(
                    exception=pe.exception,
                    errorClass="SCALAR_SUBQUERY_TOO_MANY_ROWS",
                    messageParameters={},
                )

            with self.subTest("non-equal"):
                df1 = self.spark.table("l").alias("t1")
                df2 = self.spark.table("l").alias("t2")

                assertDataFrameEqual(
                    df1.select(
                        "a",
                        (
                            df2.where(sf.col("t2.a") < sf.col("t1.a").outer())
                            .select(sf.sum("b"))
                            .scalar()
                            .alias("sum_b")
                        ),
                    ),
                    self.spark.sql(
                        """select a, (select sum(b) from l t2 where t2.a < t1.a) sum_b from l t1"""
                    ),
                )

            with self.subTest("disjunctive"):
                assertDataFrameEqual(
                    self.spark.table("l")
                    .where(
                        self.spark.table("r")
                        .where(
                            ((sf.col("a").outer() == sf.col("c")) & (sf.col("d") == sf.lit(2.0)))
                            | ((sf.col("a").outer() == sf.col("c")) & (sf.col("d") == sf.lit(1.0)))
                        )
                        .select(sf.count(sf.lit(1)))
                        .scalar()
                        > 0
                    )
                    .select("a"),
                    self.spark.sql("""
                        select a
                        from   l
                        where  (select count(*)
                                from   r
                                where (a = c and d = 2.0) or (a = c and d = 1.0)) > 0
                        """),
                )

    def test_exists_subquery(self):
        with self.temp_view("l", "r"):
            self.df1.createOrReplaceTempView("l")
            self.df2.createOrReplaceTempView("r")

            with self.subTest("EXISTS"):
                for cond in [
                    sf.col("a").outer() == sf.col("c"),
                    (sf.col("a") == sf.col("c")).outer(),
                    sf.expr("a = c").outer(),
                ]:
                    with self.subTest(cond=cond):
                        assertDataFrameEqual(
                            self.spark.table("l").where(self.spark.table("r").where(cond).exists()),
                            self.spark.sql(
                                """select * from l where exists (select * from r where l.a = r.c)"""
                            ),
                        )

                        assertDataFrameEqual(
                            self.spark.table("l").where(
                                self.spark.table("r").where(cond).exists()
                                & (sf.col("a") <= sf.lit(2))
                            ),
                            self.spark.sql("""
                        select * from l where exists (select * from r where l.a = r.c) and l.a <= 2
                        """),
                        )

            with self.subTest("NOT EXISTS"):
                assertDataFrameEqual(
                    self.spark.table("l").where(
                        ~self.spark.table("r").where(sf.col("a").outer() == sf.col("c")).exists()
                    ),
                    self.spark.sql(
                        """select * from l where not exists (select * from r where l.a = r.c)"""
                    ),
                )

                assertDataFrameEqual(
                    self.spark.table("l").where(
                        ~(
                            self.spark.table("r")
                            .where(
                                (sf.col("a").outer() == sf.col("c"))
                                & (sf.col("b").outer() < sf.col("d"))
                            )
                            .exists()
                        )
                    ),
                    self.spark.sql("""
                        select * from l
                            where not exists (select * from r where l.a = r.c and l.b < r.d)
                        """),
                )

            with self.subTest("EXISTS within OR"):
                assertDataFrameEqual(
                    self.spark.table("l").where(
                        self.spark.table("r").where(sf.col("a").outer() == sf.col("c")).exists()
                        | self.spark.table("r").where(sf.col("a").outer() == sf.col("c")).exists()
                    ),
                    self.spark.sql("""
                        select * from l where exists (select * from r where l.a = r.c)
                            or exists (select * from r where l.a = r.c)
                        """),
                )

                assertDataFrameEqual(
                    self.spark.table("l").where(
                        self.spark.table("r")
                        .where(
                            (sf.col("a").outer() == sf.col("c"))
                            & (sf.col("b").outer() < sf.col("d"))
                        )
                        .exists()
                        | self.spark.table("r").where(sf.col("a").outer() == sf.col("c")).exists()
                    ),
                    self.spark.sql("""
                        select * from l where exists (select * from r where l.a = r.c and l.b < r.d)
                            or exists (select * from r where l.a = r.c)
                        """),
                )

    def test_in_subquery(self):
        with self.temp_view("l", "r", "t"):
            self.df1.createOrReplaceTempView("l")
            self.df2.createOrReplaceTempView("r")
            self.spark.table("r").filter(
                sf.col("c").isNotNull() & sf.col("d").isNotNull()
            ).createOrReplaceTempView("t")

            with self.subTest("IN"):
                assertDataFrameEqual(
                    self.spark.table("l").where(
                        sf.col("l.a").isin(self.spark.table("r").select(sf.col("c")))
                    ),
                    self.spark.sql("""select * from l where l.a in (select c from r)"""),
                )
                assertDataFrameEqual(
                    self.spark.table("l").where(
                        sf.col("l.a").isin(
                            self.spark.table("r")
                            .where(sf.col("l.b").outer() < sf.col("r.d"))
                            .select(sf.col("c"))
                        )
                    ),
                    self.spark.sql(
                        """select * from l where l.a in (select c from r where l.b < r.d)"""
                    ),
                )
                assertDataFrameEqual(
                    self.spark.table("l").where(
                        sf.col("l.a").isin(self.spark.table("r").select("c"))
                        & (sf.col("l.a") > sf.lit(2))
                        & sf.col("l.b").isNotNull()
                    ),
                    self.spark.sql("""
                        select * from l
                        where l.a in (select c from r) and l.a > 2 and l.b is not null
                        """),
                )

            with self.subTest("IN with struct"), self.temp_view("ll", "rr"):
                self.spark.table("l").select(
                    "*", sf.struct("a", "b").alias("sab")
                ).createOrReplaceTempView("ll")
                self.spark.table("r").select(
                    "*", sf.struct(sf.col("c").alias("a"), sf.col("d").alias("b")).alias("scd")
                ).createOrReplaceTempView("rr")

                for col, values in [
                    (sf.col("sab"), "sab"),
                    (sf.struct(sf.struct(sf.col("a"), sf.col("b"))), "struct(struct(a, b))"),
                ]:
                    for df, query in [
                        (self.spark.table("rr").select(sf.col("scd")), "select scd from rr"),
                        (
                            self.spark.table("rr").select(
                                sf.struct(sf.col("c").alias("a"), sf.col("d").alias("b"))
                            ),
                            "select struct(c as a, d as b) from rr",
                        ),
                        (
                            self.spark.table("rr").select(sf.struct(sf.col("c"), sf.col("d"))),
                            "select struct(c, d) from rr",
                        ),
                    ]:
                        sql_query = f"""select a, b from ll where {values} in ({query})"""
                        with self.subTest(sql_query=sql_query):
                            assertDataFrameEqual(
                                self.spark.table("ll").where(col.isin(df)).select("a", "b"),
                                self.spark.sql(sql_query),
                            )

            with self.subTest("NOT IN"):
                assertDataFrameEqual(
                    self.spark.table("l").where(
                        ~sf.col("a").isin(self.spark.table("r").select("c"))
                    ),
                    self.spark.sql("""select * from l where a not in (select c from r)"""),
                )
                assertDataFrameEqual(
                    self.spark.table("l").where(
                        ~sf.col("a").isin(
                            self.spark.table("r").where(sf.col("c").isNotNull()).select(sf.col("c"))
                        )
                    ),
                    self.spark.sql(
                        """select * from l where a not in (select c from r where c is not null)"""
                    ),
                )
                assertDataFrameEqual(
                    self.spark.table("l").where(
                        (
                            ~sf.struct(sf.col("a"), sf.col("b")).isin(
                                self.spark.table("t").select(sf.col("c"), sf.col("d"))
                            )
                        )
                        & (sf.col("a") < sf.lit(4))
                    ),
                    self.spark.sql(
                        """select * from l where (a, b) not in (select c, d from t) and a < 4"""
                    ),
                )
                assertDataFrameEqual(
                    self.spark.table("l").where(
                        ~sf.struct(sf.col("a"), sf.col("b")).isin(
                            self.spark.table("r")
                            .where(sf.col("c") > sf.lit(10))
                            .select(sf.col("c"), sf.col("d"))
                        )
                    ),
                    self.spark.sql(
                        """select * from l where (a, b) not in (select c, d from r where c > 10)"""
                    ),
                )

            with self.subTest("IN within OR"):
                assertDataFrameEqual(
                    self.spark.table("l").where(
                        sf.col("l.a").isin(self.spark.table("r").select("c"))
                        | (
                            sf.col("l.a").isin(
                                self.spark.table("r")
                                .where(sf.col("l.b").outer() < sf.col("r.d"))
                                .select(sf.col("c"))
                            )
                        )
                    ),
                    self.spark.sql("""
                        select * from l
                        where l.a in (select c from r) or l.a in (select c from r where l.b < r.d)
                        """),
                )
                assertDataFrameEqual(
                    self.spark.table("l").where(
                        (~sf.col("a").isin(self.spark.table("r").select(sf.col("c"))))
                        | (
                            ~sf.col("a").isin(
                                self.spark.table("r")
                                .where(sf.col("c").isNotNull())
                                .select(sf.col("c"))
                            )
                        )
                    ),
                    self.spark.sql("""
                        select * from l
                        where a not in (select c from r)
                        or a not in (select c from r where c is not null)
                        """),
                )

            with self.subTest("complex IN"):
                assertDataFrameEqual(
                    self.spark.table("l").where(
                        ~sf.struct(sf.col("a"), sf.col("b")).isin(
                            self.spark.table("r").select(sf.col("c"), sf.col("d"))
                        )
                    ),
                    self.spark.sql("""select * from l where (a, b) not in (select c, d from r)"""),
                )
                assertDataFrameEqual(
                    self.spark.table("l").where(
                        (
                            ~sf.struct(sf.col("a"), sf.col("b")).isin(
                                self.spark.table("t").select(sf.col("c"), sf.col("d"))
                            )
                        )
                        & ((sf.col("a") + sf.col("b")).isNotNull())
                    ),
                    self.spark.sql("""
                        select * from l
                        where (a, b) not in (select c, d from t) and (a + b) is not null
                        """),
                )

            with self.subTest("same column in subquery"):
                assertDataFrameEqual(
                    self.spark.table("l")
                    .alias("l1")
                    .where(
                        sf.col("a").isin(
                            self.spark.table("l")
                            .where(sf.col("a") < sf.lit(3))
                            .groupBy(sf.col("a"))
                            .agg({})
                        )
                    )
                    .select(sf.col("a")),
                    self.spark.sql(
                        """select a from l l1 where a in (select a from l where a < 3 group by a)"""
                    ),
                )

            with self.subTest("col IN (NULL)"):
                assertDataFrameEqual(
                    self.spark.table("l").where(sf.col("a").isin(None)),
                    self.spark.sql("""SELECT * FROM l WHERE a IN (NULL)"""),
                )

    def test_scalar_subquery_with_missing_outer_reference(self):
        with self.temp_view("l", "r"):
            self.df1.createOrReplaceTempView("l")
            self.df2.createOrReplaceTempView("r")

            with self.assertRaises(AnalysisException) as pe:
                self.spark.table("l").select(
                    "a",
                    (
                        self.spark.table("r")
                        .where(sf.col("c") == sf.col("a"))
                        .select(sf.sum("d"))
                        .scalar()
                    ),
                ).collect()

            self.check_error(
                exception=pe.exception,
                errorClass="UNRESOLVED_COLUMN.WITH_SUGGESTION",
                messageParameters={"objectName": "`a`", "proposal": "`c`, `d`"},
                query_context_type=QueryContextType.DataFrame,
                fragment="col",
            )

    def table1(self):
        t1 = self.spark.sql("VALUES (0, 1), (1, 2) AS t1(c1, c2)")
        t1.createOrReplaceTempView("t1")
        return self.spark.table("t1")

    def table2(self):
        t2 = self.spark.sql("VALUES (0, 2), (0, 3) AS t2(c1, c2)")
        t2.createOrReplaceTempView("t2")
        return self.spark.table("t2")

    def table3(self):
        t3 = self.spark.sql(
            "VALUES (0, ARRAY(0, 1)), (1, ARRAY(2)), (2, ARRAY()), (null, ARRAY(4)) AS t3(c1, c2)"
        )
        t3.createOrReplaceTempView("t3")
        return self.spark.table("t3")

    def test_lateral_join_with_single_column_select(self):
        with self.temp_view("t1", "t2"):
            t1 = self.table1()
            t2 = self.table2()

            assertDataFrameEqual(
                t1.lateralJoin(self.spark.range(1).select(sf.col("c1").outer())),
                self.spark.sql("""SELECT * FROM t1, LATERAL (SELECT c1)"""),
            )
            assertDataFrameEqual(
                t1.lateralJoin(t2.select(sf.col("t1.c1").outer())),
                self.spark.sql("""SELECT * FROM t1, LATERAL (SELECT t1.c1 FROM t2)"""),
            )
            assertDataFrameEqual(
                t1.lateralJoin(t2.select(sf.col("t1.c1").outer() + sf.col("t2.c1"))),
                self.spark.sql("""SELECT * FROM t1, LATERAL (SELECT t1.c1 + t2.c1 FROM t2)"""),
            )

    def test_lateral_join_with_star_expansion(self):
        with self.temp_view("t1", "t2"):
            t1 = self.table1()
            t2 = self.table2()

            assertDataFrameEqual(
                t1.lateralJoin(self.spark.range(1).select().select(sf.col("*"))),
                self.spark.sql("""SELECT * FROM t1, LATERAL (SELECT *)"""),
            )
            assertDataFrameEqual(
                t1.lateralJoin(t2.select(sf.col("*"))),
                self.spark.sql("""SELECT * FROM t1, LATERAL (SELECT * FROM t2)"""),
            )
            assertDataFrameEqual(
                t1.lateralJoin(t2.select(sf.col("t1.*").outer(), sf.col("t2.*"))),
                self.spark.sql("""SELECT * FROM t1, LATERAL (SELECT t1.*, t2.* FROM t2)"""),
            )
            assertDataFrameEqual(
                t1.lateralJoin(t2.alias("t1").select(sf.col("t1.*"))),
                self.spark.sql("""SELECT * FROM t1, LATERAL (SELECT t1.* FROM t2 AS t1)"""),
            )

    def test_lateral_join_with_different_join_types(self):
        with self.temp_view("t1"):
            t1 = self.table1()

            assertDataFrameEqual(
                t1.lateralJoin(
                    self.spark.range(1).select(
                        (sf.col("c1").outer() + sf.col("c2").outer()).alias("c3")
                    ),
                    sf.col("c2") == sf.col("c3"),
                ),
                self.spark.sql(
                    """SELECT * FROM t1 JOIN LATERAL (SELECT c1 + c2 AS c3) ON c2 = c3"""
                ),
            )
            assertDataFrameEqual(
                t1.lateralJoin(
                    self.spark.range(1).select(
                        (sf.col("c1").outer() + sf.col("c2").outer()).alias("c3")
                    ),
                    sf.col("c2") == sf.col("c3"),
                    "left",
                ),
                self.spark.sql(
                    """SELECT * FROM t1 LEFT JOIN LATERAL (SELECT c1 + c2 AS c3) ON c2 = c3"""
                ),
            )
            assertDataFrameEqual(
                t1.lateralJoin(
                    self.spark.range(1).select(
                        (sf.col("c1").outer() + sf.col("c2").outer()).alias("c3")
                    ),
                    how="cross",
                ),
                self.spark.sql("""SELECT * FROM t1 CROSS JOIN LATERAL (SELECT c1 + c2 AS c3)"""),
            )

            with self.assertRaises(AnalysisException) as pe:
                t1.lateralJoin(
                    self.spark.range(1).select(
                        (sf.col("c1").outer() + sf.col("c2").outer()).alias("c3")
                    ),
                    how="right",
                ).collect()

            self.check_error(
                pe.exception,
                errorClass="UNSUPPORTED_JOIN_TYPE",
                messageParameters={
                    "typ": "right",
                    "supported": "'inner', 'leftouter', 'left', 'left_outer', 'cross'",
                },
            )

    def test_lateral_join_with_subquery_alias(self):
        with self.temp_view("t1"):
            t1 = self.table1()

            assertDataFrameEqual(
                t1.lateralJoin(
                    self.spark.range(1)
                    .select(sf.col("c1").outer(), sf.col("c2").outer())
                    .toDF("a", "b")
                    .alias("s")
                ).select("a", "b"),
                self.spark.sql("""SELECT a, b FROM t1, LATERAL (SELECT c1, c2) s(a, b)"""),
            )

    def test_lateral_join_with_correlated_predicates(self):
        with self.temp_view("t1", "t2"):
            t1 = self.table1()
            t2 = self.table2()

            assertDataFrameEqual(
                t1.lateralJoin(
                    t2.where(sf.col("t1.c1").outer() == sf.col("t2.c1")).select(sf.col("c2"))
                ),
                self.spark.sql(
                    """SELECT * FROM t1, LATERAL (SELECT c2 FROM t2 WHERE t1.c1 = t2.c1)"""
                ),
            )
            assertDataFrameEqual(
                t1.lateralJoin(
                    t2.where(sf.col("t1.c1").outer() < sf.col("t2.c1")).select(sf.col("c2"))
                ),
                self.spark.sql(
                    """SELECT * FROM t1, LATERAL (SELECT c2 FROM t2 WHERE t1.c1 < t2.c1)"""
                ),
            )

    def test_lateral_join_with_aggregation_and_correlated_predicates(self):
        with self.temp_view("t1", "t2"):
            t1 = self.table1()
            t2 = self.table2()

            assertDataFrameEqual(
                t1.lateralJoin(
                    t2.where(sf.col("t1.c2").outer() < sf.col("t2.c2")).select(
                        sf.max(sf.col("c2")).alias("m")
                    )
                ),
                self.spark.sql("""
                    SELECT * FROM t1, LATERAL (SELECT max(c2) AS m FROM t2 WHERE t1.c2 < t2.c2)
                    """),
            )

    def test_lateral_join_reference_preceding_from_clause_items(self):
        with self.temp_view("t1", "t2"):
            t1 = self.table1()
            t2 = self.table2()

            assertDataFrameEqual(
                t1.join(t2).lateralJoin(
                    self.spark.range(1).select(sf.col("t1.c2").outer() + sf.col("t2.c2").outer())
                ),
                self.spark.sql("""SELECT * FROM t1 JOIN t2 JOIN LATERAL (SELECT t1.c2 + t2.c2)"""),
            )

    def test_multiple_lateral_joins(self):
        with self.temp_view("t1"):
            t1 = self.table1()

            assertDataFrameEqual(
                t1.lateralJoin(
                    self.spark.range(1).select(
                        (sf.col("c1").outer() + sf.col("c2").outer()).alias("a")
                    )
                )
                .lateralJoin(
                    self.spark.range(1).select(
                        (sf.col("c1").outer() - sf.col("c2").outer()).alias("b")
                    )
                )
                .lateralJoin(
                    self.spark.range(1).select(
                        (sf.col("a").outer() * sf.col("b").outer()).alias("c")
                    )
                ),
                self.spark.sql("""
                    SELECT * FROM t1,
                    LATERAL (SELECT c1 + c2 AS a),
                    LATERAL (SELECT c1 - c2 AS b),
                    LATERAL (SELECT a * b AS c)
                    """),
            )

    def test_lateral_join_in_between_regular_joins(self):
        with self.temp_view("t1", "t2"):
            t1 = self.table1()
            t2 = self.table2()

            assertDataFrameEqual(
                t1.lateralJoin(
                    t2.where(sf.col("t1.c1").outer() == sf.col("t2.c1"))
                    .select(sf.col("c2"))
                    .alias("s"),
                    how="left",
                ).join(t1.alias("t3"), sf.col("s.c2") == sf.col("t3.c2"), how="left"),
                self.spark.sql("""
                    SELECT * FROM t1
                    LEFT OUTER JOIN LATERAL (SELECT c2 FROM t2 WHERE t1.c1 = t2.c1) s
                    LEFT OUTER JOIN t1 t3 ON s.c2 = t3.c2
                    """),
            )

    def test_nested_lateral_joins(self):
        with self.temp_view("t1", "t2"):
            t1 = self.table1()
            t2 = self.table2()

            assertDataFrameEqual(
                t1.lateralJoin(t2.lateralJoin(self.spark.range(1).select(sf.col("c1").outer()))),
                self.spark.sql(
                    """SELECT * FROM t1, LATERAL (SELECT * FROM t2, LATERAL (SELECT c1))"""
                ),
            )
            assertDataFrameEqual(
                t1.lateralJoin(
                    self.spark.range(1)
                    .select((sf.col("c1").outer() + sf.lit(1)).alias("c1"))
                    .lateralJoin(self.spark.range(1).select(sf.col("c1").outer()))
                ),
                self.spark.sql("""
                    SELECT * FROM t1,
                    LATERAL (SELECT * FROM (SELECT c1 + 1 AS c1), LATERAL (SELECT c1))
                    """),
            )

    def test_scalar_subquery_inside_lateral_join(self):
        with self.temp_view("t1", "t2"):
            t1 = self.table1()
            t2 = self.table2()

            assertDataFrameEqual(
                t1.lateralJoin(
                    self.spark.range(1).select(
                        sf.col("c2").outer(), t2.select(sf.min(sf.col("c2"))).scalar()
                    )
                ),
                self.spark.sql(
                    """SELECT * FROM t1, LATERAL (SELECT c2, (SELECT MIN(c2) FROM t2))"""
                ),
            )
            assertDataFrameEqual(
                t1.lateralJoin(
                    self.spark.range(1)
                    .select(sf.col("c1").outer().alias("a"))
                    .select(
                        t2.where(sf.col("c1") == sf.col("a").outer())
                        .select(sf.sum(sf.col("c2")))
                        .scalar()
                    )
                ),
                self.spark.sql("""
                    SELECT * FROM t1, LATERAL (
                        SELECT (SELECT SUM(c2) FROM t2 WHERE c1 = a) FROM (SELECT c1 AS a)
                    )
                    """),
            )

    def test_lateral_join_inside_subquery(self):
        with self.temp_view("t1", "t2"):
            t1 = self.table1()
            t2 = self.table2()

            assertDataFrameEqual(
                t1.where(
                    sf.col("c1")
                    == (
                        t2.lateralJoin(self.spark.range(1).select(sf.col("c1").outer().alias("a")))
                        .select(sf.min(sf.col("a")))
                        .scalar()
                    )
                ),
                self.spark.sql("""
                    SELECT * FROM t1 WHERE c1 = (SELECT MIN(a) FROM t2, LATERAL (SELECT c1 AS a))
                    """),
            )
            assertDataFrameEqual(
                t1.where(
                    sf.col("c1")
                    == (
                        t2.lateralJoin(self.spark.range(1).select(sf.col("c1").outer().alias("a")))
                        .where(sf.col("c1") == sf.col("t1.c1").outer())
                        .select(sf.min(sf.col("a")))
                        .scalar()
                    )
                ),
                self.spark.sql("""
                    SELECT * FROM t1
                    WHERE c1 = (SELECT MIN(a) FROM t2, LATERAL (SELECT c1 AS a) WHERE c1 = t1.c1)
                    """),
            )

    def test_lateral_join_with_table_valued_functions(self):
        with self.temp_view("t1", "t3"):
            t1 = self.table1()
            t3 = self.table3()

            assertDataFrameEqual(
                t1.lateralJoin(self.spark.tvf.range(3)),
                self.spark.sql("""SELECT * FROM t1, LATERAL RANGE(3)"""),
            )
            assertDataFrameEqual(
                t1.lateralJoin(
                    self.spark.tvf.explode(sf.array(sf.col("c1").outer(), sf.col("c2").outer()))
                ).toDF("c1", "c2", "c3"),
                self.spark.sql("""SELECT * FROM t1, LATERAL EXPLODE(ARRAY(c1, c2)) t2(c3)"""),
            )
            assertDataFrameEqual(
                t3.lateralJoin(self.spark.tvf.explode_outer(sf.col("c2").outer())).toDF(
                    "c1", "c2", "v"
                ),
                self.spark.sql("""SELECT * FROM t3, LATERAL EXPLODE_OUTER(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() + 1).alias("v"))),
                self.spark.sql(
                    """SELECT * FROM EXPLODE(ARRAY(1, 2)) t(v), LATERAL (SELECT v + 1 AS v)"""
                ),
            )

    def test_lateral_join_with_table_valued_functions_and_join_conditions(self):
        with self.temp_view("t1", "t3"):
            t1 = self.table1()
            t3 = self.table3()

            assertDataFrameEqual(
                t1.lateralJoin(
                    self.spark.tvf.explode(sf.array(sf.col("c1").outer(), sf.col("c2").outer())),
                    sf.col("c1") == sf.col("col"),
                ).toDF("c1", "c2", "c3"),
                self.spark.sql(
                    """SELECT * FROM t1 JOIN LATERAL EXPLODE(ARRAY(c1, c2)) t(c3) ON t1.c1 = c3"""
                ),
            )
            assertDataFrameEqual(
                t3.lateralJoin(
                    self.spark.tvf.explode(sf.col("c2").outer()),
                    sf.col("c1") == sf.col("col"),
                ).toDF("c1", "c2", "c3"),
                self.spark.sql("""SELECT * FROM t3 JOIN LATERAL EXPLODE(c2) t(c3) ON t3.c1 = c3"""),
            )
            assertDataFrameEqual(
                t3.lateralJoin(
                    self.spark.tvf.explode(sf.col("c2").outer()),
                    sf.col("c1") == sf.col("col"),
                    "left",
                ).toDF("c1", "c2", "c3"),
                self.spark.sql(
                    """SELECT * FROM t3 LEFT JOIN LATERAL EXPLODE(c2) t(c3) ON t3.c1 = c3"""
                ),
            )

    def test_subquery_with_generator_and_tvf(self):
        with self.temp_view("t1"):
            t1 = self.table1()

            assertDataFrameEqual(
                self.spark.range(1).select(sf.explode(t1.select(sf.collect_list("c2")).scalar())),
                self.spark.sql("""SELECT EXPLODE((SELECT COLLECT_LIST(c2) FROM t1))"""),
            )
            assertDataFrameEqual(
                self.spark.tvf.explode(t1.select(sf.collect_list("c2")).scalar()),
                self.spark.sql("""SELECT * FROM EXPLODE((SELECT COLLECT_LIST(c2) FROM t1))"""),
            )

    def test_subquery_in_join_condition(self):
        with self.temp_view("t1", "t2"):
            t1 = self.table1()
            t2 = self.table2()

            assertDataFrameEqual(
                t1.join(t2, sf.col("t1.c1") == t1.select(sf.max("c1")).scalar()),
                self.spark.sql("""SELECT * FROM t1 JOIN t2 ON t1.c1 = (SELECT MAX(c1) FROM t1)"""),
            )

    def test_subquery_in_unpivot(self):
        self.check_subquery_in_unpivot(QueryContextType.DataFrame, "exists")

    def check_subquery_in_unpivot(self, query_context_type, fragment):
        with self.temp_view("t1", "t2"):
            t1 = self.table1()
            t2 = self.table2()

            with self.assertRaises(AnalysisException) as pe:
                t1.unpivot("c1", t2.exists(), "c1", "c2").collect()

            self.check_error(
                exception=pe.exception,
                errorClass=(
                    "UNSUPPORTED_SUBQUERY_EXPRESSION_CATEGORY.UNSUPPORTED_IN_EXISTS_SUBQUERY"
                ),
                messageParameters={"treeNode": "Expand.*"},
                query_context_type=query_context_type,
                fragment=fragment,
                matchPVals=True,
            )

    def test_subquery_in_transpose(self):
        with self.temp_view("t1"):
            t1 = self.table1()

            with self.assertRaises(AnalysisException) as pe:
                t1.transpose(t1.select(sf.max("c1")).scalar()).collect()

            self.check_error(
                exception=pe.exception,
                errorClass="TRANSPOSE_INVALID_INDEX_COLUMN",
                messageParameters={"reason": "Index column must be an atomic attribute"},
            )

    def test_subquery_in_with_columns(self):
        with self.temp_view("t1"):
            t1 = self.table1()

            assertDataFrameEqual(
                t1.withColumn(
                    "scalar",
                    self.spark.range(1)
                    .select(sf.col("c1").outer() + sf.col("c2").outer())
                    .scalar(),
                ),
                t1.select("*", (sf.col("c1") + sf.col("c2")).alias("scalar")),
            )
            assertDataFrameEqual(
                t1.withColumn(
                    "scalar",
                    self.spark.range(1)
                    .withColumn("c1", sf.col("c1").outer())
                    .select(sf.col("c1") + sf.col("c2").outer())
                    .scalar(),
                ),
                t1.select("*", (sf.col("c1") + sf.col("c2")).alias("scalar")),
            )
            assertDataFrameEqual(
                t1.withColumn(
                    "scalar",
                    self.spark.range(1)
                    .select(sf.col("c1").outer().alias("c1"))
                    .withColumn("c2", sf.col("c2").outer())
                    .select(sf.col("c1") + sf.col("c2"))
                    .scalar(),
                ),
                t1.select("*", (sf.col("c1") + sf.col("c2")).alias("scalar")),
            )

    def test_subquery_in_with_columns_renamed(self):
        with self.temp_view("t1"):
            t1 = self.table1()

            assertDataFrameEqual(
                t1.withColumn(
                    "scalar",
                    self.spark.range(1)
                    .select(sf.col("c1").outer().alias("c1"), sf.col("c2").outer().alias("c2"))
                    .withColumnsRenamed({"c1": "x", "c2": "y"})
                    .select(sf.col("x") + sf.col("y"))
                    .scalar(),
                ),
                t1.select("*", (sf.col("c1").alias("x") + sf.col("c2").alias("y")).alias("scalar")),
            )

    def test_subquery_in_drop(self):
        with self.temp_view("t1"):
            t1 = self.table1()

            assertDataFrameEqual(t1.drop(self.spark.range(1).select(sf.lit("c1")).scalar()), t1)

    def test_subquery_in_repartition(self):
        with self.temp_view("t1"):
            t1 = self.table1()

            assertDataFrameEqual(t1.repartition(self.spark.range(1).select(sf.lit(1)).scalar()), t1)


class SubqueryTests(SubqueryTestsMixin, ReusedSQLTestCase):
    pass


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

    main()