#
# 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 glob
import os
import pydoc
import shutil
import tempfile
import warnings
import unittest
import io
from contextlib import redirect_stdout

from pyspark.sql import Row, functions, DataFrame
from pyspark.sql.functions import (
    col,
    lit,
    count,
    struct,
    date_format,
    to_date,
    array,
    explode,
)
from pyspark.sql.types import (
    StringType,
    IntegerType,
    LongType,
    StructType,
    StructField,
)
from pyspark.storagelevel import StorageLevel
from pyspark.errors import (
    AnalysisException,
    IllegalArgumentException,
    PySparkTypeError,
    PySparkValueError,
)
from pyspark.testing import assertDataFrameEqual
from pyspark.testing.sqlutils import (
    ReusedSQLTestCase,
    SPARK_HOME,
)
from pyspark.testing.utils import (
    have_pyarrow,
    have_pandas,
    pandas_requirement_message,
    pyarrow_requirement_message,
)


class DataFrameTestsMixin:
    def test_range(self):
        self.assertEqual(self.spark.range(1, 1).count(), 0)
        self.assertEqual(self.spark.range(1, 0, -1).count(), 1)
        self.assertEqual(self.spark.range(0, 1 << 40, 1 << 39).count(), 2)
        self.assertEqual(self.spark.range(-2).count(), 0)
        self.assertEqual(self.spark.range(3).count(), 3)

    def test_table(self):
        with self.assertRaises(PySparkTypeError) as pe:
            self.spark.table(None)

        self.check_error(
            exception=pe.exception,
            errorClass="NOT_EXPECTED_TYPE",
            messageParameters={
                "expected_type": "str",
                "arg_name": "tableName",
                "arg_type": "NoneType",
            },
        )

    def test_dataframe_star(self):
        df1 = self.spark.createDataFrame([{"a": 1}])
        df2 = self.spark.createDataFrame([{"a": 1, "b": "v"}])
        df3 = df2.withColumnsRenamed({"a": "x", "b": "y"})

        df = df1.join(df2)
        self.assertEqual(df.columns, ["a", "a", "b"])
        self.assertEqual(df.select(df1["*"]).columns, ["a"])
        self.assertEqual(df.select(df2["*"]).columns, ["a", "b"])

        df = df1.join(df2).withColumn("c", lit(0))
        self.assertEqual(df.columns, ["a", "a", "b", "c"])
        self.assertEqual(df.select(df1["*"]).columns, ["a"])
        self.assertEqual(df.select(df2["*"]).columns, ["a", "b"])

        df = df1.join(df2, "a")
        self.assertEqual(df.columns, ["a", "b"])
        self.assertEqual(df.select(df1["*"]).columns, ["a"])
        self.assertEqual(df.select(df2["*"]).columns, ["a", "b"])

        df = df1.join(df2, "a").withColumn("c", lit(0))
        self.assertEqual(df.columns, ["a", "b", "c"])
        self.assertEqual(df.select(df1["*"]).columns, ["a"])
        self.assertEqual(df.select(df2["*"]).columns, ["a", "b"])

        df = df2.join(df3)
        self.assertEqual(df.columns, ["a", "b", "x", "y"])
        self.assertEqual(df.select(df2["*"]).columns, ["a", "b"])
        self.assertEqual(df.select(df3["*"]).columns, ["x", "y"])

        df = df2.join(df3).withColumn("c", lit(0))
        self.assertEqual(df.columns, ["a", "b", "x", "y", "c"])
        self.assertEqual(df.select(df2["*"]).columns, ["a", "b"])
        self.assertEqual(df.select(df3["*"]).columns, ["x", "y"])

    def test_count_star(self):
        df1 = self.spark.createDataFrame([{"a": 1}])
        df2 = self.spark.createDataFrame([{"a": 1, "b": "v"}])
        df3 = df2.select(struct("a", "b").alias("s"))

        self.assertEqual(df1.select(count(df1["*"])).columns, ["count(1)"])
        self.assertEqual(df1.select(count(col("*"))).columns, ["count(1)"])

        self.assertEqual(df2.select(count(df2["*"])).columns, ["count(1)"])
        self.assertEqual(df2.select(count(col("*"))).columns, ["count(1)"])

        self.assertEqual(df3.select(count(df3["*"])).columns, ["count(1)"])
        self.assertEqual(df3.select(count(col("*"))).columns, ["count(1)"])

    def test_lateral_column_alias(self):
        df1 = self.spark.range(10).select(
            (col("id") + lit(1)).alias("x"), (col("x") + lit(1)).alias("y")
        )
        df2 = self.spark.range(10).select(col("id").alias("x"))
        df3 = df1.join(df2, df1.x == df2.x).select(df1.y)
        self.assertEqual(df3.columns, ["y"])
        self.assertEqual(df3.count(), 9)

    def test_duplicated_column_names(self):
        df = self.spark.createDataFrame([(1, 2)], ["c", "c"])
        row = df.select("*").first()
        self.assertEqual(1, row[0])
        self.assertEqual(2, row[1])
        self.assertEqual("Row(c=1, c=2)", str(row))
        # Cannot access columns
        self.assertRaises(AnalysisException, lambda: df.select(df[0]).first())
        self.assertRaises(AnalysisException, lambda: df.select(df.c).first())
        self.assertRaises(AnalysisException, lambda: df.select(df["c"]).first())

    def test_help_command(self):
        # Regression test for SPARK-5464
        rdd = self.sc.parallelize(['{"foo":"bar"}', '{"foo":"baz"}'])
        df = self.spark.read.json(rdd)
        # render_doc() reproduces the help() exception without printing output
        self.check_help_command(df)

    def check_help_command(self, df):
        pydoc.render_doc(df)
        pydoc.render_doc(df.foo)
        pydoc.render_doc(df.take(1))

    def test_drop(self):
        df = self.spark.createDataFrame([("A", 50, "Y"), ("B", 60, "Y")], ["name", "age", "active"])
        self.assertEqual(df.drop("active").columns, ["name", "age"])
        self.assertEqual(df.drop("active", "nonexistent_column").columns, ["name", "age"])
        self.assertEqual(df.drop("name", "age", "active").columns, [])
        self.assertEqual(df.drop(col("name")).columns, ["age", "active"])
        self.assertEqual(df.drop(col("name"), col("age")).columns, ["active"])
        self.assertEqual(df.drop(col("name"), col("age"), col("random")).columns, ["active"])

    def test_drop_col_from_different_dataframe(self):
        df1 = self.spark.range(10)
        df2 = df1.withColumn("v0", lit(0))

        # drop df2["id"] from df2
        self.assertEqual(df2.drop(df2["id"]).columns, ["v0"])

        # drop df1["id"] from df2, which is semantically equal to df2["id"]
        # note that df1.drop(df2["id"]) works in Classic, but not in Connect
        self.assertEqual(df2.drop(df1["id"]).columns, ["v0"])

        df3 = df2.select("*", lit(1).alias("v1"))

        # drop df3["id"] from df3
        self.assertEqual(df3.drop(df3["id"]).columns, ["v0", "v1"])

        # drop df2["id"] from df3, which is semantically equal to df3["id"]
        self.assertEqual(df3.drop(df2["id"]).columns, ["v0", "v1"])

        # drop df1["id"] from df3, which is semantically equal to df3["id"]
        self.assertEqual(df3.drop(df1["id"]).columns, ["v0", "v1"])

        # drop df3["v0"] from df3
        self.assertEqual(df3.drop(df3["v0"]).columns, ["id", "v1"])

        # drop df2["v0"] from df3, which is semantically equal to df3["v0"]
        self.assertEqual(df3.drop(df2["v0"]).columns, ["id", "v1"])

    def test_drop_join(self):
        left_df = self.spark.createDataFrame(
            [(1, "a"), (2, "b"), (3, "c")],
            ["join_key", "value1"],
        )
        right_df = self.spark.createDataFrame(
            [(1, "aa"), (2, "bb"), (4, "dd")],
            ["join_key", "value2"],
        )
        joined_df = left_df.join(
            right_df,
            on=left_df["join_key"] == right_df["join_key"],
            how="left",
        )

        dropped_1 = joined_df.drop(left_df["join_key"])
        self.assertEqual(dropped_1.columns, ["value1", "join_key", "value2"])
        self.assertEqual(
            dropped_1.sort("value1").collect(),
            [
                Row(value1="a", join_key=1, value2="aa"),
                Row(value1="b", join_key=2, value2="bb"),
                Row(value1="c", join_key=None, value2=None),
            ],
        )

        dropped_2 = joined_df.drop(right_df["join_key"])
        self.assertEqual(dropped_2.columns, ["join_key", "value1", "value2"])
        self.assertEqual(
            dropped_2.sort("value1").collect(),
            [
                Row(join_key=1, value1="a", value2="aa"),
                Row(join_key=2, value1="b", value2="bb"),
                Row(join_key=3, value1="c", value2=None),
            ],
        )

    def test_with_columns_renamed(self):
        df = self.spark.createDataFrame([("Alice", 50), ("Alice", 60)], ["name", "age"])

        # rename both columns
        renamed_df1 = df.withColumnsRenamed({"name": "naam", "age": "leeftijd"})
        self.assertEqual(renamed_df1.columns, ["naam", "leeftijd"])

        # rename one column with one missing name
        renamed_df2 = df.withColumnsRenamed({"name": "naam", "address": "adres"})
        self.assertEqual(renamed_df2.columns, ["naam", "age"])

        # negative test for incorrect type
        with self.assertRaises(PySparkTypeError) as pe:
            df.withColumnsRenamed(("name", "x"))

        self.check_error(
            exception=pe.exception,
            errorClass="NOT_EXPECTED_TYPE",
            messageParameters={"expected_type": "dict", "arg_name": "colsMap", "arg_type": "tuple"},
        )

    def test_with_columns_renamed_with_duplicated_names(self):
        df1 = self.spark.createDataFrame([(1, "v1")], ["id", "value"])
        df2 = self.spark.createDataFrame([(1, "x", "v2")], ["id", "a", "value"])
        join = df2.join(df1, on=["id"], how="left")

        self.assertEqual(
            join.withColumnRenamed("id", "value").columns,
            join.withColumnsRenamed({"id": "value"}).columns,
        )
        self.assertEqual(
            join.withColumnRenamed("a", "b").columns,
            join.withColumnsRenamed({"a": "b"}).columns,
        )
        self.assertEqual(
            join.withColumnRenamed("value", "new_value").columns,
            join.withColumnsRenamed({"value": "new_value"}).columns,
        )
        self.assertEqual(
            join.withColumnRenamed("x", "y").columns,
            join.withColumnsRenamed({"x": "y"}).columns,
        )

    def test_ordering_of_with_columns_renamed(self):
        df = self.spark.range(10)

        df1 = df.withColumnsRenamed({"id": "a", "a": "b"})
        self.assertEqual(df1.columns, ["b"])

        df2 = df.withColumnsRenamed({"a": "b", "id": "a"})
        self.assertEqual(df2.columns, ["a"])

    def test_drop_duplicates(self):
        # SPARK-36034 test that drop duplicates throws a type error when in correct type provided
        df = self.spark.createDataFrame([("Alice", 50), ("Alice", 60)], ["name", "age"])

        # shouldn't drop a non-null row
        self.assertEqual(df.dropDuplicates().count(), 2)

        self.assertEqual(df.dropDuplicates(["name"]).count(), 1)

        self.assertEqual(df.dropDuplicates(["name", "age"]).count(), 2)

        with self.assertRaises(PySparkTypeError) as pe:
            df.dropDuplicates("name")

        self.check_error(
            exception=pe.exception,
            errorClass="NOT_EXPECTED_TYPE",
            messageParameters={
                "expected_type": "list or tuple",
                "arg_name": "subset",
                "arg_type": "str",
            },
        )

        # Should raise proper error when taking non-string values
        with self.assertRaises(PySparkTypeError) as pe:
            df.dropDuplicates([None]).show()

        self.check_error(
            exception=pe.exception,
            errorClass="NOT_EXPECTED_TYPE",
            messageParameters={
                "expected_type": "str",
                "arg_name": "subset",
                "arg_type": "NoneType",
            },
        )

        with self.assertRaises(PySparkTypeError) as pe:
            df.dropDuplicates([1]).show()

        self.check_error(
            exception=pe.exception,
            errorClass="NOT_EXPECTED_TYPE",
            messageParameters={"expected_type": "str", "arg_name": "subset", "arg_type": "int"},
        )

    def test_drop_duplicates_with_ambiguous_reference(self):
        df1 = self.spark.createDataFrame([(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"])
        df2 = self.spark.createDataFrame([Row(height=80, name="Tom"), Row(height=85, name="Bob")])
        df3 = df1.join(df2, df1.name == df2.name, "inner")

        self.assertEqual(df3.drop("name", "age").columns, ["height"])
        self.assertEqual(df3.drop("name", df3.age, "unknown").columns, ["height"])
        self.assertEqual(df3.drop("name", "age", df3.height).columns, [])

    def test_drop_empty_column(self):
        df = self.spark.createDataFrame([(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"])

        self.assertEqual(df.drop().columns, ["age", "name"])
        self.assertEqual(df.drop(*[]).columns, ["age", "name"])

    def test_drop_column_name_with_dot(self):
        df = (
            self.spark.range(1, 3)
            .withColumn("first.name", lit("Peter"))
            .withColumn("city.name", lit("raleigh"))
            .withColumn("state", lit("nc"))
        )

        self.assertEqual(df.drop("first.name").columns, ["id", "city.name", "state"])
        self.assertEqual(df.drop("city.name").columns, ["id", "first.name", "state"])
        self.assertEqual(df.drop("first.name", "city.name").columns, ["id", "state"])
        self.assertEqual(
            df.drop("first.name", "city.name", "unknown.unknown").columns, ["id", "state"]
        )
        self.assertEqual(
            df.drop("unknown.unknown").columns, ["id", "first.name", "city.name", "state"]
        )

    def test_with_column_with_existing_name(self):
        keys = self.df.withColumn("key", self.df.key).select("key").collect()
        self.assertEqual([r.key for r in keys], list(range(100)))

    # regression test for SPARK-10417
    def test_column_iterator(self):
        def foo():
            for x in self.df.key:
                break

        self.assertRaises(TypeError, foo)

    def test_with_columns(self):
        # With single column
        keys = self.df.withColumns({"key": self.df.key}).select("key").collect()
        self.assertEqual([r.key for r in keys], list(range(100)))

        # With key and value columns
        kvs = (
            self.df.withColumns({"key": self.df.key, "value": self.df.value})
            .select("key", "value")
            .collect()
        )
        self.assertEqual([(r.key, r.value) for r in kvs], [(i, str(i)) for i in range(100)])

        # Columns rename
        kvs = (
            self.df.withColumns({"key_alias": self.df.key, "value_alias": self.df.value})
            .select("key_alias", "value_alias")
            .collect()
        )
        self.assertEqual(
            [(r.key_alias, r.value_alias) for r in kvs], [(i, str(i)) for i in range(100)]
        )

        # Type check
        self.assertRaises(TypeError, self.df.withColumns, ["key"])
        self.assertRaises(Exception, self.df.withColumns)

    def test_generic_hints(self):
        df1 = self.spark.range(10e10).toDF("id")
        df2 = self.spark.range(10e10).toDF("id")

        self.assertIsInstance(df1.hint("broadcast"), type(df1))

        # Dummy rules
        self.assertIsInstance(df1.hint("broadcast", "foo", "bar"), type(df1))

        with io.StringIO() as buf, redirect_stdout(buf):
            df1.join(df2.hint("broadcast"), "id").explain(True)
            self.assertEqual(1, buf.getvalue().count("BroadcastHashJoin"))

    def test_coalesce_hints_with_string_parameter(self):
        with self.sql_conf({"spark.sql.adaptive.coalescePartitions.enabled": False}):
            df = self.spark.createDataFrame(
                zip(["A", "B"] * 2**9, range(2**10)),
                StructType([StructField("a", StringType()), StructField("n", IntegerType())]),
            )
            with io.StringIO() as buf, redirect_stdout(buf):
                # COALESCE
                coalesce = df.hint("coalesce", 2)
                coalesce.explain(True)
                output = buf.getvalue()
                self.assertGreaterEqual(output.count("Coalesce 2"), 1)
                buf.truncate(0)
                buf.seek(0)

                # REPARTITION_BY_RANGE
                range_partitioned = df.hint("REPARTITION_BY_RANGE", 2, "a")
                range_partitioned.explain(True)
                output = buf.getvalue()
                self.assertGreaterEqual(output.count("REPARTITION_BY_NUM"), 1)
                buf.truncate(0)
                buf.seek(0)

                # REBALANCE
                rebalanced1 = df.hint("REBALANCE", "a")  # just check this doesn't error
                rebalanced1.explain(True)
                rebalanced2 = df.hint("REBALANCE", 2)
                rebalanced2.explain(True)
                rebalanced3 = df.hint("REBALANCE", 2, "a")
                rebalanced3.explain(True)
                rebalanced4 = df.hint("REBALANCE", functions.col("a"))
                rebalanced4.explain(True)
                output = buf.getvalue()
                self.assertGreaterEqual(output.count("REBALANCE_PARTITIONS_BY_NONE"), 1)
                self.assertGreaterEqual(output.count("REBALANCE_PARTITIONS_BY_COL"), 3)

    # add tests for SPARK-23647 (test more types for hint)
    def test_extended_hint_types(self):
        df = self.spark.range(10e10).toDF("id")
        such_a_nice_list = ["itworks1", "itworks2", "itworks3"]
        int_list = [1, 2, 3]
        hinted_df = df.hint("my awesome hint", 1.2345, "what", such_a_nice_list, int_list)

        self.assertIsInstance(df.hint("broadcast", []), type(df))
        self.assertIsInstance(df.hint("broadcast", ["foo", "bar"]), type(df))

        with io.StringIO() as buf, redirect_stdout(buf):
            hinted_df.explain(True)
            explain_output = buf.getvalue()
            self.assertGreaterEqual(explain_output.count("1.2345"), 1)
            self.assertGreaterEqual(explain_output.count("what"), 1)
            self.assertGreaterEqual(explain_output.count("itworks"), 1)

    def test_sample(self):
        with self.assertRaises(PySparkTypeError) as pe:
            self.spark.range(1).sample()

        self.check_error(
            exception=pe.exception,
            errorClass="NOT_EXPECTED_TYPE",
            messageParameters={
                "expected_type": "bool, float or int",
                "arg_name": "withReplacement (optional), fraction (required) and seed (optional)",
                "arg_type": "NoneType, NoneType, NoneType",
            },
        )

        self.assertRaises(TypeError, lambda: self.spark.range(1).sample("a"))

        self.assertRaises(TypeError, lambda: self.spark.range(1).sample(seed="abc"))

        self.assertRaises(
            IllegalArgumentException, lambda: self.spark.range(1).sample(-1.0).count()
        )

    def test_sample_with_random_seed(self):
        df = self.spark.range(10000).sample(0.1)
        cnts = [df.count() for i in range(10)]
        self.assertEqual(1, len(set(cnts)))

    def test_toDF_with_string(self):
        df = self.spark.createDataFrame([("John", 30), ("Alice", 25), ("Bob", 28)])
        data = [("John", 30), ("Alice", 25), ("Bob", 28)]

        result = df.toDF("key", "value")
        self.assertEqual(result.schema.simpleString(), "struct<key:string,value:bigint>")
        self.assertEqual(result.collect(), data)

        with self.assertRaises(PySparkTypeError) as pe:
            df.toDF("key", None)

        self.check_error(
            exception=pe.exception,
            errorClass="NOT_EXPECTED_TYPE",
            messageParameters={
                "expected_type": "list[str]",
                "arg_name": "cols",
                "arg_type": "NoneType",
            },
        )

    def test_toDF_with_schema_string(self):
        data = [Row(key=i, value=str(i)) for i in range(100)]
        rdd = self.sc.parallelize(data, 5)

        df = rdd.toDF("key: int, value: string")
        self.assertEqual(df.schema.simpleString(), "struct<key:int,value:string>")
        self.assertEqual(df.collect(), data)

        # different but compatible field types can be used.
        df = rdd.toDF("key: string, value: string")
        self.assertEqual(df.schema.simpleString(), "struct<key:string,value:string>")
        self.assertEqual(df.collect(), [Row(key=str(i), value=str(i)) for i in range(100)])

        # field names can differ.
        df = rdd.toDF(" a: int, b: string ")
        self.assertEqual(df.schema.simpleString(), "struct<a:int,b:string>")
        self.assertEqual(df.collect(), data)

        # number of fields must match.
        self.assertRaisesRegex(
            Exception,
            "FIELD_STRUCT_LENGTH_MISMATCH",
            lambda: rdd.coalesce(1).toDF("key: int").collect(),
        )

        # field types mismatch will cause exception at runtime.
        self.assertRaisesRegex(
            Exception,
            "FIELD_DATA_TYPE_UNACCEPTABLE",
            lambda: rdd.coalesce(1).toDF("key: float, value: string").collect(),
        )

        # flat schema values will be wrapped into row.
        df = rdd.map(lambda row: row.key).toDF("int")
        self.assertEqual(df.schema.simpleString(), "struct<value:int>")
        self.assertEqual(df.collect(), [Row(key=i) for i in range(100)])

        # users can use DataType directly instead of data type string.
        df = rdd.map(lambda row: row.key).toDF(IntegerType())
        self.assertEqual(df.schema.simpleString(), "struct<value:int>")
        self.assertEqual(df.collect(), [Row(key=i) for i in range(100)])

    def test_create_df_with_collation(self):
        schema = StructType([StructField("name", StringType("UNICODE_CI"), True)])
        df = self.spark.createDataFrame([("Alice",), ("alice",)], schema)

        self.assertEqual(df.select("name").distinct().count(), 1)

    def test_print_schema(self):
        df = self.spark.createDataFrame([(1, (2, 2))], ["a", "b"])

        with io.StringIO() as buf, redirect_stdout(buf):
            df.printSchema(1)
            self.assertEqual(1, buf.getvalue().count("long"))
            self.assertEqual(0, buf.getvalue().count("_1"))
            self.assertEqual(0, buf.getvalue().count("_2"))

            buf.truncate(0)
            buf.seek(0)

            df.printSchema(2)
            self.assertEqual(3, buf.getvalue().count("long"))
            self.assertEqual(1, buf.getvalue().count("_1"))
            self.assertEqual(1, buf.getvalue().count("_2"))

    def test_join_without_on(self):
        df1 = self.spark.range(1).toDF("a")
        df2 = self.spark.range(1).toDF("b")

        with self.sql_conf({"spark.sql.crossJoin.enabled": False}):
            self.assertRaises(AnalysisException, lambda: df1.join(df2, how="inner").collect())

        with self.sql_conf({"spark.sql.crossJoin.enabled": True}):
            actual = df1.join(df2, how="inner").collect()
            expected = [Row(a=0, b=0)]
            self.assertEqual(actual, expected)

    # Regression test for invalid join methods when on is None, Spark-14761
    def test_invalid_join_method(self):
        df1 = self.spark.createDataFrame([("Alice", 5), ("Bob", 8)], ["name", "age"])
        df2 = self.spark.createDataFrame([("Alice", 80), ("Bob", 90)], ["name", "height"])
        self.assertRaises(AnalysisException, lambda: df1.join(df2, how="invalid-join-type"))

    # Cartesian products require cross join syntax
    def test_require_cross(self):
        df1 = self.spark.createDataFrame([(1, "1")], ("key", "value"))
        df2 = self.spark.createDataFrame([(1, "1")], ("key", "value"))

        with self.sql_conf({"spark.sql.crossJoin.enabled": False}):
            # joins without conditions require cross join syntax
            self.assertRaises(AnalysisException, lambda: df1.join(df2).collect())

            # works with crossJoin
            self.assertEqual(1, df1.crossJoin(df2).count())

    def test_cache_dataframe(self):
        df = self.spark.createDataFrame([(2, 2), (3, 3)])
        try:
            self.assertEqual(df.storageLevel, StorageLevel.NONE)

            df.cache()
            self.assertEqual(df.storageLevel, StorageLevel.MEMORY_AND_DISK_DESER)

            df.unpersist()
            self.assertEqual(df.storageLevel, StorageLevel.NONE)

            df.persist()
            self.assertEqual(df.storageLevel, StorageLevel.MEMORY_AND_DISK_DESER)

            df.unpersist(blocking=True)
            self.assertEqual(df.storageLevel, StorageLevel.NONE)

            df.persist(StorageLevel.DISK_ONLY)
            self.assertEqual(df.storageLevel, StorageLevel.DISK_ONLY)
        finally:
            df.unpersist()
            self.assertEqual(df.storageLevel, StorageLevel.NONE)

    def test_cache_table(self):
        spark = self.spark
        tables = ["tab1", "tab2", "tab3"]
        with self.temp_view(*tables):
            for i, tab in enumerate(tables):
                spark.createDataFrame([(2, i), (3, i)]).createOrReplaceTempView(tab)
                self.assertFalse(spark.catalog.isCached(tab))
            spark.catalog.cacheTable("tab1")
            spark.catalog.cacheTable("tab3", StorageLevel.OFF_HEAP)
            self.assertTrue(spark.catalog.isCached("tab1"))
            self.assertFalse(spark.catalog.isCached("tab2"))
            self.assertTrue(spark.catalog.isCached("tab3"))
            spark.catalog.cacheTable("tab2")
            spark.catalog.uncacheTable("tab1")
            spark.catalog.uncacheTable("tab3")
            self.assertFalse(spark.catalog.isCached("tab1"))
            self.assertTrue(spark.catalog.isCached("tab2"))
            self.assertFalse(spark.catalog.isCached("tab3"))
            spark.catalog.clearCache()
            self.assertFalse(spark.catalog.isCached("tab1"))
            self.assertFalse(spark.catalog.isCached("tab2"))
            self.assertFalse(spark.catalog.isCached("tab3"))
            self.assertRaisesRegex(
                AnalysisException,
                "does_not_exist",
                lambda: spark.catalog.isCached("does_not_exist"),
            )
            self.assertRaisesRegex(
                AnalysisException,
                "does_not_exist",
                lambda: spark.catalog.cacheTable("does_not_exist"),
            )
            self.assertRaisesRegex(
                AnalysisException,
                "does_not_exist",
                lambda: spark.catalog.uncacheTable("does_not_exist"),
            )

    def test_repr_behaviors(self):
        import re

        pattern = re.compile(r"^ *\|", re.MULTILINE)
        df = self.spark.createDataFrame([(1, "1"), (22222, "22222")], ("key", "value"))

        # test when eager evaluation is enabled and _repr_html_ will not be called
        with self.sql_conf({"spark.sql.repl.eagerEval.enabled": True}):
            expected1 = """+-----+-----+
                ||  key|value|
                |+-----+-----+
                ||    1|    1|
                ||22222|22222|
                |+-----+-----+
                |"""
            self.assertEqual(re.sub(pattern, "", expected1), df.__repr__())
            with self.sql_conf({"spark.sql.repl.eagerEval.truncate": 3}):
                expected2 = """+---+-----+
                ||key|value|
                |+---+-----+
                ||  1|    1|
                ||222|  222|
                |+---+-----+
                |"""
                self.assertEqual(re.sub(pattern, "", expected2), df.__repr__())
                with self.sql_conf({"spark.sql.repl.eagerEval.maxNumRows": 1}):
                    expected3 = """+---+-----+
                    ||key|value|
                    |+---+-----+
                    ||  1|    1|
                    |+---+-----+
                    |only showing top 1 row"""
                    self.assertEqual(re.sub(pattern, "", expected3), df.__repr__())

        # test when eager evaluation is enabled and _repr_html_ will be called
        with self.sql_conf({"spark.sql.repl.eagerEval.enabled": True}):
            expected1 = """<table border='1'>
                |<tr><th>key</th><th>value</th></tr>
                |<tr><td>1</td><td>1</td></tr>
                |<tr><td>22222</td><td>22222</td></tr>
                |</table>
                |"""
            self.assertEqual(re.sub(pattern, "", expected1), df._repr_html_())
            with self.sql_conf({"spark.sql.repl.eagerEval.truncate": 3}):
                expected2 = """<table border='1'>
                    |<tr><th>key</th><th>value</th></tr>
                    |<tr><td>1</td><td>1</td></tr>
                    |<tr><td>222</td><td>222</td></tr>
                    |</table>
                    |"""
                self.assertEqual(re.sub(pattern, "", expected2), df._repr_html_())
                with self.sql_conf({"spark.sql.repl.eagerEval.maxNumRows": 1}):
                    expected3 = """<table border='1'>
                        |<tr><th>key</th><th>value</th></tr>
                        |<tr><td>1</td><td>1</td></tr>
                        |</table>
                        |only showing top 1 row
                        |"""
                    self.assertEqual(re.sub(pattern, "", expected3), df._repr_html_())

        # test when eager evaluation is disabled and _repr_html_ will be called
        with self.sql_conf({"spark.sql.repl.eagerEval.enabled": False}):
            expected = "DataFrame[key: bigint, value: string]"
            self.assertEqual(None, df._repr_html_())
            self.assertEqual(expected, df.__repr__())
            with self.sql_conf({"spark.sql.repl.eagerEval.truncate": 3}):
                self.assertEqual(None, df._repr_html_())
                self.assertEqual(expected, df.__repr__())
                with self.sql_conf({"spark.sql.repl.eagerEval.maxNumRows": 1}):
                    self.assertEqual(None, df._repr_html_())
                    self.assertEqual(expected, df.__repr__())

    def test_same_semantics_error(self):
        with self.assertRaises(PySparkTypeError) as pe:
            self.spark.range(10).sameSemantics(1)

        self.check_error(
            exception=pe.exception,
            errorClass="NOT_EXPECTED_TYPE",
            messageParameters={
                "expected_type": "DataFrame",
                "arg_name": "other",
                "arg_type": "int",
            },
        )

    def test_input_files(self):
        tpath = tempfile.mkdtemp()
        shutil.rmtree(tpath)
        try:
            self.spark.range(1, 100, 1, 10).write.parquet(tpath)
            # read parquet file and get the input files list
            input_files_list = self.spark.read.parquet(tpath).inputFiles()

            # input files list should contain 10 entries
            self.assertEqual(len(input_files_list), 10)
            # all file paths in list must contain tpath
            for file_path in input_files_list:
                self.assertTrue(tpath in file_path)
        finally:
            shutil.rmtree(tpath)

    def test_df_show(self):
        # SPARK-35408: ensure better diagnostics if incorrect parameters are passed
        # to DataFrame.show

        df = self.spark.createDataFrame([("foo",)])
        df.show(5)
        df.show(5, True)
        df.show(5, 1, True)
        df.show(n=5, truncate="1", vertical=False)
        df.show(n=5, truncate=1.5, vertical=False)

        with self.assertRaises(PySparkTypeError) as pe:
            df.show(True)

        self.check_error(
            exception=pe.exception,
            errorClass="NOT_EXPECTED_TYPE",
            messageParameters={"expected_type": "int", "arg_name": "n", "arg_type": "bool"},
        )

        with self.assertRaises(PySparkTypeError) as pe:
            df.show(vertical="foo")

        self.check_error(
            exception=pe.exception,
            errorClass="NOT_EXPECTED_TYPE",
            messageParameters={"expected_type": "bool", "arg_name": "vertical", "arg_type": "str"},
        )

        with self.assertRaises(PySparkTypeError) as pe:
            df.show(truncate="foo")

        self.check_error(
            exception=pe.exception,
            errorClass="NOT_EXPECTED_TYPE",
            messageParameters={"expected_type": "bool", "arg_name": "truncate", "arg_type": "str"},
        )

    def test_df_merge_into(self):
        filename_pattern = (
            "sql/catalyst/target/scala-*/test-classes/org/apache/spark/sql/connector/catalog/"
            "InMemoryRowLevelOperationTableCatalog.class"
        )
        if not bool(glob.glob(os.path.join(SPARK_HOME, filename_pattern))):
            raise unittest.SkipTest(
                "org.apache.spark.sql.connector.catalog.InMemoryRowLevelOperationTableCatalog' "
                "is not available. Will skip the related tests"
            )

        try:
            # InMemoryRowLevelOperationTableCatalog is a test catalog that is included in the
            # catalyst-test package. If Spark complains that it can't find this class, make sure
            # the catalyst-test JAR does exist in "SPARK_HOME/assembly/target/scala-x.xx/jars"
            # directory. If not, build it with `build/sbt test:assembly` and copy it over.
            self.spark.conf.set(
                "spark.sql.catalog.testcat",
                "org.apache.spark.sql.connector.catalog.InMemoryRowLevelOperationTableCatalog",
            )
            with self.table("testcat.ns1.target"):

                def reset_target_table():
                    self.spark.createDataFrame(
                        [(1, "Alice"), (2, "Bob")], ["id", "name"]
                    ).write.mode("overwrite").saveAsTable("testcat.ns1.target")

                source = self.spark.createDataFrame([(1, "Charlie"), (3, "David")], ["id", "name"])  # type: DataFrame

                from pyspark.sql.functions import col

                # Match -> update, NotMatch -> insert, NotMatchedBySource -> delete
                reset_target_table()
                # fmt: off
                source.mergeInto("testcat.ns1.target", source.id == col("target.id")) \
                    .whenMatched(source.id == 1).update({"name": source.name}) \
                    .whenNotMatched().insert({"id": source.id, "name": source.name}) \
                    .whenNotMatchedBySource().delete() \
                    .merge()
                # fmt: on
                self.assertEqual(
                    self.spark.table("testcat.ns1.target").orderBy("id").collect(),
                    [Row(id=1, name="Charlie"), Row(id=3, name="David")],
                )

                # Match -> updateAll, NotMatch -> insertAll, NotMatchedBySource -> update
                reset_target_table()
                # fmt: off
                source.mergeInto("testcat.ns1.target", source.id == col("target.id")) \
                    .whenMatched(source.id == 1).updateAll() \
                    .whenNotMatched(source.id == 3).insertAll() \
                    .whenNotMatchedBySource(col("target.id") == lit(2)) \
                      .update({"name": lit("not_matched")}) \
                    .merge()
                # fmt: on
                self.assertEqual(
                    self.spark.table("testcat.ns1.target").orderBy("id").collect(),
                    [
                        Row(id=1, name="Charlie"),
                        Row(id=2, name="not_matched"),
                        Row(id=3, name="David"),
                    ],
                )

                # Match -> delete, NotMatchedBySource -> delete
                reset_target_table()
                # fmt: off
                self.spark.createDataFrame([(1, "AliceJr")], ["id", "name"]) \
                    .write.mode("append").saveAsTable("testcat.ns1.target")
                source.mergeInto("testcat.ns1.target", source.id == col("target.id")) \
                    .whenMatched(col("target.name") != lit("AliceJr")).delete() \
                    .whenNotMatchedBySource().delete() \
                    .merge()
                # fmt: on
                self.assertEqual(
                    self.spark.table("testcat.ns1.target").orderBy("id").collect(),
                    [Row(id=1, name="AliceJr")],
                )
        finally:
            self.spark.conf.unset("spark.sql.catalog.testcat")

    @unittest.skipIf(
        not have_pandas or not have_pyarrow,
        pandas_requirement_message or pyarrow_requirement_message,
    )
    def test_pandas_api(self):
        import pandas as pd
        from pandas.testing import assert_frame_equal

        sdf = self.spark.createDataFrame([("a", 1), ("b", 2), ("c", 3)], ["Col1", "Col2"])
        psdf_from_sdf = sdf.pandas_api()
        psdf_from_sdf_with_index = sdf.pandas_api(index_col="Col1")
        pdf = pd.DataFrame({"Col1": ["a", "b", "c"], "Col2": [1, 2, 3]})
        pdf_with_index = pdf.set_index("Col1")

        assert_frame_equal(pdf, psdf_from_sdf.to_pandas())
        assert_frame_equal(pdf_with_index, psdf_from_sdf_with_index.to_pandas())

    def test_to(self):
        schema = StructType(
            [StructField("i", StringType(), True), StructField("j", IntegerType(), True)]
        )
        df = self.spark.createDataFrame([("a", 1)], schema)

        schema1 = StructType([StructField("j", StringType()), StructField("i", StringType())])
        df1 = df.to(schema1)
        self.assertEqual(schema1, df1.schema)
        self.assertEqual(df.count(), df1.count())

        schema2 = StructType([StructField("j", LongType())])
        df2 = df.to(schema2)
        self.assertEqual(schema2, df2.schema)
        self.assertEqual(df.count(), df2.count())

        schema3 = StructType([StructField("struct", schema1, False)])
        df3 = df.select(struct("i", "j").alias("struct")).to(schema3)
        self.assertEqual(schema3, df3.schema)
        self.assertEqual(df.count(), df3.count())

        # incompatible field nullability
        schema4 = StructType([StructField("j", LongType(), False)])
        self.assertRaisesRegex(
            AnalysisException, "NULLABLE_COLUMN_OR_FIELD", lambda: df.to(schema4).count()
        )

        # field cannot upcast
        schema5 = StructType([StructField("i", LongType())])
        self.assertRaisesRegex(
            AnalysisException, "INVALID_COLUMN_OR_FIELD_DATA_TYPE", lambda: df.to(schema5).count()
        )

    def test_colregex(self):
        with self.assertRaises(PySparkTypeError) as pe:
            self.spark.range(10).colRegex(10)

        self.check_error(
            exception=pe.exception,
            errorClass="NOT_EXPECTED_TYPE",
            messageParameters={"expected_type": "str", "arg_name": "colName", "arg_type": "int"},
        )

    def test_where(self):
        with self.assertRaises(PySparkTypeError) as pe:
            self.spark.range(10).where(10)

        self.check_error(
            exception=pe.exception,
            errorClass="NOT_EXPECTED_TYPE",
            messageParameters={
                "expected_type": "Column or str",
                "arg_name": "condition",
                "arg_type": "int",
            },
        )

    def test_duplicate_field_names(self):
        data = [
            Row(Row("a", 1), Row(2, 3, "b", 4, "c", "d")),
            Row(Row("w", 6), Row(7, 8, "x", 9, "y", "z")),
        ]
        schema = (
            StructType()
            .add("struct", StructType().add("x", StringType()).add("x", IntegerType()))
            .add(
                "struct",
                StructType()
                .add("a", IntegerType())
                .add("x", IntegerType())
                .add("x", StringType())
                .add("y", IntegerType())
                .add("y", StringType())
                .add("x", StringType()),
            )
        )
        df = self.spark.createDataFrame(data, schema=schema)

        self.assertEqual(df.schema, schema)
        self.assertEqual(df.collect(), data)

    def test_union_classmethod_usage(self):
        df = self.spark.range(1)
        self.assertEqual(DataFrame.union(df, df).collect(), [Row(id=0), Row(id=0)])

    def test_isinstance_dataframe(self):
        self.assertIsInstance(self.spark.range(1), DataFrame)

    def test_local_checkpoint_dataframe(self):
        with io.StringIO() as buf, redirect_stdout(buf):
            self.spark.range(1).localCheckpoint().explain()
            self.assertIn("ExistingRDD", buf.getvalue())

    def test_local_checkpoint_dataframe_with_storage_level(self):
        # We don't have a way to reach into the server and assert the storage level server side, but
        # this test should cover for unexpected errors in the API.
        df = self.spark.range(10).localCheckpoint(eager=True, storageLevel=StorageLevel.DISK_ONLY)
        df.collect()

    def test_socket_leak(self):
        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter("always", ResourceWarning)
            df = self.spark.range(10)
            df.collect()

            df = self.spark.range(10)
            for _ in df.toLocalIterator():
                pass

        self.assertEqual(w, [])

    def test_transpose(self):
        df = self.spark.createDataFrame([{"a": "x", "b": "y", "c": "z"}])

        # default index column
        transposed_df = df.transpose()
        expected_schema = StructType(
            [StructField("key", StringType(), False), StructField("x", StringType(), True)]
        )
        expected_data = [Row(key="b", x="y"), Row(key="c", x="z")]
        expected_df = self.spark.createDataFrame(expected_data, schema=expected_schema)
        assertDataFrameEqual(transposed_df, expected_df, checkRowOrder=True)

        # specified index column
        transposed_df = df.transpose("c")
        expected_schema = StructType(
            [StructField("key", StringType(), False), StructField("z", StringType(), True)]
        )
        expected_data = [Row(key="a", z="x"), Row(key="b", z="y")]
        expected_df = self.spark.createDataFrame(expected_data, schema=expected_schema)
        assertDataFrameEqual(transposed_df, expected_df, checkRowOrder=True)

        # enforce transpose max values
        with self.sql_conf({"spark.sql.transposeMaxValues": 0}):
            with self.assertRaises(AnalysisException) as pe:
                df.transpose().collect()
            self.check_error(
                exception=pe.exception,
                errorClass="TRANSPOSE_EXCEED_ROW_LIMIT",
                messageParameters={"maxValues": "0", "config": "spark.sql.transposeMaxValues"},
            )

        # enforce ascending order based on index column values for transposed columns
        df = self.spark.createDataFrame([{"a": "z"}, {"a": "y"}, {"a": "x"}])
        transposed_df = df.transpose()
        expected_schema = StructType(
            [
                StructField("key", StringType(), False),
                StructField("x", StringType(), True),
                StructField("y", StringType(), True),
                StructField("z", StringType(), True),
            ]
        )  # z, y, x -> x, y, z
        expected_df = self.spark.createDataFrame([], schema=expected_schema)
        assertDataFrameEqual(transposed_df, expected_df, checkRowOrder=True)

        # enforce AtomicType Attribute for index column values
        df = self.spark.createDataFrame([{"a": ["x", "x"], "b": "y", "c": "z"}])
        with self.assertRaises(AnalysisException) as pe:
            df.transpose().collect()
        self.check_error(
            exception=pe.exception,
            errorClass="TRANSPOSE_INVALID_INDEX_COLUMN",
            messageParameters={
                "reason": "Index column must be of atomic type, "
                "but found: ArrayType(StringType,true)"
            },
        )

        # enforce least common type for non-index columns
        df = self.spark.createDataFrame([{"a": "x", "b": "y", "c": 1}])
        with self.assertRaises(AnalysisException) as pe:
            df.transpose().collect()
        self.check_error(
            exception=pe.exception,
            errorClass="TRANSPOSE_NO_LEAST_COMMON_TYPE",
            messageParameters={"dt1": '"STRING"', "dt2": '"BIGINT"'},
        )

    def test_transpose_with_invalid_index_columns(self):
        # SPARK-50602: invalid index columns
        df = self.spark.createDataFrame([{"a": "x", "b": "y", "c": "z"}])

        with self.assertRaises(AnalysisException) as pe:
            df.transpose(col("a") + 1).collect()
        self.check_error(
            exception=pe.exception,
            errorClass="TRANSPOSE_INVALID_INDEX_COLUMN",
            messageParameters={"reason": "Index column must be an atomic attribute"},
        )

    def test_metadata_column(self):
        with self.sql_conf(
            {"spark.sql.catalog.testcat": "org.apache.spark.sql.connector.catalog.InMemoryCatalog"}
        ):
            tbl = "testcat.t"
            with self.table(tbl):
                self.spark.sql(f"""
                    CREATE TABLE {tbl} (index bigint, data string)
                    PARTITIONED BY (bucket(4, index), index)
                    """)
                self.spark.sql(f"""INSERT INTO {tbl} VALUES (1, 'a'), (2, 'b'), (3, 'c')""")

                df = self.spark.sql(f"""SELECT * FROM {tbl}""")
                assertDataFrameEqual(
                    df.select(df.metadataColumn("index")),
                    [Row(0), Row(0), Row(0)],
                )

    def test_with_column_and_generator(self):
        # SPARK-51451: Generators should be available with withColumn
        df = self.spark.createDataFrame([("082017",)], ["dt"]).select(
            to_date(col("dt"), "MMyyyy").alias("dt")
        )
        df_dt = df.withColumn("dt", date_format(col("dt"), "MM/dd/yyyy"))
        monthArray = [lit(x) for x in range(0, 12)]
        df_month_y = df_dt.withColumn("month_y", explode(array(monthArray)))

        assertDataFrameEqual(
            df_month_y,
            [Row(dt="08/01/2017", month_y=i) for i in range(12)],
        )

        df_dt_month_y = df.withColumns(
            {
                "dt": date_format(col("dt"), "MM/dd/yyyy"),
                "month_y": explode(array(monthArray)),
            }
        )

        assertDataFrameEqual(
            df_dt_month_y,
            [Row(dt="08/01/2017", month_y=i) for i in range(12)],
        )

    def test_query_execution_unsupported_in_classic(self):
        with self.assertRaises(PySparkValueError) as pe:
            self.spark.range(1).executionInfo

        self.check_error(
            exception=pe.exception,
            errorClass="CLASSIC_OPERATION_NOT_SUPPORTED_ON_DF",
            messageParameters={"member": "queryExecution"},
        )

    def test_to_json(self):
        df = self.spark.range(10)

        with self.sql_conf({"spark.sql.pyspark.toJSON.returnDataFrame": False}):
            from pyspark import RDD

            rdd = df.toJSON()
            self.assertIsInstance(rdd, RDD)
            self.assertEqual(rdd.count(), 10)

        with self.sql_conf({"spark.sql.pyspark.toJSON.returnDataFrame": True}):
            df = df.toJSON()
            self.assertIsInstance(df, DataFrame)
            self.assertEqual(df.select("value").count(), 10)

    def test_zip_with_index(self):
        df = self.spark.createDataFrame([("a", 1), ("b", 2), ("c", 3)], ["letter", "number"])

        # Default column name "index"
        result = df.zipWithIndex()
        self.assertEqual(result.columns, ["letter", "number", "index"])
        rows = result.collect()
        self.assertEqual(len(rows), 3)
        indices = [row["index"] for row in rows]
        self.assertEqual(sorted(indices), [0, 1, 2])

        # Custom column name
        result = df.zipWithIndex("row_id")
        self.assertEqual(result.columns, ["letter", "number", "row_id"])
        rows = result.collect()
        indices = [row["row_id"] for row in rows]
        self.assertEqual(sorted(indices), [0, 1, 2])

        # Duplicate column name causes AMBIGUOUS_REFERENCE on select
        result = df.zipWithIndex("letter")
        with self.assertRaises(AnalysisException) as ctx:
            result.select("letter").collect()
        self.assertEqual(ctx.exception.getCondition(), "AMBIGUOUS_REFERENCE")

        # Duplicate column name causes COLUMN_ALREADY_EXISTS on write
        with tempfile.TemporaryDirectory() as d:
            with self.assertRaises(AnalysisException) as ctx:
                result.write.parquet(d)
            self.assertEqual(ctx.exception.getCondition(), "COLUMN_ALREADY_EXISTS")


class DataFrameTests(DataFrameTestsMixin, ReusedSQLTestCase):
    pass


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

    main()