python/pyspark/sql/tests/connect/test_connect_basic.py - master - githubmirror/spark

python/pyspark/sql/tests/connect/test_connect_basic.py
1 671 строка63 KB
Ruifeng Zheng
[SPARK-58402][PYTHON] Consolidate UNKNOWN_EXPLAIN_MODE into VALUE_NOT_ALLOWED
28 июл 2026, 15:31
28 июл 2026, 15:312df516c
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import os
import gc
import unittest
import shutil
import tempfile
import io
from contextlib import redirect_stdout
import datetime

from pyspark.util import is_remote_only
from pyspark.errors import PySparkTypeError, PySparkValueError
from pyspark.sql.types import (
    StructType,
    StructField,
    LongType,
    StringType,
    IntegerType,
    MapType,
    ArrayType,
    Row,
)
from pyspark.testing import assertDataFrameEqual
from pyspark.testing.utils import eventually
from pyspark.testing.connectutils import (
    should_test_connect,
    connect_requirement_message,
    ReusedMixedTestCase,
)
from pyspark.testing.pandasutils import PandasOnSparkTestUtils

if should_test_connect:
    from pyspark.sql.connect.proto import ExecutePlanResponse, Expression as ProtoExpression
    from pyspark.sql.connect.column import Column
    from pyspark.sql.dataframe import DataFrame
    from pyspark.sql.connect.dataframe import DataFrame as CDataFrame
    from pyspark.sql import functions as SF
    from pyspark.sql.connect import functions as CF
    from pyspark.errors.exceptions.connect import AnalysisException, SparkConnectException


@unittest.skipIf(
    not should_test_connect or is_remote_only(),
    connect_requirement_message or "Requires JVM access",
)
class SparkConnectSQLTestCase(ReusedMixedTestCase, PandasOnSparkTestUtils):
    """Parent test fixture class for all Spark Connect related
    test cases."""

    @classmethod
    def setUpClass(cls):
        super().setUpClass()

        cls.testData = [Row(key=i, value=str(i)) for i in range(100)]
        cls.testDataStr = [Row(key=str(i)) for i in range(100)]
        cls.df = cls.spark.sparkContext.parallelize(cls.testData).toDF()
        cls.df_text = cls.spark.sparkContext.parallelize(cls.testDataStr).toDF()

        cls.tbl_name = "test_connect_basic_table_1"
        cls.tbl_name2 = "test_connect_basic_table_2"
        cls.tbl_name3 = "test_connect_basic_table_3"
        cls.tbl_name4 = "test_connect_basic_table_4"
        cls.tbl_name_empty = "test_connect_basic_table_empty"

        # Cleanup test data
        cls.spark_connect_clean_up_test_data()
        # Load test data
        cls.spark_connect_load_test_data()

    @classmethod
    def tearDownClass(cls):
        try:
            cls.spark_connect_clean_up_test_data()
        finally:
            super().tearDownClass()

    @classmethod
    def spark_connect_load_test_data(cls):
        df = cls.spark.createDataFrame([(x, f"{x}") for x in range(100)], ["id", "name"])
        # Since we might create multiple Spark sessions, we need to create global temporary view
        # that is specifically maintained in the "global_temp" schema.
        df.write.saveAsTable(cls.tbl_name)
        df2 = cls.spark.createDataFrame(
            [(x, f"{x}", 2 * x) for x in range(100)], ["col1", "col2", "col3"]
        )
        df2.write.saveAsTable(cls.tbl_name2)
        df3 = cls.spark.createDataFrame([(x, f"{x}") for x in range(100)], ["id", "test\n_column"])
        df3.write.saveAsTable(cls.tbl_name3)
        df4 = cls.spark.createDataFrame(
            [(x, {"a": x}, [x, x * 2]) for x in range(100)], ["id", "map_column", "array_column"]
        )
        df4.write.saveAsTable(cls.tbl_name4)
        empty_table_schema = StructType(
            [
                StructField("firstname", StringType(), True),
                StructField("middlename", StringType(), True),
                StructField("lastname", StringType(), True),
            ]
        )
        emptyRDD = cls.spark.sparkContext.emptyRDD()
        empty_df = cls.spark.createDataFrame(emptyRDD, empty_table_schema)
        empty_df.write.saveAsTable(cls.tbl_name_empty)

    @classmethod
    def spark_connect_clean_up_test_data(cls):
        cls.spark.sql("DROP TABLE IF EXISTS {}".format(cls.tbl_name))
        cls.spark.sql("DROP TABLE IF EXISTS {}".format(cls.tbl_name2))
        cls.spark.sql("DROP TABLE IF EXISTS {}".format(cls.tbl_name3))
        cls.spark.sql("DROP TABLE IF EXISTS {}".format(cls.tbl_name4))
        cls.spark.sql("DROP TABLE IF EXISTS {}".format(cls.tbl_name_empty))


class SparkConnectBasicTests(SparkConnectSQLTestCase):
    def test_toJSON(self):
        sdf = self.spark.range(10).withColumn("s", SF.col("id").cast("string"))
        cdf = self.connect.range(10).withColumn("s", CF.col("id").cast("string"))

        str1 = sdf.toJSON().collect()
        str2 = [row.value for row in cdf.toJSON().collect()]
        self.assertEqual(str1, str2)

    def test_serialization(self):
        from pyspark.cloudpickle import dumps, loads

        cdf = self.connect.range(10)
        data = dumps(cdf)
        cdf2 = loads(data)
        self.assertEqual(cdf.collect(), cdf2.collect())

    def test_serialization_II(self):
        from pyspark.serializers import CPickleSerializer

        pickle_ser = CPickleSerializer()

        cdf = self.connect.range(10)
        data = pickle_ser.dumps(cdf)
        cdf2 = pickle_ser.loads(data)
        self.assertEqual(cdf.collect(), cdf2.collect())

    def test_window_spec_serialization(self):
        from pyspark.sql.connect.window import Window
        from pyspark.serializers import CPickleSerializer

        pickle_ser = CPickleSerializer()
        w = Window.partitionBy("some_string").orderBy("value")
        b = pickle_ser.dumps(w)
        w2 = pickle_ser.loads(b)
        self.assertEqual(str(w), str(w2))

    def test_df_getattr_behavior(self):
        cdf = self.connect.range(10)
        sdf = self.spark.range(10)

        sdf._simple_extension = 10
        cdf._simple_extension = 10

        self.assertEqual(sdf._simple_extension, cdf._simple_extension)
        self.assertEqual(type(sdf._simple_extension), type(cdf._simple_extension))

        self.assertTrue(hasattr(cdf, "_simple_extension"))

        with self.temp_env({"PYSPARK_VALIDATE_COLUMN_NAME_LEGACY": None}):
            self.assertTrue(hasattr(cdf, "_simple_extension_does_not_exsit"))
            self.assertIsInstance(getattr(cdf, "_simple_extension_does_not_exsit"), Column)
            self.assertIsInstance(cdf._simple_extension_does_not_exsit, Column)

            # For name starting with '__', still validate it eagerly
            self.assertFalse(hasattr(cdf, "__simple_extension_does_not_exsit"))

        with self.temp_env({"PYSPARK_VALIDATE_COLUMN_NAME_LEGACY": "1"}):
            self.assertFalse(hasattr(cdf, "_simple_extension_does_not_exsit"))

            self.assertFalse(hasattr(cdf, "__simple_extension_does_not_exsit"))

    def test_df_get_item(self):
        # SPARK-41779: test __getitem__

        query = """
            SELECT * FROM VALUES
            (true, 1, NULL), (false, NULL, 2.0), (NULL, 3, 3.0)
            AS tab(a, b, c)
            """

        # +-----+----+----+
        # |    a|   b|   c|
        # +-----+----+----+
        # | true|   1|NULL|
        # |false|NULL| 2.0|
        # | NULL|   3| 3.0|
        # +-----+----+----+

        cdf = self.connect.sql(query)
        sdf = self.spark.sql(query)

        # filter
        self.assert_eq(
            cdf[cdf.a].toPandas(),
            sdf[sdf.a].toPandas(),
        )
        self.assert_eq(
            cdf[cdf.b.isin(2, 3)].toPandas(),
            sdf[sdf.b.isin(2, 3)].toPandas(),
        )
        self.assert_eq(
            cdf[cdf.c > 1.5].toPandas(),
            sdf[sdf.c > 1.5].toPandas(),
        )

        # select
        self.assert_eq(
            cdf[[cdf.a, "b", cdf.c]].toPandas(),
            sdf[[sdf.a, "b", sdf.c]].toPandas(),
        )
        self.assert_eq(
            cdf[(cdf.a, "b", cdf.c)].toPandas(),
            sdf[(sdf.a, "b", sdf.c)].toPandas(),
        )

        # select by index
        self.assertTrue(isinstance(cdf[0], Column))
        self.assertTrue(isinstance(cdf[1], Column))
        self.assertTrue(isinstance(cdf[2], Column))

        self.assert_eq(
            cdf[[cdf[0], cdf[1], cdf[2]]].toPandas(),
            sdf[[sdf[0], sdf[1], sdf[2]]].toPandas(),
        )

        # check error
        with self.assertRaises(PySparkTypeError) as pe:
            cdf[1.5]

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

        with self.assertRaises(PySparkTypeError) as pe:
            cdf[None]

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

        with self.assertRaises(PySparkTypeError) as pe:
            cdf[cdf]

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

    def test_join_condition_column_list_columns(self):
        left_connect_df = self.connect.read.table(self.tbl_name)
        right_connect_df = self.connect.read.table(self.tbl_name2)
        left_spark_df = self.spark.read.table(self.tbl_name)
        right_spark_df = self.spark.read.table(self.tbl_name2)
        joined_plan = left_connect_df.join(
            other=right_connect_df, on=left_connect_df.id == right_connect_df.col1, how="inner"
        )
        joined_plan2 = left_spark_df.join(
            other=right_spark_df, on=left_spark_df.id == right_spark_df.col1, how="inner"
        )
        self.assert_eq(joined_plan.toPandas(), joined_plan2.toPandas())

        joined_plan3 = left_connect_df.join(
            other=right_connect_df,
            on=[
                left_connect_df.id == right_connect_df.col1,
                left_connect_df.name == right_connect_df.col2,
            ],
            how="inner",
        )
        joined_plan4 = left_spark_df.join(
            other=right_spark_df,
            on=[left_spark_df.id == right_spark_df.col1, left_spark_df.name == right_spark_df.col2],
            how="inner",
        )
        self.assert_eq(joined_plan3.toPandas(), joined_plan4.toPandas())

    def test_join_ambiguous_cols(self):
        # SPARK-41812: test join with ambiguous columns
        data1 = [Row(id=1, value="foo"), Row(id=2, value=None)]
        cdf1 = self.connect.createDataFrame(data1)
        sdf1 = self.spark.createDataFrame(data1)

        data2 = [Row(value="bar"), Row(value=None), Row(value="foo")]
        cdf2 = self.connect.createDataFrame(data2)
        sdf2 = self.spark.createDataFrame(data2)

        cdf3 = cdf1.join(cdf2, cdf1["value"] == cdf2["value"])
        sdf3 = sdf1.join(sdf2, sdf1["value"] == sdf2["value"])

        assertDataFrameEqual(cdf3, sdf3)

        cdf4 = cdf1.join(cdf2, cdf1["value"].eqNullSafe(cdf2["value"]))
        sdf4 = sdf1.join(sdf2, sdf1["value"].eqNullSafe(sdf2["value"]))

        assertDataFrameEqual(cdf4, sdf4)

        cdf5 = cdf1.join(
            cdf2, (cdf1["value"] == cdf2["value"]) & (cdf1["value"].eqNullSafe(cdf2["value"]))
        )
        sdf5 = sdf1.join(
            sdf2, (sdf1["value"] == sdf2["value"]) & (sdf1["value"].eqNullSafe(sdf2["value"]))
        )

        assertDataFrameEqual(cdf5, sdf5)

        cdf6 = cdf1.join(cdf2, cdf1["value"] == cdf2["value"]).select(cdf1.value)
        sdf6 = sdf1.join(sdf2, sdf1["value"] == sdf2["value"]).select(sdf1.value)

        assertDataFrameEqual(cdf6, sdf6)

        cdf7 = cdf1.join(cdf2, cdf1["value"] == cdf2["value"]).select(cdf2.value)
        sdf7 = sdf1.join(sdf2, sdf1["value"] == sdf2["value"]).select(sdf2.value)

        assertDataFrameEqual(cdf7, sdf7)

    def test_join_with_cte(self):
        cte_query = "with dt as (select 1 as ida) select ida as id from dt"

        sdf1 = self.spark.range(10)
        sdf2 = self.spark.sql(cte_query)
        sdf3 = sdf1.join(sdf2, sdf1.id == sdf2.id)

        cdf1 = self.connect.range(10)
        cdf2 = self.connect.sql(cte_query)
        cdf3 = cdf1.join(cdf2, cdf1.id == cdf2.id)

        assertDataFrameEqual(cdf3, sdf3)

    def test_with_columns_renamed(self):
        # SPARK-41312: test DataFrame.withColumnsRenamed()
        self.assertEqual(
            self.connect.read.table(self.tbl_name).withColumnRenamed("id", "id_new").schema,
            self.spark.read.table(self.tbl_name).withColumnRenamed("id", "id_new").schema,
        )
        self.assertEqual(
            self.connect.read.table(self.tbl_name)
            .withColumnsRenamed({"id": "id_new", "name": "name_new"})
            .schema,
            self.spark.read.table(self.tbl_name)
            .withColumnsRenamed({"id": "id_new", "name": "name_new"})
            .schema,
        )

    def test_simple_explain_string(self):
        df = self.connect.read.table(self.tbl_name).limit(10)
        result = df._explain_string()
        self.assertGreater(len(result), 0)

    def _check_print_schema(self, query: str):
        with io.StringIO() as buf, redirect_stdout(buf):
            self.spark.sql(query).printSchema()
            print1 = buf.getvalue()
        with io.StringIO() as buf, redirect_stdout(buf):
            self.connect.sql(query).printSchema()
            print2 = buf.getvalue()
        self.assertEqual(print1, print2, query)

        for level in [-1, 0, 1, 2, 3, 4]:
            with io.StringIO() as buf, redirect_stdout(buf):
                self.spark.sql(query).printSchema(level)
                print1 = buf.getvalue()
            with io.StringIO() as buf, redirect_stdout(buf):
                self.connect.sql(query).printSchema(level)
                print2 = buf.getvalue()
            self.assertEqual(print1, print2, query)

    def test_schema(self):
        schema = self.connect.read.table(self.tbl_name).schema
        self.assertEqual(
            StructType(
                [StructField("id", LongType(), True), StructField("name", StringType(), True)]
            ),
            schema,
        )

        # test FloatType, DoubleType, DecimalType, StringType, BooleanType, NullType
        query = """
            SELECT * FROM VALUES
            (float(1.0), double(1.0), 1.0, "1", true, NULL),
            (float(2.0), double(2.0), 2.0, "2", false, NULL),
            (float(3.0), double(3.0), NULL, "3", false, NULL)
            AS tab(a, b, c, d, e, f)
            """
        self.assertEqual(
            self.spark.sql(query).schema,
            self.connect.sql(query).schema,
        )
        self._check_print_schema(query)

        # test TimestampType, DateType
        query = """
            SELECT * FROM VALUES
            (TIMESTAMP('2019-04-12 15:50:00'), DATE('2022-02-22')),
            (TIMESTAMP('2019-04-12 15:50:00'), NULL),
            (NULL, DATE('2022-02-22'))
            AS tab(a, b)
            """
        self.assertEqual(
            self.spark.sql(query).schema,
            self.connect.sql(query).schema,
        )
        self._check_print_schema(query)

        # test DayTimeIntervalType
        query = """ SELECT INTERVAL '100 10:30' DAY TO MINUTE AS interval """
        self.assertEqual(
            self.spark.sql(query).schema,
            self.connect.sql(query).schema,
        )
        self._check_print_schema(query)

        # test MapType
        query = """
            SELECT * FROM VALUES
            (MAP('a', 'ab'), MAP('a', 'ab'), MAP(1, 2, 3, 4)),
            (MAP('x', 'yz'), MAP('x', NULL), NULL),
            (MAP('c', 'de'), NULL, MAP(-1, NULL, -3, -4))
            AS tab(a, b, c)
            """
        self.assertEqual(
            self.spark.sql(query).schema,
            self.connect.sql(query).schema,
        )
        self._check_print_schema(query)

        # test ArrayType
        query = """
            SELECT * FROM VALUES
            (ARRAY('a', 'ab'), ARRAY(1, 2, 3), ARRAY(1, NULL, 3)),
            (ARRAY('x', NULL), NULL, ARRAY(1, 3)),
            (NULL, ARRAY(-1, -2, -3), Array())
            AS tab(a, b, c)
            """
        self.assertEqual(
            self.spark.sql(query).schema,
            self.connect.sql(query).schema,
        )
        self._check_print_schema(query)

        # test StructType
        query = """
            SELECT STRUCT(a, b, c, d), STRUCT(e, f, g), STRUCT(STRUCT(a, b), STRUCT(h)) FROM VALUES
            (float(1.0), double(1.0), 1.0, "1", true, NULL, ARRAY(1, NULL, 3), MAP(1, 2, 3, 4)),
            (float(2.0), double(2.0), 2.0, "2", false, NULL, ARRAY(1, 3), MAP(1, NULL, 3, 4)),
            (float(3.0), double(3.0), NULL, "3", false, NULL, ARRAY(NULL), NULL)
            AS tab(a, b, c, d, e, f, g, h)
            """
        self.assertEqual(
            self.spark.sql(query).schema,
            self.connect.sql(query).schema,
        )
        self._check_print_schema(query)

    def test_to(self):
        # SPARK-41464: test DataFrame.to()

        cdf = self.connect.read.table(self.tbl_name)
        df = self.spark.read.table(self.tbl_name)

        def assert_eq_schema(cdf: CDataFrame, df: DataFrame, schema: StructType):
            cdf_to = cdf.to(schema)
            df_to = df.to(schema)
            self.assertEqual(cdf_to.schema, df_to.schema)
            self.assert_eq(cdf_to.toPandas(), df_to.toPandas())

        # The schema has not changed
        schema = StructType(
            [
                StructField("id", IntegerType(), True),
                StructField("name", StringType(), True),
            ]
        )

        assert_eq_schema(cdf, df, schema)

        # Change schema with struct
        schema2 = StructType([StructField("struct", schema, False)])

        cdf_to = cdf.select(CF.struct("id", "name").alias("struct")).to(schema2)
        df_to = df.select(SF.struct("id", "name").alias("struct")).to(schema2)

        self.assertEqual(cdf_to.schema, df_to.schema)

        # Change the column name
        schema = StructType(
            [
                StructField("col1", IntegerType(), True),
                StructField("col2", StringType(), True),
            ]
        )

        assert_eq_schema(cdf, df, schema)

        # Change the column data type
        schema = StructType(
            [
                StructField("id", StringType(), True),
                StructField("name", StringType(), True),
            ]
        )

        assert_eq_schema(cdf, df, schema)

        # Reduce the column quantity and change data type
        schema = StructType(
            [
                StructField("id", LongType(), True),
            ]
        )

        assert_eq_schema(cdf, df, schema)

        # incompatible field nullability
        schema = StructType([StructField("id", LongType(), False)])
        self.assertRaisesRegex(
            AnalysisException,
            "NULLABLE_COLUMN_OR_FIELD",
            lambda: cdf.to(schema).toPandas(),
        )

        # field cannot upcast
        schema = StructType([StructField("name", LongType())])
        self.assertRaisesRegex(
            AnalysisException,
            "INVALID_COLUMN_OR_FIELD_DATA_TYPE",
            lambda: cdf.to(schema).toPandas(),
        )

        schema = StructType(
            [
                StructField("id", IntegerType(), True),
                StructField("name", IntegerType(), True),
            ]
        )
        self.assertRaisesRegex(
            AnalysisException,
            "INVALID_COLUMN_OR_FIELD_DATA_TYPE",
            lambda: cdf.to(schema).toPandas(),
        )

        # Test map type and array type
        schema = StructType(
            [
                StructField("id", StringType(), True),
                StructField("my_map", MapType(StringType(), IntegerType(), False), True),
                StructField("my_array", ArrayType(IntegerType(), False), True),
            ]
        )
        cdf = self.connect.read.table(self.tbl_name4)
        df = self.spark.read.table(self.tbl_name4)

        assert_eq_schema(cdf, df, schema)

    def test_toDF(self):
        # SPARK-41310: test DataFrame.toDF()
        self.assertEqual(
            self.connect.read.table(self.tbl_name).toDF("col1", "col2").schema,
            self.spark.read.table(self.tbl_name).toDF("col1", "col2").schema,
        )

    def test_print_schema(self):
        # SPARK-41216: Test print schema
        tree_str = self.connect.sql("SELECT 1 AS X, 2 AS Y").schema.treeString()
        # root
        #  |-- X: integer (nullable = false)
        #  |-- Y: integer (nullable = false)
        expected = "root\n |-- X: integer (nullable = false)\n |-- Y: integer (nullable = false)\n"
        self.assertEqual(tree_str, expected)

    def test_is_local(self):
        # SPARK-41216: Test is local
        self.assertTrue(self.connect.sql("SHOW DATABASES").isLocal())
        self.assertFalse(self.connect.read.table(self.tbl_name).isLocal())

    def test_is_streaming(self):
        # SPARK-41216: Test is streaming
        self.assertFalse(self.connect.read.table(self.tbl_name).isStreaming)
        self.assertFalse(self.connect.sql("SELECT 1 AS X LIMIT 0").isStreaming)

    def test_input_files(self):
        # SPARK-41216: Test input files
        tmpPath = tempfile.mkdtemp()
        shutil.rmtree(tmpPath)
        try:
            self.df_text.write.text(tmpPath)

            input_files_list1 = (
                self.spark.read.format("text").schema("id STRING").load(path=tmpPath).inputFiles()
            )
            input_files_list2 = (
                self.connect.read.format("text").schema("id STRING").load(path=tmpPath).inputFiles()
            )

            self.assertTrue(len(input_files_list1) > 0)
            self.assertEqual(len(input_files_list1), len(input_files_list2))
            for file_path in input_files_list2:
                self.assertTrue(file_path in input_files_list1)
        finally:
            shutil.rmtree(tmpPath)

    def test_limit_offset(self):
        df = self.connect.read.table(self.tbl_name)
        pd = df.limit(10).offset(1).toPandas()
        self.assertEqual(9, len(pd.index))
        pd2 = df.offset(98).limit(10).toPandas()
        self.assertEqual(2, len(pd2.index))

    def test_tail(self):
        df = self.connect.read.table(self.tbl_name)
        df2 = self.spark.read.table(self.tbl_name)
        self.assertEqual(df.tail(10), df2.tail(10))

    def test_sql(self):
        pdf = self.connect.sql("SELECT 1").toPandas()
        self.assertEqual(1, len(pdf.index))

    def test_sql_with_named_args(self):
        sqlText = "SELECT *, element_at(:m, 'a') FROM range(10) WHERE id > :minId"
        df = self.connect.sql(
            sqlText, args={"minId": 7, "m": CF.create_map(CF.lit("a"), CF.lit(1))}
        )
        df2 = self.spark.sql(sqlText, args={"minId": 7, "m": SF.create_map(SF.lit("a"), SF.lit(1))})
        self.assert_eq(df.toPandas(), df2.toPandas())

    def test_namedargs_with_global_limit(self):
        sqlText = """SELECT * FROM VALUES (TIMESTAMP('2022-12-25 10:30:00'), 1) as tab(date, val)
         where val = :val"""
        df = self.connect.sql(sqlText, args={"val": 1})
        df2 = self.spark.sql(sqlText, args={"val": 1})
        self.assert_eq(df.toPandas(), df2.toPandas())

        self.assert_eq(df.first()[0], datetime.datetime(2022, 12, 25, 10, 30))
        self.assert_eq(df.first().date, datetime.datetime(2022, 12, 25, 10, 30))
        self.assert_eq(df.first()[1], 1)
        self.assert_eq(df.first().val, 1)

    def test_sql_with_pos_args(self):
        sqlText = "SELECT *, element_at(?, 1) FROM range(10) WHERE id > ?"
        df = self.connect.sql(sqlText, args=[CF.array(CF.lit(1)), 7])
        df2 = self.spark.sql(sqlText, args=[SF.array(SF.lit(1)), 7])
        self.assert_eq(df.toPandas(), df2.toPandas())

    def test_sql_with_invalid_args(self):
        sqlText = "SELECT ?, ?, ?"
        for session in [self.connect, self.spark]:
            with self.assertRaises(PySparkTypeError) as pe:
                session.sql(sqlText, args={1, 2, 3})

            self.check_error(
                exception=pe.exception,
                errorClass="INVALID_TYPE",
                messageParameters={"arg_name": "args", "arg_type": "set"},
            )

    def test_deduplicate(self):
        # SPARK-41326: test distinct and dropDuplicates.
        df = self.connect.read.table(self.tbl_name)
        df2 = self.spark.read.table(self.tbl_name)
        self.assert_eq(df.distinct().toPandas(), df2.distinct().toPandas())
        self.assert_eq(df.dropDuplicates().toPandas(), df2.dropDuplicates().toPandas())
        self.assert_eq(
            df.dropDuplicates(["name"]).toPandas(), df2.dropDuplicates(["name"]).toPandas()
        )

    def test_drop(self):
        # SPARK-41169: test drop
        query = """
            SELECT * FROM VALUES
            (false, 1, NULL), (false, NULL, 2), (NULL, 3, 3)
            AS tab(a, b, c)
            """

        cdf = self.connect.sql(query)
        sdf = self.spark.sql(query)
        self.assert_eq(
            cdf.drop("a").toPandas(),
            sdf.drop("a").toPandas(),
        )
        self.assert_eq(
            cdf.drop("a", "b").toPandas(),
            sdf.drop("a", "b").toPandas(),
        )
        self.assert_eq(
            cdf.drop("a", "x").toPandas(),
            sdf.drop("a", "x").toPandas(),
        )
        self.assert_eq(
            cdf.drop(cdf.a, "x").toPandas(),
            sdf.drop(sdf.a, "x").toPandas(),
        )

    def test_subquery_alias(self) -> None:
        # SPARK-40938: test subquery alias.
        plan_text = (
            self.connect.read.table(self.tbl_name)
            .alias("special_alias")
            ._explain_string(extended=True)
        )
        self.assertTrue("special_alias" in plan_text)

    def test_sort(self):
        # SPARK-41332: test sort
        query = """
            SELECT * FROM VALUES
            (false, 1, NULL), (false, NULL, 2.0), (NULL, 3, 3.0)
            AS tab(a, b, c)
            """
        # +-----+----+----+
        # |    a|   b|   c|
        # +-----+----+----+
        # |false|   1|NULL|
        # |false|NULL| 2.0|
        # | NULL|   3| 3.0|
        # +-----+----+----+

        cdf = self.connect.sql(query)
        sdf = self.spark.sql(query)
        self.assert_eq(
            cdf.sort("a").toPandas(),
            sdf.sort("a").toPandas(),
        )
        self.assert_eq(
            cdf.sort("c").toPandas(),
            sdf.sort("c").toPandas(),
        )
        self.assert_eq(
            cdf.sort("b").toPandas(),
            sdf.sort("b").toPandas(),
        )
        self.assert_eq(
            cdf.sort(cdf.c, "b").toPandas(),
            sdf.sort(sdf.c, "b").toPandas(),
        )
        self.assert_eq(
            cdf.sort(cdf.c.desc(), "b").toPandas(),
            sdf.sort(sdf.c.desc(), "b").toPandas(),
        )
        self.assert_eq(
            cdf.sort(cdf.c.desc(), cdf.a.asc()).toPandas(),
            sdf.sort(sdf.c.desc(), sdf.a.asc()).toPandas(),
        )

    def test_range(self):
        self.assert_eq(
            self.connect.range(start=0, end=10).toPandas(),
            self.spark.range(start=0, end=10).toPandas(),
        )
        self.assert_eq(
            self.connect.range(start=0, end=10, step=3).toPandas(),
            self.spark.range(start=0, end=10, step=3).toPandas(),
        )
        self.assert_eq(
            self.connect.range(start=0, end=10, step=3, numPartitions=2).toPandas(),
            self.spark.range(start=0, end=10, step=3, numPartitions=2).toPandas(),
        )
        # SPARK-41301
        self.assert_eq(
            self.connect.range(10).toPandas(), self.connect.range(start=0, end=10).toPandas()
        )

    def test_create_global_temp_view(self):
        # SPARK-41127: test global temp view creation.
        with self.temp_view("view_1"):
            self.connect.sql("SELECT 1 AS X LIMIT 0").createGlobalTempView("view_1")
            self.connect.sql("SELECT 2 AS X LIMIT 1").createOrReplaceGlobalTempView("view_1")
            self.assertTrue(self.spark.catalog.tableExists("global_temp.view_1"))

            # Test when creating a view which is already exists but
            self.assertTrue(self.spark.catalog.tableExists("global_temp.view_1"))
            with self.assertRaises(AnalysisException):
                self.connect.sql("SELECT 1 AS X LIMIT 0").createGlobalTempView("view_1")

    def test_create_session_local_temp_view(self):
        # SPARK-41372: test session local temp view creation.
        with self.temp_view("view_local_temp"):
            self.connect.sql("SELECT 1 AS X").createTempView("view_local_temp")
            self.assertEqual(self.connect.sql("SELECT * FROM view_local_temp").count(), 1)
            self.connect.sql("SELECT 1 AS X LIMIT 0").createOrReplaceTempView("view_local_temp")
            self.assertEqual(self.connect.sql("SELECT * FROM view_local_temp").count(), 0)

            # Test when creating a view which is already exists but
            with self.assertRaises(AnalysisException):
                self.connect.sql("SELECT 1 AS X LIMIT 0").createTempView("view_local_temp")

    def test_select_expr(self):
        # SPARK-41201: test selectExpr API.
        self.assert_eq(
            self.connect.read.table(self.tbl_name).selectExpr("id * 2").toPandas(),
            self.spark.read.table(self.tbl_name).selectExpr("id * 2").toPandas(),
        )
        self.assert_eq(
            self.connect.read.table(self.tbl_name)
            .selectExpr(["id * 2", "cast(name as long) as name"])
            .toPandas(),
            self.spark.read.table(self.tbl_name)
            .selectExpr(["id * 2", "cast(name as long) as name"])
            .toPandas(),
        )

        self.assert_eq(
            self.connect.read.table(self.tbl_name)
            .selectExpr("id * 2", "cast(name as long) as name")
            .toPandas(),
            self.spark.read.table(self.tbl_name)
            .selectExpr("id * 2", "cast(name as long) as name")
            .toPandas(),
        )

    def test_select_star(self):
        data = [Row(a=1, b=Row(c=2, d=Row(e=3)))]

        # +---+--------+
        # |  a|       b|
        # +---+--------+
        # |  1|{2, {3}}|
        # +---+--------+

        cdf = self.connect.createDataFrame(data=data)
        sdf = self.spark.createDataFrame(data=data)

        self.assertEqual(
            cdf.select("*").collect(),
            sdf.select("*").collect(),
        )
        self.assertEqual(
            cdf.select("a", "*").collect(),
            sdf.select("a", "*").collect(),
        )
        self.assertEqual(
            cdf.select("a", "b").collect(),
            sdf.select("a", "b").collect(),
        )
        self.assertEqual(
            cdf.select("a", "b.*").collect(),
            sdf.select("a", "b.*").collect(),
        )

    def test_union_by_name(self):
        # SPARK-41832: Test unionByName
        data1 = [(1, 2, 3)]
        data2 = [(6, 2, 5)]
        df1_connect = self.connect.createDataFrame(data1, ["a", "b", "c"])
        df2_connect = self.connect.createDataFrame(data2, ["a", "b", "c"])
        union_df_connect = df1_connect.unionByName(df2_connect)

        df1_spark = self.spark.createDataFrame(data1, ["a", "b", "c"])
        df2_spark = self.spark.createDataFrame(data2, ["a", "b", "c"])
        union_df_spark = df1_spark.unionByName(df2_spark)

        self.assert_eq(union_df_connect.toPandas(), union_df_spark.toPandas())

        df2_connect = self.connect.createDataFrame(data2, ["a", "B", "C"])
        union_df_connect = df1_connect.unionByName(df2_connect, allowMissingColumns=True)

        df2_spark = self.spark.createDataFrame(data2, ["a", "B", "C"])
        union_df_spark = df1_spark.unionByName(df2_spark, allowMissingColumns=True)

        self.assert_eq(union_df_connect.toPandas(), union_df_spark.toPandas())

    def test_observe(self):
        # SPARK-41527: test DataFrame.observe()
        observation_name = "my_metric"

        self.assert_eq(
            self.connect.read.table(self.tbl_name)
            .filter("id > 3")
            .observe(observation_name, CF.min("id"), CF.max("id"), CF.sum("id"))
            .toPandas(),
            self.spark.read.table(self.tbl_name)
            .filter("id > 3")
            .observe(observation_name, SF.min("id"), SF.max("id"), SF.sum("id"))
            .toPandas(),
        )

        from pyspark.sql.connect.observation import Observation as ConnectObservation
        from pyspark.sql.observation import Observation

        cobservation = ConnectObservation(observation_name)
        observation = Observation(observation_name)

        cdf = (
            self.connect.read.table(self.tbl_name)
            .filter("id > 3")
            .observe(cobservation, CF.min("id"), CF.max("id"), CF.sum("id"))
            .toPandas()
        )
        df = (
            self.spark.read.table(self.tbl_name)
            .filter("id > 3")
            .observe(observation, SF.min("id"), SF.max("id"), SF.sum("id"))
            .toPandas()
        )

        self.assert_eq(cdf, df)

        self.assertEqual(cobservation.get, observation.get)

        observed_metrics = cdf.attrs["observed_metrics"]
        self.assert_eq(len(observed_metrics), 1)
        self.assert_eq(observed_metrics[0].name, observation_name)
        self.assert_eq(len(observed_metrics[0].metrics), 3)
        for metric in observed_metrics[0].metrics:
            self.assertIsInstance(metric, ProtoExpression.Literal)
        values = list(map(lambda metric: metric.long, observed_metrics[0].metrics))
        self.assert_eq(values, [4, 99, 4944])

        with self.assertRaises(PySparkValueError) as pe:
            self.connect.read.table(self.tbl_name).observe(observation_name)

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

        with self.assertRaises(PySparkTypeError) as pe:
            self.connect.read.table(self.tbl_name).observe(observation_name, CF.lit(1), "id")

        self.check_error(
            exception=pe.exception,
            errorClass="NOT_EXPECTED_TYPE",
            messageParameters={
                "expected_type": "list[Column]",
                "arg_name": "exprs",
                "arg_type": "tuple",
            },
        )

    def test_with_columns(self):
        # SPARK-41256: test withColumn(s).
        self.assert_eq(
            self.connect.read.table(self.tbl_name).withColumn("id", CF.lit(False)).toPandas(),
            self.spark.read.table(self.tbl_name).withColumn("id", SF.lit(False)).toPandas(),
        )

        self.assert_eq(
            self.connect.read.table(self.tbl_name)
            .withColumns({"id": CF.lit(False), "col_not_exist": CF.lit(False)})
            .toPandas(),
            self.spark.read.table(self.tbl_name)
            .withColumns(
                {
                    "id": SF.lit(False),
                    "col_not_exist": SF.lit(False),
                }
            )
            .toPandas(),
        )

    def test_hint(self):
        # SPARK-41349: Test hint
        self.assert_eq(
            self.connect.read.table(self.tbl_name).hint("COALESCE", 3000).toPandas(),
            self.spark.read.table(self.tbl_name).hint("COALESCE", 3000).toPandas(),
        )

        # Hint with unsupported name will be ignored
        self.assert_eq(
            self.connect.read.table(self.tbl_name).hint("illegal").toPandas(),
            self.spark.read.table(self.tbl_name).hint("illegal").toPandas(),
        )

        # Hint with all supported parameter values
        such_a_nice_list = ["itworks1", "itworks2", "itworks3"]
        self.assert_eq(
            self.connect.read.table(self.tbl_name).hint("my awesome hint", 1.2345, 2).toPandas(),
            self.spark.read.table(self.tbl_name).hint("my awesome hint", 1.2345, 2).toPandas(),
        )

        # Hint with unsupported parameter values
        with self.assertRaises(AnalysisException):
            self.connect.read.table(self.tbl_name).hint("REPARTITION", "id+1").toPandas()

        # Hint with unsupported parameter types
        with self.assertRaises(TypeError):
            self.connect.read.table(self.tbl_name).hint("REPARTITION", range(5)).toPandas()

        # Hint with unsupported parameter types
        with self.assertRaises(TypeError):
            self.connect.read.table(self.tbl_name).hint(
                "my awesome hint", 1.2345, 2, such_a_nice_list, range(6)
            ).toPandas()

        # Hint with wrong combination
        with self.assertRaises(AnalysisException):
            self.connect.read.table(self.tbl_name).hint("REPARTITION", "id", 3).toPandas()

    def test_join_hint(self):
        cdf1 = self.connect.createDataFrame([(2, "Alice"), (5, "Bob")], schema=["age", "name"])
        cdf2 = self.connect.createDataFrame(
            [Row(height=80, name="Tom"), Row(height=85, name="Bob")]
        )

        self.assertTrue(
            "BroadcastHashJoin" in cdf1.join(cdf2.hint("BROADCAST"), "name")._explain_string()
        )
        self.assertTrue("SortMergeJoin" in cdf1.join(cdf2.hint("MERGE"), "name")._explain_string())
        self.assertTrue(
            "ShuffledHashJoin" in cdf1.join(cdf2.hint("SHUFFLE_HASH"), "name")._explain_string()
        )

    def test_extended_hint_types(self):
        cdf = self.connect.range(100).toDF("id")

        cdf.hint(
            "my awesome hint",
            1.2345,
            "what",
            ["itworks1", "itworks2", "itworks3"],
        ).show()

        with self.assertRaises(PySparkTypeError) as pe:
            cdf.hint(
                "my awesome hint",
                1.2345,
                "what",
                {"itworks1": "itworks2"},
            ).show()

        self.check_error(
            exception=pe.exception,
            errorClass="INVALID_ITEM_FOR_CONTAINER",
            messageParameters={
                "arg_name": "parameters",
                "allowed_types": "str, float, int, Column, list[str], list[float], list[int]",
                "item_type": "dict",
            },
        )

    def test_empty_dataset(self):
        # SPARK-41005: Test arrow based collection with empty dataset.
        self.assertTrue(
            self.connect.sql("SELECT 1 AS X LIMIT 0")
            .toPandas()
            .equals(self.spark.sql("SELECT 1 AS X LIMIT 0").toPandas())
        )
        pdf = self.connect.sql("SELECT 1 AS X LIMIT 0").toPandas()
        self.assertEqual(0, len(pdf))  # empty dataset
        self.assertEqual(1, len(pdf.columns))  # one column
        self.assertEqual("X", pdf.columns[0])

    def test_is_empty(self):
        # SPARK-41212: Test is empty
        self.assertFalse(self.connect.sql("SELECT 1 AS X").isEmpty())
        self.assertTrue(self.connect.sql("SELECT 1 AS X LIMIT 0").isEmpty())

    def test_is_empty_with_unsupported_types(self):
        df = self.spark.sql("SELECT INTERVAL '10-8' YEAR TO MONTH AS interval")
        self.assertEqual(df.count(), 1)
        self.assertFalse(df.isEmpty())

    def test_session(self):
        self.assertEqual(self.connect, self.connect.sql("SELECT 1").sparkSession)

    def test_show(self):
        # SPARK-41111: Test the show method
        show_str = self.connect.sql("SELECT 1 AS X, 2 AS Y")._show_string()
        # +---+---+
        # |  X|  Y|
        # +---+---+
        # |  1|  2|
        # +---+---+
        expected = "+---+---+\n|  X|  Y|\n+---+---+\n|  1|  2|\n+---+---+\n"
        self.assertEqual(show_str, expected)

    def test_repr(self):
        # SPARK-41213: Test the __repr__ method
        query = """SELECT * FROM VALUES (1L, NULL), (3L, "Z") AS tab(a, b)"""
        self.assertEqual(
            self.connect.sql(query).__repr__(),
            self.spark.sql(query).__repr__(),
        )

    def test_explain_string(self):
        # SPARK-41122: test explain API.
        plan_str = self.connect.sql("SELECT 1")._explain_string(extended=True)
        self.assertTrue("Parsed Logical Plan" in plan_str)
        self.assertTrue("Analyzed Logical Plan" in plan_str)
        self.assertTrue("Optimized Logical Plan" in plan_str)
        self.assertTrue("Physical Plan" in plan_str)

        with self.assertRaises(PySparkValueError) as pe:
            self.connect.sql("SELECT 1")._explain_string(mode="unknown")
        self.check_error(
            exception=pe.exception,
            errorClass="VALUE_NOT_ALLOWED",
            messageParameters={
                "arg_name": "explain_mode",
                "allowed_values": "['simple', 'extended', 'codegen', 'cost', 'formatted']",
            },
        )

    def test_count(self) -> None:
        # SPARK-41308: test count() API.
        self.assertEqual(
            self.connect.read.table(self.tbl_name).count(),
            self.spark.read.table(self.tbl_name).count(),
        )

    def test_simple_transform(self) -> None:
        """SPARK-41203: Support DF.transform"""

        def transform_df(input_df: CDataFrame) -> CDataFrame:
            return input_df.select((CF.col("id") + CF.lit(10)).alias("id"))

        df = self.connect.range(1, 100)
        result_left = df.transform(transform_df).collect()
        result_right = self.connect.range(11, 110).collect()
        self.assertEqual(result_right, result_left)

        # Check assertion.
        with self.assertRaises(AssertionError):
            df.transform(lambda x: 2)  # type: ignore

    def test_alias(self) -> None:
        """Testing supported and unsupported alias"""
        col0 = (
            self.connect.range(1, 10)
            .select(CF.col("id").alias("name", metadata={"max": 99}))
            .schema.names[0]
        )
        self.assertEqual("name", col0)

        with self.assertRaises(SparkConnectException) as exc:
            self.connect.range(1, 10).select(CF.col("id").alias("this", "is", "not")).collect()
        self.assertIn("(this, is, not)", str(exc.exception))

    def test_alias_metadata(self):
        df = self.connect.createDataFrame([("",)], ["a"])
        df = df.withMetadata("a", {"foo": "bar"})
        self.assertEqual(df.schema["a"].metadata, {"foo": "bar"})

        # SPARK-51426: Ensure setting metadata to `{}` clears it
        df = df.select([CF.col("a").alias("a", metadata={})])
        self.assertEqual(df.schema["a"].metadata, {})

        df = df.withMetadata("a", {"baz": "burr"})
        self.assertEqual(df.schema["a"].metadata, {"baz": "burr"})

        df = df.withMetadata("a", {})
        self.assertEqual(df.schema["a"].metadata, {})

    def test_column_regexp(self) -> None:
        # SPARK-41438: test dataframe.colRegex()
        ndf = self.connect.read.table(self.tbl_name3)
        df = self.spark.read.table(self.tbl_name3)

        self.assert_eq(
            ndf.select(ndf.colRegex("`tes.*\n.*mn`")).toPandas(),
            df.select(df.colRegex("`tes.*\n.*mn`")).toPandas(),
        )

    def test_repartition(self) -> None:
        # SPARK-41354: test dataframe.repartition(numPartitions)
        self.assert_eq(
            self.connect.read.table(self.tbl_name).repartition(10).toPandas(),
            self.spark.read.table(self.tbl_name).repartition(10).toPandas(),
        )

        self.assert_eq(
            self.connect.read.table(self.tbl_name).coalesce(10).toPandas(),
            self.spark.read.table(self.tbl_name).coalesce(10).toPandas(),
        )

    def test_repartition_by_expression(self) -> None:
        # SPARK-41354: test dataframe.repartition(expressions)
        self.assert_eq(
            self.connect.read.table(self.tbl_name).repartition(10, "id").toPandas(),
            self.spark.read.table(self.tbl_name).repartition(10, "id").toPandas(),
        )

        self.assert_eq(
            self.connect.read.table(self.tbl_name).repartition("id").toPandas(),
            self.spark.read.table(self.tbl_name).repartition("id").toPandas(),
        )

        # repartition with unsupported parameter values
        with self.assertRaises(AnalysisException):
            self.connect.read.table(self.tbl_name).repartition("id+1").toPandas()

    def test_repartition_by_range(self) -> None:
        # SPARK-41354: test dataframe.repartitionByRange(expressions)
        cdf = self.connect.read.table(self.tbl_name)
        sdf = self.spark.read.table(self.tbl_name)

        self.assert_eq(
            cdf.repartitionByRange(10, "id").toPandas(),
            sdf.repartitionByRange(10, "id").toPandas(),
        )

        self.assert_eq(
            cdf.repartitionByRange("id").toPandas(),
            sdf.repartitionByRange("id").toPandas(),
        )

        self.assert_eq(
            cdf.repartitionByRange(cdf.id.desc()).toPandas(),
            sdf.repartitionByRange(sdf.id.desc()).toPandas(),
        )

        # repartitionByRange with unsupported parameter values
        with self.assertRaises(AnalysisException):
            self.connect.read.table(self.tbl_name).repartitionByRange("id+1").toPandas()

    def test_crossjoin(self):
        # SPARK-41227: Test CrossJoin
        connect_df = self.connect.read.table(self.tbl_name)
        spark_df = self.spark.read.table(self.tbl_name)
        self.assert_eq(
            set(
                connect_df.select("id")
                .join(other=connect_df.select("name"), how="cross")
                .toPandas()
            ),
            set(spark_df.select("id").join(other=spark_df.select("name"), how="cross").toPandas()),
        )
        self.assert_eq(
            set(connect_df.select("id").crossJoin(other=connect_df.select("name")).toPandas()),
            set(spark_df.select("id").crossJoin(other=spark_df.select("name")).toPandas()),
        )

    def test_self_join(self):
        # SPARK-47713: this query fails in classic spark
        df1 = self.connect.createDataFrame([(1, "a")], schema=["i", "j"])
        df1_filter = df1.filter(df1.i > 0)
        df2 = df1.join(df1_filter, df1.i == 1)
        self.assertEqual(df2.count(), 1)
        self.assertEqual(df2.columns, ["i", "j", "i", "j"])
        self.assertEqual(list(df2.first()), [1, "a", 1, "a"])

    def test_with_metadata(self):
        cdf = self.connect.createDataFrame(data=[(2, "Alice"), (5, "Bob")], schema=["age", "name"])
        self.assertEqual(cdf.schema["age"].metadata, {})
        self.assertEqual(cdf.schema["name"].metadata, {})

        cdf1 = cdf.withMetadata(columnName="age", metadata={"max_age": 5})
        self.assertEqual(cdf1.schema["age"].metadata, {"max_age": 5})

        cdf2 = cdf.withMetadata(columnName="name", metadata={"names": ["Alice", "Bob"]})
        self.assertEqual(cdf2.schema["name"].metadata, {"names": ["Alice", "Bob"]})

        with self.assertRaises(PySparkTypeError) as pe:
            cdf.withMetadata(columnName="name", metadata=["magic"])

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

    def test_version(self):
        self.assertEqual(
            self.connect.version,
            self.spark.version,
        )

    def test_same_semantics(self):
        plan = self.connect.sql("SELECT 1")
        other = self.connect.sql("SELECT 1")
        self.assertTrue(plan.sameSemantics(other))

    def test_semantic_hash(self):
        plan = self.connect.sql("SELECT 1")
        other = self.connect.sql("SELECT 1")
        self.assertEqual(
            plan.semanticHash(),
            other.semanticHash(),
        )

    def test_sql_with_command(self):
        # SPARK-42705: spark.sql should return values from the command.
        self.assertEqual(
            self.connect.sql("show functions").collect(), self.spark.sql("show functions").collect()
        )

    def test_df_caache(self):
        df = self.connect.range(10)
        df.cache()
        self.assert_eq(10, df.count())
        self.assertTrue(df.is_cached)

    def test_parse_col_name(self):
        from pyspark.sql.connect.types import parse_attr_name

        self.assert_eq(parse_attr_name(""), [""])

        self.assert_eq(parse_attr_name("a"), ["a"])
        self.assert_eq(parse_attr_name("`a`"), ["a"])
        self.assert_eq(parse_attr_name("`a"), None)
        self.assert_eq(parse_attr_name("a`"), None)

        self.assert_eq(parse_attr_name("`a`.b"), ["a", "b"])
        self.assert_eq(parse_attr_name("`a`.`b`"), ["a", "b"])
        self.assert_eq(parse_attr_name("`a```.b"), ["a`", "b"])
        self.assert_eq(parse_attr_name("`a``.b"), None)

        self.assert_eq(parse_attr_name("a.b.c"), ["a", "b", "c"])
        self.assert_eq(parse_attr_name("`a`.`b`.`c`"), ["a", "b", "c"])
        self.assert_eq(parse_attr_name("a.`b`.c"), ["a", "b", "c"])

        self.assert_eq(parse_attr_name("`a.b.c`"), ["a.b.c"])
        self.assert_eq(parse_attr_name("a.`b.c`"), ["a", "b.c"])
        self.assert_eq(parse_attr_name("`a.b`.c"), ["a.b", "c"])
        self.assert_eq(parse_attr_name("`a.b.c"), None)
        self.assert_eq(parse_attr_name("a.b.c`"), None)
        self.assert_eq(parse_attr_name("`a.`b.`c"), None)
        self.assert_eq(parse_attr_name("a`.b`.c`"), None)

        self.assert_eq(parse_attr_name("`ab..c`e.f"), None)

    def test_verify_col_name(self):
        from pyspark.sql.connect.types import verify_col_name

        cdf = (
            self.connect.range(10)
            .withColumn("v", CF.lit(123))
            .withColumn("s", CF.struct("id", "v"))
            .withColumn("m", CF.struct("s", "v"))
            .withColumn("a", CF.array("s"))
        )

        # root
        # |-- id: long (nullable = false)
        # |-- v: integer (nullable = false)
        # |-- s: struct (nullable = false)
        # |    |-- id: long (nullable = false)
        # |    |-- v: integer (nullable = false)
        # |-- m: struct (nullable = false)
        # |    |-- s: struct (nullable = false)
        # |    |    |-- id: long (nullable = false)
        # |    |    |-- v: integer (nullable = false)
        # |    |-- v: integer (nullable = false)
        # |-- a: array (nullable = false)
        # |    |-- element: struct (containsNull = false)
        # |    |    |-- id: long (nullable = false)
        # |    |    |-- v: integer (nullable = false)

        self.assertTrue(verify_col_name("id", cdf.schema))
        self.assertTrue(verify_col_name("`id`", cdf.schema))

        self.assertTrue(verify_col_name("v", cdf.schema))
        self.assertTrue(verify_col_name("`v`", cdf.schema))

        self.assertFalse(verify_col_name("x", cdf.schema))
        self.assertFalse(verify_col_name("`x`", cdf.schema))

        self.assertTrue(verify_col_name("s", cdf.schema))
        self.assertTrue(verify_col_name("`s`", cdf.schema))
        self.assertTrue(verify_col_name("s.id", cdf.schema))
        self.assertTrue(verify_col_name("s.`id`", cdf.schema))
        self.assertTrue(verify_col_name("`s`.id", cdf.schema))
        self.assertTrue(verify_col_name("`s`.`id`", cdf.schema))
        self.assertFalse(verify_col_name("`s.id`", cdf.schema))

        self.assertTrue(verify_col_name("m", cdf.schema))
        self.assertTrue(verify_col_name("`m`", cdf.schema))
        self.assertTrue(verify_col_name("m.s.id", cdf.schema))
        self.assertTrue(verify_col_name("m.s.`id`", cdf.schema))
        self.assertTrue(verify_col_name("m.`s`.id", cdf.schema))
        self.assertTrue(verify_col_name("`m`.`s`.`id`", cdf.schema))
        self.assertFalse(verify_col_name("m.`s.id`", cdf.schema))

        self.assertTrue(verify_col_name("a", cdf.schema))
        self.assertTrue(verify_col_name("`a`", cdf.schema))
        self.assertTrue(verify_col_name("a.`v`", cdf.schema))
        self.assertTrue(verify_col_name("a.`v`", cdf.schema))
        self.assertTrue(verify_col_name("`a`.v", cdf.schema))
        self.assertTrue(verify_col_name("`a`.`v`", cdf.schema))
        self.assertFalse(verify_col_name("`a`.`x`", cdf.schema))

        cdf = (
            self.connect.range(10)
            .withColumn("v", CF.lit(123))
            .withColumn("s.s", CF.struct("id", "v"))
            .withColumn("m`", CF.struct("`s.s`", "v"))
        )

        # root
        # |-- id: long (nullable = false)
        # |-- v: string (nullable = false)
        # |-- s.s: struct (nullable = false)
        # |    |-- id: long (nullable = false)
        # |    |-- v: string (nullable = false)
        # |-- m`: struct (nullable = false)
        # |    |-- s.s: struct (nullable = false)
        # |    |    |-- id: long (nullable = false)
        # |    |    |-- v: string (nullable = false)
        # |    |-- v: string (nullable = false)

        self.assertFalse(verify_col_name("s", cdf.schema))
        self.assertFalse(verify_col_name("`s`", cdf.schema))
        self.assertFalse(verify_col_name("s.s", cdf.schema))
        self.assertFalse(verify_col_name("s.`s`", cdf.schema))
        self.assertFalse(verify_col_name("`s`.s", cdf.schema))
        self.assertTrue(verify_col_name("`s.s`", cdf.schema))

        self.assertFalse(verify_col_name("m", cdf.schema))
        self.assertFalse(verify_col_name("`m`", cdf.schema))
        self.assertTrue(verify_col_name("`m```", cdf.schema))

        self.assertFalse(verify_col_name("`m```.s", cdf.schema))
        self.assertFalse(verify_col_name("`m```.`s`", cdf.schema))
        self.assertFalse(verify_col_name("`m```.s.s", cdf.schema))
        self.assertFalse(verify_col_name("`m```.s.`s`", cdf.schema))
        self.assertTrue(verify_col_name("`m```.`s.s`", cdf.schema))

        self.assertFalse(verify_col_name("`m```.s.s.v", cdf.schema))
        self.assertFalse(verify_col_name("`m```.s.`s`.v", cdf.schema))
        self.assertTrue(verify_col_name("`m```.`s.s`.v", cdf.schema))
        self.assertTrue(verify_col_name("`m```.`s.s`.`v`", cdf.schema))

    def test_truncate_message(self):
        cdf1 = self.connect.createDataFrame(
            [
                ("a B c"),
                ("X y Z"),
            ],
            ["a" * 4096],
        )
        plan1 = cdf1._plan.to_proto(self.connect._client)

        proto_string_1 = self.connect._client._proto_to_string(plan1, False)
        self.assertTrue(len(proto_string_1) > 10000, len(proto_string_1))
        proto_string_truncated_1 = self.connect._client._proto_to_string(plan1, True)
        self.assertTrue(len(proto_string_truncated_1) < 4000, len(proto_string_truncated_1))

        cdf2 = cdf1.select("a" * 4096, "a" * 4096, "a" * 4096)
        plan2 = cdf2._plan.to_proto(self.connect._client)

        proto_string_2 = self.connect._client._proto_to_string(plan2, False)
        self.assertTrue(len(proto_string_2) > 20000, len(proto_string_2))
        proto_string_truncated_2 = self.connect._client._proto_to_string(plan2, True)
        self.assertTrue(len(proto_string_truncated_2) < 8000, len(proto_string_truncated_2))

        cdf3 = cdf1.select("a" * 4096)
        for _ in range(64):
            cdf3 = cdf3.select("a" * 4096)
        plan3 = cdf3._plan.to_proto(self.connect._client)

        proto_string_3 = self.connect._client._proto_to_string(plan3, False)
        self.assertTrue(len(proto_string_3) > 128000, len(proto_string_3))
        proto_string_truncated_3 = self.connect._client._proto_to_string(plan3, True)
        self.assertTrue(len(proto_string_truncated_3) < 64000, len(proto_string_truncated_3))

    def test_plan_compression(self):
        self.assertTrue(self.connect._client._zstd_module is not None)
        self.connect.range(1).count()
        default_plan_compression_threshold = self.connect._client._plan_compression_threshold
        self.assertTrue(default_plan_compression_threshold > 0)
        self.assertEqual(self.connect._client._plan_compression_algorithm, "ZSTD")
        try:
            self.connect._client._plan_compression_threshold = 1000

            # Small plan should not be compressed
            cdf1 = self.connect.range(1).select(CF.lit("Apache Spark"))
            plan1 = cdf1._plan.to_proto(self.connect._client)
            self.assertTrue(plan1.root is not None)
            self.assertEqual(cdf1.count(), 1)

            # Large plan should be compressed
            cdf2 = self.connect.range(1).select(CF.lit("Apache Spark" * 1000))
            plan2 = cdf2._plan.to_proto(self.connect._client)
            self.assertTrue(plan2.compressed_operation is not None)
            # Test compressed relation
            self.assertEqual(cdf2.count(), 1)
            # Test compressed command
            cdf2.createOrReplaceTempView("temp_view_cdf2")
            self.assertEqual(self.connect.sql("SELECT * FROM temp_view_cdf2").count(), 1)
        finally:
            self.connect._client._plan_compression_threshold = default_plan_compression_threshold


class SparkConnectGCTests(SparkConnectSQLTestCase):
    @classmethod
    def setUpClass(cls):
        cls.origin = os.getenv("USER", None)
        os.environ["USER"] = "SparkConnectGCTests"
        super().setUpClass()

    @classmethod
    def tearDownClass(cls):
        super().tearDownClass()
        if cls.origin is not None:
            os.environ["USER"] = cls.origin
        else:
            del os.environ["USER"]

    def test_garbage_collection_checkpoint(self):
        # SPARK-48258: Make sure garbage-collecting DataFrame remove the paired state
        # in Spark Connect server
        df = self.connect.range(10).localCheckpoint()
        self.assertIsNotNone(df._plan._relation_id)
        cached_remote_relation_id = df._plan._relation_id

        jvm = self.spark._jvm
        session_holder = getattr(
            getattr(
                jvm.org.apache.spark.sql.connect.service,
                "SparkConnectService$",
            ),
            "MODULE$",
        ).getOrCreateIsolatedSession(self.connect.client._user_id, self.connect.client._session_id)

        # Check the state exists.
        self.assertIsNotNone(
            session_holder.dataFrameCache().getOrDefault(cached_remote_relation_id, None)
        )

        del df
        gc.collect()

        def condition():
            # Check the state was removed up on garbage-collection.
            self.assertIsNone(
                session_holder.dataFrameCache().getOrDefault(cached_remote_relation_id, None)
            )

        eventually(catch_assertions=True)(condition)()

    def test_garbage_collection_derived_checkpoint(self):
        # SPARK-48258: Should keep the cached remote relation when derived DataFrames exist
        df = self.connect.range(10).localCheckpoint()
        self.assertIsNotNone(df._plan._relation_id)
        derived = df.repartition(10)
        cached_remote_relation_id = df._plan._relation_id

        jvm = self.spark._jvm
        session_holder = getattr(
            getattr(
                jvm.org.apache.spark.sql.connect.service,
                "SparkConnectService$",
            ),
            "MODULE$",
        ).getOrCreateIsolatedSession(self.connect.client._user_id, self.connect.client._session_id)

        # Check the state exists.
        self.assertIsNotNone(
            session_holder.dataFrameCache().getOrDefault(cached_remote_relation_id, None)
        )

        del df
        gc.collect()

        def condition():
            self.assertIsNone(
                session_holder.dataFrameCache().getOrDefault(cached_remote_relation_id, None)
            )

        # Should not remove the cache
        with self.assertRaises(AssertionError):
            eventually(catch_assertions=True, timeout=5)(condition)()

        del derived
        gc.collect()

        eventually(catch_assertions=True)(condition)()

    def test_arrow_batch_result_chunking(self):
        # Two cases are tested here:
        # (a) client preferred chunk size is set: the server should respect it
        # (b) client preferred chunk size is not set: the server should use its own max chunk size
        for preferred_chunk_size_optional, max_chunk_size_optional in ((1024, None), (None, 1024)):
            sql_query = "select id, CAST(id + 0.5 AS DOUBLE) from range(0, 2000, 1, 4)"
            cdf = self.connect.sql(sql_query)
            sdf = self.spark.sql(sql_query)

            original_verify_response_integrity = self.connect._client._verify_response_integrity
            captured_chunks = []

            def patched_verify_response_integrity(response):
                original_verify_response_integrity(response)
                if isinstance(response, ExecutePlanResponse) and response.HasField("arrow_batch"):
                    captured_chunks.append(response.arrow_batch)

            try:
                # Patch the response verifier for testing to access the chunked arrow batch
                # responses.
                self.connect._client._verify_response_integrity = patched_verify_response_integrity
                # Override the chunk size to 1024 bytes for testing
                if preferred_chunk_size_optional:
                    self.connect._client._preferred_arrow_chunk_size = preferred_chunk_size_optional
                if max_chunk_size_optional:
                    self.connect.conf.set(
                        "spark.connect.session.resultChunking.maxChunkSize",
                        max_chunk_size_optional,
                    )

                # Execute the query, and assert the results are correct.
                assertDataFrameEqual(cdf, sdf)

                # Verify the metadata of arrow batch chunks.
                def split_into_batches(chunks):
                    batches = []
                    i = 0
                    n = len(chunks)
                    while i < n:
                        num_chunks = chunks[i].num_chunks_in_batch
                        batch = chunks[i : i + num_chunks]
                        batches.append(batch)
                        i += num_chunks
                    return batches

                batches = split_into_batches(captured_chunks)
                # There are 4 batches (partitions) in total.
                self.assertEqual(len(batches), 4)
                for batch in batches:
                    # In this example, the max chunk size is set to a small value, so each Arrow
                    # batch should be split into multiple chunks.
                    self.assertTrue(len(batch) > 5)
                    row_count = batch[0].row_count
                    row_start_offset = batch[0].start_offset
                    for i, chunk in enumerate(batch):
                        self.assertEqual(chunk.chunk_index, i)
                        self.assertEqual(chunk.num_chunks_in_batch, len(batch))
                        self.assertEqual(chunk.row_count, row_count)
                        self.assertEqual(chunk.start_offset, row_start_offset)
                        self.assertTrue(len(chunk.data) > 0)
                        self.assertTrue(
                            len(chunk.data)
                            <= (preferred_chunk_size_optional or max_chunk_size_optional)
                        )
            finally:
                self.connect._client._verify_response_integrity = original_verify_response_integrity
                self.connect._client._preferred_arrow_chunk_size = None
                self.connect.conf.unset("spark.connect.session.resultChunking.maxChunkSize")


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

    main()