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

python/pyspark/sql/tests/connect/test_connect_error.py
289 строк11 KB
Hyukjin Kwon
[SPARK-57989][CONNECT][PYTHON] Raise NOT_IMPLEMENTED for year-month interval in Connect collect, matching Classic's default
07 июл 2026, 10:49
07 июл 2026, 10:4926660a9
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import unittest

from pyspark.errors import PySparkAttributeError
from pyspark.errors.exceptions.base import SessionNotSameException
from pyspark.sql.types import Row
from pyspark.sql import functions as F
from pyspark.errors import PySparkNotImplementedError, PySparkTypeError
from pyspark.testing.connectutils import ReusedConnectTestCase
from pyspark.util import is_remote_only


class SparkConnectErrorTests(ReusedConnectTestCase):
    def test_recursion_handling_for_plan_logging(self):
        """SPARK-45852 - Test that we can handle recursion in plan logging."""
        cdf = self.spark.range(1)
        for x in range(400):
            cdf = cdf.withColumn(f"col_{x}", F.lit(x))

        # Calling schema will trigger logging the message that will in turn trigger the message
        # conversion into protobuf that will then trigger the recursion error.
        self.assertIsNotNone(cdf.schema)

        result = self.spark._client._proto_to_string(cdf._plan.to_proto(self.spark._client))
        self.assertIn("recursion", result)

    def test_error_handling(self):
        from pyspark.errors.exceptions.connect import AnalysisException

        # SPARK-41533 Proper error handling for Spark Connect
        df = self.spark.range(10).select("id2")
        with self.assertRaises(AnalysisException):
            df.collect()

    def test_invalid_column(self):
        from pyspark.errors.exceptions.connect import AnalysisException

        # SPARK-41812: fail df1.select(df2.col)
        data1 = [Row(a=1, b=2, c=3)]
        cdf1 = self.spark.createDataFrame(data1)

        data2 = [Row(a=2, b=0)]
        cdf2 = self.spark.createDataFrame(data2)

        with self.assertRaises(AnalysisException):
            cdf1.select(cdf2.a).schema

        with self.assertRaises(AnalysisException):
            cdf2.withColumn("x", cdf1.a + 1).schema

        # Can find the target plan node, but the resolved attribute is not part of the outer
        # operator's output (e.g. it was dropped by a projection).
        with self.assertRaisesRegex(
            AnalysisException,
            "CANNOT_RESOLVE_DATAFRAME_COLUMN",
        ):
            cdf3 = cdf1.select(cdf1.a)
            cdf3.select(cdf1.b).schema

        # Same pattern: the column from a different DataFrame is overwritten by `withColumn`.
        with self.assertRaisesRegex(
            AnalysisException,
            "CANNOT_RESOLVE_DATAFRAME_COLUMN",
        ):
            cdf3 = cdf1.withColumn("a", F.lit(0))
            cdf3.select(cdf1.a).schema

        # SPARK-56547: `df["id"]` references the original `id`, but `df2`'s `id` was shadowed
        # by `withColumn`. Previously surfaced a misleading
        # `UNRESOLVED_COLUMN.WITH_SUGGESTION` error suggesting the same column name.
        df = self.spark.range(10)
        df2 = df.withColumn("id", F.col("id") + 1)
        with self.assertRaisesRegex(
            AnalysisException,
            "CANNOT_RESOLVE_DATAFRAME_COLUMN",
        ):
            df2.select(df["id"]).schema

        # Can not find the target plan node by plan id
        with self.assertRaisesRegex(
            AnalysisException,
            "CANNOT_RESOLVE_DATAFRAME_COLUMN",
        ):
            cdf1.select(cdf2.a).schema

    def test_invalid_star(self):
        from pyspark.errors.exceptions.connect import AnalysisException

        data1 = [Row(a=1, b=2, c=3)]
        cdf1 = self.spark.createDataFrame(data1)

        data2 = [Row(a=2, b=0)]
        cdf2 = self.spark.createDataFrame(data2)

        # Can find the target plan node, but fail to resolve with it
        with self.assertRaisesRegex(
            AnalysisException,
            "CANNOT_RESOLVE_DATAFRAME_COLUMN",
        ):
            cdf3 = cdf1.select(cdf1.a)
            cdf3.select(cdf1["*"]).schema

        # Can find the target plan node, but fail to resolve with it
        with self.assertRaisesRegex(
            AnalysisException,
            "CANNOT_RESOLVE_DATAFRAME_COLUMN",
        ):
            # column 'a has been replaced
            cdf3 = cdf1.withColumn("a", F.lit(0))
            cdf3.select(cdf1["*"]).schema

        # Can not find the target plan node by plan id
        with self.assertRaisesRegex(
            AnalysisException,
            "CANNOT_RESOLVE_DATAFRAME_COLUMN",
        ):
            cdf1.select(cdf2["*"]).schema

        # cdf1["*"] exists on both side
        with self.assertRaisesRegex(
            AnalysisException,
            "AMBIGUOUS_COLUMN_REFERENCE",
        ):
            cdf1.join(cdf1).select(cdf1["*"]).schema

    def test_deduplicate_within_watermark_in_batch(self):
        from pyspark.errors.exceptions.connect import AnalysisException

        table_name = "tmp_table_for_test_deduplicate_within_watermark_in_batch"
        with self.table(table_name):
            self.spark.createDataFrame(
                [Row(key=i, value=str(i)) for i in range(100)]
            ).write.saveAsTable(table_name)

            with self.assertRaisesRegex(
                AnalysisException,
                "dropDuplicatesWithinWatermark is not supported with batch DataFrames/DataSets",
            ):
                self.spark.read.table(table_name).dropDuplicatesWithinWatermark().toPandas()

    def test_different_spark_session_join_or_union(self):
        from pyspark.sql.connect.session import SparkSession as RemoteSparkSession

        df = self.spark.range(10).limit(3)

        spark2 = RemoteSparkSession(connection="sc://localhost")
        df2 = spark2.range(10).limit(3)

        with self.assertRaises(SessionNotSameException) as e1:
            df.union(df2).collect()
        self.check_error(
            exception=e1.exception,
            errorClass="SESSION_NOT_SAME",
            messageParameters={},
        )

        with self.assertRaises(SessionNotSameException) as e2:
            df.unionByName(df2).collect()
        self.check_error(
            exception=e2.exception,
            errorClass="SESSION_NOT_SAME",
            messageParameters={},
        )

        with self.assertRaises(SessionNotSameException) as e3:
            df.join(df2).collect()
        self.check_error(
            exception=e3.exception,
            errorClass="SESSION_NOT_SAME",
            messageParameters={},
        )

    @unittest.skipIf(is_remote_only(), "Disabled for remote only")
    def test_unsupported_functions(self):
        # SPARK-41225: Disable unsupported functions.
        df = self.spark.range(10)
        with self.assertRaises(NotImplementedError):
            df.rdd

    def test_unsupported_jvm_attribute(self):
        # Unsupported jvm attributes for Spark session.
        unsupported_attrs = ["_jsc", "_jconf", "_jvm", "_jsparkSession"]
        spark_session = self.spark
        for attr in unsupported_attrs:
            with self.assertRaises(PySparkAttributeError) as pe:
                getattr(spark_session, attr)

            self.check_error(
                exception=pe.exception,
                errorClass="JVM_ATTRIBUTE_NOT_SUPPORTED",
                messageParameters={"attr_name": attr},
            )

        # Unsupported jvm attributes for DataFrame.
        unsupported_attrs = ["_jseq", "_jdf", "_jmap", "_jcols"]
        cdf = self.spark.range(10)
        for attr in unsupported_attrs:
            with self.assertRaises(PySparkAttributeError) as pe:
                getattr(cdf, attr)

            self.check_error(
                exception=pe.exception,
                errorClass="JVM_ATTRIBUTE_NOT_SUPPORTED",
                messageParameters={"attr_name": attr},
            )

        # Unsupported jvm attributes for Column.
        with self.assertRaises(PySparkAttributeError) as pe:
            getattr(cdf.id, "_jc")

        self.check_error(
            exception=pe.exception,
            errorClass="JVM_ATTRIBUTE_NOT_SUPPORTED",
            messageParameters={"attr_name": "_jc"},
        )

        # Unsupported jvm attributes for DataFrameReader.
        with self.assertRaises(PySparkAttributeError) as pe:
            getattr(spark_session.read, "_jreader")

        self.check_error(
            exception=pe.exception,
            errorClass="JVM_ATTRIBUTE_NOT_SUPPORTED",
            messageParameters={"attr_name": "_jreader"},
        )

    def test_column_cannot_be_constructed_from_string(self):
        from pyspark.sql.connect.column import Column

        with self.assertRaises(TypeError):
            Column("col")

    def test_select_none(self):
        with self.assertRaises(PySparkTypeError) as e1:
            self.spark.range(1).select(None)

        self.check_error(
            exception=e1.exception,
            errorClass="NOT_EXPECTED_TYPE",
            messageParameters={
                "expected_type": "Column or str",
                "arg_name": "cols",
                "arg_type": "NoneType",
            },
        )

    def test_ym_interval_in_collect(self):
        # PyArrow cannot materialize Arrow YEAR_MONTH intervals, so collecting a year-month
        # interval over Spark Connect raises NOT_IMPLEMENTED (matching the default of the classic
        # PySpark path). Unlike classic, PYSPARK_YM_INTERVAL_LEGACY is not honored here.
        with self.assertRaises(PySparkNotImplementedError):
            self.spark.sql("SELECT INTERVAL '10-8' YEAR TO MONTH AS interval").first()

        # A year-month interval nested inside an array is rejected the same way (the schema-level
        # check recurses into array/map/struct element types).
        with self.assertRaises(PySparkNotImplementedError):
            self.spark.sql("SELECT array(INTERVAL '10-8' YEAR TO MONTH) AS interval").first()

    def test_ym_interval_empty_collect(self):
        # Even an empty result raises NOT_IMPLEMENTED rather than returning []. PyArrow cannot
        # build the INTERVAL_MONTHS array class at all -- `to_pylist()` raises `KeyError: 21` from
        # get_array_class_from_type regardless of row count -- so the schema-level check covers
        # empty results too, giving a clean error instead of an opaque PyArrow KeyError. This is
        # one place Spark Connect diverges from classic PySpark, which returns [] for an empty
        # result (it never reaches PyArrow materialization).
        with self.assertRaises(PySparkNotImplementedError):
            self.spark.sql("SELECT INTERVAL '10-8' YEAR TO MONTH AS interval").limit(0).collect()


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

    main()