/
githubmirror
/
spark
Обзор
Документация
Войти
/
githubmirror
/
spark
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
master
python/pyspark/sql/tests/test_dataframe_query_context.py
518 строк
22 KB
Tian Gao
[SPARK-55612][PYTHON][TESTS] Add test_dataframe_query_context to modules
20 фев 2026, 07:26
20 фев 2026, 07:26
434aa7a
Код
Авторство
О чём код?
# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from pyspark.errors import ( AnalysisException, ArithmeticException, QueryContextType, NumberFormatException, ) from pyspark.sql import functions as sf from pyspark.testing.sqlutils import ( ReusedSQLTestCase, ) class DataFrameQueryContextTestsMixin: def test_dataframe_query_context(self): # SPARK-47274: Add more useful contexts for PySpark DataFrame API errors. with self.sql_conf({"spark.sql.ansi.enabled": True}): df = self.spark.range(10) # DataFrameQueryContext with pysparkLoggingInfo - divide with self.assertRaises(ArithmeticException) as pe: df.withColumn("div_zero", df.id / 0).collect() self.check_error( exception=pe.exception, errorClass="DIVIDE_BY_ZERO", messageParameters={"config": '"spark.sql.ansi.enabled"'}, query_context_type=QueryContextType.DataFrame, fragment="__truediv__", ) # DataFrameQueryContext with pysparkLoggingInfo - plus with self.assertRaises(NumberFormatException) as pe: df.withColumn("plus_invalid_type", df.id + "string").collect() self.check_error( exception=pe.exception, errorClass="CAST_INVALID_INPUT", messageParameters={ "ansiConfig": '"spark.sql.ansi.enabled"', "expression": "'string'", "sourceType": '"STRING"', "targetType": '"BIGINT"', }, query_context_type=QueryContextType.DataFrame, fragment="__add__", ) # DataFrameQueryContext with pysparkLoggingInfo - minus with self.assertRaises(NumberFormatException) as pe: df.withColumn("minus_invalid_type", df.id - "string").collect() self.check_error( exception=pe.exception, errorClass="CAST_INVALID_INPUT", messageParameters={ "ansiConfig": '"spark.sql.ansi.enabled"', "expression": "'string'", "sourceType": '"STRING"', "targetType": '"BIGINT"', }, query_context_type=QueryContextType.DataFrame, fragment="__sub__", ) # DataFrameQueryContext with pysparkLoggingInfo - multiply with self.assertRaises(NumberFormatException) as pe: df.withColumn("multiply_invalid_type", df.id * "string").collect() self.check_error( exception=pe.exception, errorClass="CAST_INVALID_INPUT", messageParameters={ "ansiConfig": '"spark.sql.ansi.enabled"', "expression": "'string'", "sourceType": '"STRING"', "targetType": '"BIGINT"', }, query_context_type=QueryContextType.DataFrame, fragment="__mul__", ) # DataFrameQueryContext with pysparkLoggingInfo - mod with self.assertRaises(NumberFormatException) as pe: df.withColumn("mod_invalid_type", df.id % "string").collect() self.check_error( exception=pe.exception, errorClass="CAST_INVALID_INPUT", messageParameters={ "ansiConfig": '"spark.sql.ansi.enabled"', "expression": "'string'", "sourceType": '"STRING"', "targetType": '"BIGINT"', }, query_context_type=QueryContextType.DataFrame, fragment="__mod__", ) # DataFrameQueryContext with pysparkLoggingInfo - equalTo with self.assertRaises(NumberFormatException) as pe: df.withColumn("equalTo_invalid_type", df.id == "string").collect() self.check_error( exception=pe.exception, errorClass="CAST_INVALID_INPUT", messageParameters={ "ansiConfig": '"spark.sql.ansi.enabled"', "expression": "'string'", "sourceType": '"STRING"', "targetType": '"BIGINT"', }, query_context_type=QueryContextType.DataFrame, fragment="__eq__", ) # DataFrameQueryContext with pysparkLoggingInfo - lt with self.assertRaises(NumberFormatException) as pe: df.withColumn("lt_invalid_type", df.id < "string").collect() self.check_error( exception=pe.exception, errorClass="CAST_INVALID_INPUT", messageParameters={ "ansiConfig": '"spark.sql.ansi.enabled"', "expression": "'string'", "sourceType": '"STRING"', "targetType": '"BIGINT"', }, query_context_type=QueryContextType.DataFrame, fragment="__lt__", ) # DataFrameQueryContext with pysparkLoggingInfo - leq with self.assertRaises(NumberFormatException) as pe: df.withColumn("leq_invalid_type", df.id <= "string").collect() self.check_error( exception=pe.exception, errorClass="CAST_INVALID_INPUT", messageParameters={ "ansiConfig": '"spark.sql.ansi.enabled"', "expression": "'string'", "sourceType": '"STRING"', "targetType": '"BIGINT"', }, query_context_type=QueryContextType.DataFrame, fragment="__le__", ) # DataFrameQueryContext with pysparkLoggingInfo - geq with self.assertRaises(NumberFormatException) as pe: df.withColumn("geq_invalid_type", df.id >= "string").collect() self.check_error( exception=pe.exception, errorClass="CAST_INVALID_INPUT", messageParameters={ "ansiConfig": '"spark.sql.ansi.enabled"', "expression": "'string'", "sourceType": '"STRING"', "targetType": '"BIGINT"', }, query_context_type=QueryContextType.DataFrame, fragment="__ge__", ) # DataFrameQueryContext with pysparkLoggingInfo - gt with self.assertRaises(NumberFormatException) as pe: df.withColumn("gt_invalid_type", df.id > "string").collect() self.check_error( exception=pe.exception, errorClass="CAST_INVALID_INPUT", messageParameters={ "ansiConfig": '"spark.sql.ansi.enabled"', "expression": "'string'", "sourceType": '"STRING"', "targetType": '"BIGINT"', }, query_context_type=QueryContextType.DataFrame, fragment="__gt__", ) # DataFrameQueryContext with pysparkLoggingInfo - eqNullSafe with self.assertRaises(NumberFormatException) as pe: df.withColumn("eqNullSafe_invalid_type", df.id.eqNullSafe("string")).collect() self.check_error( exception=pe.exception, errorClass="CAST_INVALID_INPUT", messageParameters={ "ansiConfig": '"spark.sql.ansi.enabled"', "expression": "'string'", "sourceType": '"STRING"', "targetType": '"BIGINT"', }, query_context_type=QueryContextType.DataFrame, fragment="eqNullSafe", ) # DataFrameQueryContext with pysparkLoggingInfo - bitwiseOR with self.assertRaises(NumberFormatException) as pe: df.withColumn("bitwiseOR_invalid_type", df.id.bitwiseOR("string")).collect() self.check_error( exception=pe.exception, errorClass="CAST_INVALID_INPUT", messageParameters={ "ansiConfig": '"spark.sql.ansi.enabled"', "expression": "'string'", "sourceType": '"STRING"', "targetType": '"BIGINT"', }, query_context_type=QueryContextType.DataFrame, fragment="bitwiseOR", ) # DataFrameQueryContext with pysparkLoggingInfo - bitwiseAND with self.assertRaises(NumberFormatException) as pe: df.withColumn("bitwiseAND_invalid_type", df.id.bitwiseAND("string")).collect() self.check_error( exception=pe.exception, errorClass="CAST_INVALID_INPUT", messageParameters={ "ansiConfig": '"spark.sql.ansi.enabled"', "expression": "'string'", "sourceType": '"STRING"', "targetType": '"BIGINT"', }, query_context_type=QueryContextType.DataFrame, fragment="bitwiseAND", ) # DataFrameQueryContext with pysparkLoggingInfo - bitwiseXOR with self.assertRaises(NumberFormatException) as pe: df.withColumn("bitwiseXOR_invalid_type", df.id.bitwiseXOR("string")).collect() self.check_error( exception=pe.exception, errorClass="CAST_INVALID_INPUT", messageParameters={ "ansiConfig": '"spark.sql.ansi.enabled"', "expression": "'string'", "sourceType": '"STRING"', "targetType": '"BIGINT"', }, query_context_type=QueryContextType.DataFrame, fragment="bitwiseXOR", ) # DataFrameQueryContext with pysparkLoggingInfo - chained (`divide` is problematic) with self.assertRaises(ArithmeticException) as pe: df.withColumn("multiply_ten", df.id * 10).withColumn( "divide_zero", df.id / 0 ).withColumn("plus_ten", df.id + 10).withColumn("minus_ten", df.id - 10).collect() self.check_error( exception=pe.exception, errorClass="DIVIDE_BY_ZERO", messageParameters={"config": '"spark.sql.ansi.enabled"'}, query_context_type=QueryContextType.DataFrame, fragment="__truediv__", ) # DataFrameQueryContext with pysparkLoggingInfo - chained (`plus` is problematic) with self.assertRaises(NumberFormatException) as pe: df.withColumn("multiply_ten", df.id * 10).withColumn( "divide_ten", df.id / 10 ).withColumn("plus_string", df.id + "string").withColumn( "minus_ten", df.id - 10 ).collect() self.check_error( exception=pe.exception, errorClass="CAST_INVALID_INPUT", messageParameters={ "ansiConfig": '"spark.sql.ansi.enabled"', "expression": "'string'", "sourceType": '"STRING"', "targetType": '"BIGINT"', }, query_context_type=QueryContextType.DataFrame, fragment="__add__", ) # DataFrameQueryContext with pysparkLoggingInfo - chained (`minus` is problematic) with self.assertRaises(NumberFormatException) as pe: df.withColumn("multiply_ten", df.id * 10).withColumn( "divide_ten", df.id / 10 ).withColumn("plus_ten", df.id + 10).withColumn( "minus_string", df.id - "string" ).collect() self.check_error( exception=pe.exception, errorClass="CAST_INVALID_INPUT", messageParameters={ "ansiConfig": '"spark.sql.ansi.enabled"', "expression": "'string'", "sourceType": '"STRING"', "targetType": '"BIGINT"', }, query_context_type=QueryContextType.DataFrame, fragment="__sub__", ) # DataFrameQueryContext with pysparkLoggingInfo - chained (`multiply` is problematic) with self.assertRaises(NumberFormatException) as pe: df.withColumn("multiply_string", df.id * "string").withColumn( "divide_ten", df.id / 10 ).withColumn("plus_ten", df.id + 10).withColumn("minus_ten", df.id - 10).collect() self.check_error( exception=pe.exception, errorClass="CAST_INVALID_INPUT", messageParameters={ "ansiConfig": '"spark.sql.ansi.enabled"', "expression": "'string'", "sourceType": '"STRING"', "targetType": '"BIGINT"', }, query_context_type=QueryContextType.DataFrame, fragment="__mul__", ) # Multiple expressions in df.select (`divide` is problematic) with self.assertRaises(ArithmeticException) as pe: df.select(df.id - 10, df.id + 4, df.id / 0, df.id * 5).collect() self.check_error( exception=pe.exception, errorClass="DIVIDE_BY_ZERO", messageParameters={"config": '"spark.sql.ansi.enabled"'}, query_context_type=QueryContextType.DataFrame, fragment="__truediv__", ) # Multiple expressions in df.select (`plus` is problematic) with self.assertRaises(NumberFormatException) as pe: df.select(df.id - 10, df.id + "string", df.id / 10, df.id * 5).collect() self.check_error( exception=pe.exception, errorClass="CAST_INVALID_INPUT", messageParameters={ "ansiConfig": '"spark.sql.ansi.enabled"', "expression": "'string'", "sourceType": '"STRING"', "targetType": '"BIGINT"', }, query_context_type=QueryContextType.DataFrame, fragment="__add__", ) # Multiple expressions in df.select (`minus` is problematic) with self.assertRaises(NumberFormatException) as pe: df.select(df.id - "string", df.id + 4, df.id / 10, df.id * 5).collect() self.check_error( exception=pe.exception, errorClass="CAST_INVALID_INPUT", messageParameters={ "ansiConfig": '"spark.sql.ansi.enabled"', "expression": "'string'", "sourceType": '"STRING"', "targetType": '"BIGINT"', }, query_context_type=QueryContextType.DataFrame, fragment="__sub__", ) # Multiple expressions in df.select (`multiply` is problematic) with self.assertRaises(NumberFormatException) as pe: df.select(df.id - 10, df.id + 4, df.id / 10, df.id * "string").collect() self.check_error( exception=pe.exception, errorClass="CAST_INVALID_INPUT", messageParameters={ "ansiConfig": '"spark.sql.ansi.enabled"', "expression": "'string'", "sourceType": '"STRING"', "targetType": '"BIGINT"', }, query_context_type=QueryContextType.DataFrame, fragment="__mul__", ) # Multiple expressions with pre-declared expressions (`divide` is problematic) a = df.id / 10 b = df.id / 0 with self.assertRaises(ArithmeticException) as pe: df.select(a, df.id + 4, b, df.id * 5).collect() self.check_error( exception=pe.exception, errorClass="DIVIDE_BY_ZERO", messageParameters={"config": '"spark.sql.ansi.enabled"'}, query_context_type=QueryContextType.DataFrame, fragment="__truediv__", ) # Multiple expressions with pre-declared expressions (`plus` is problematic) a = df.id + "string" b = df.id + 4 with self.assertRaises(NumberFormatException) as pe: df.select(df.id / 10, a, b, df.id * 5).collect() self.check_error( exception=pe.exception, errorClass="CAST_INVALID_INPUT", messageParameters={ "ansiConfig": '"spark.sql.ansi.enabled"', "expression": "'string'", "sourceType": '"STRING"', "targetType": '"BIGINT"', }, query_context_type=QueryContextType.DataFrame, fragment="__add__", ) # Multiple expressions with pre-declared expressions (`minus` is problematic) a = df.id - "string" b = df.id - 5 with self.assertRaises(NumberFormatException) as pe: df.select(a, df.id / 10, b, df.id * 5).collect() self.check_error( exception=pe.exception, errorClass="CAST_INVALID_INPUT", messageParameters={ "ansiConfig": '"spark.sql.ansi.enabled"', "expression": "'string'", "sourceType": '"STRING"', "targetType": '"BIGINT"', }, query_context_type=QueryContextType.DataFrame, fragment="__sub__", ) # Multiple expressions with pre-declared expressions (`multiply` is problematic) a = df.id * "string" b = df.id * 10 with self.assertRaises(NumberFormatException) as pe: df.select(a, df.id / 10, b, df.id + 5).collect() self.check_error( exception=pe.exception, errorClass="CAST_INVALID_INPUT", messageParameters={ "ansiConfig": '"spark.sql.ansi.enabled"', "expression": "'string'", "sourceType": '"STRING"', "targetType": '"BIGINT"', }, query_context_type=QueryContextType.DataFrame, fragment="__mul__", ) def test_sql_query_context(self): with self.sql_conf({"spark.sql.ansi.enabled": True}): # SQLQueryContext with self.assertRaises(ArithmeticException) as pe: self.spark.sql("select 10/0").collect() self.check_error( exception=pe.exception, errorClass="DIVIDE_BY_ZERO", messageParameters={"config": '"spark.sql.ansi.enabled"'}, query_context_type=QueryContextType.SQL, ) # No QueryContext with self.assertRaises(AnalysisException) as pe: self.spark.sql("select * from non-existing-table") self.check_error( exception=pe.exception, errorClass="INVALID_IDENTIFIER", messageParameters={"ident": "non-existing-table"}, query_context_type=None, ) def test_query_context_complex(self): with self.sql_conf({"spark.sql.ansi.enabled": True}): # SQLQueryContext with self.assertRaises(ArithmeticException) as pe: self.spark.sql("select (10/0)*100").collect() self.check_error( exception=pe.exception, errorClass="DIVIDE_BY_ZERO", messageParameters={"config": '"spark.sql.ansi.enabled"'}, query_context_type=QueryContextType.SQL, ) # DataFrameQueryContext df = self.spark.range(10) with self.assertRaises(ArithmeticException) as pe: df.withColumn("div_zero", (df.id / 0) * 10).collect() self.check_error( exception=pe.exception, errorClass="DIVIDE_BY_ZERO", messageParameters={"config": '"spark.sql.ansi.enabled"'}, query_context_type=QueryContextType.DataFrame, fragment="__truediv__", ) def test_dataframe_query_context_col(self): with self.assertRaises(AnalysisException) as pe: self.spark.range(1).select(sf.col("id") + sf.col("idd")).show() self.check_error( exception=pe.exception, errorClass="UNRESOLVED_COLUMN.WITH_SUGGESTION", messageParameters={"objectName": "`idd`", "proposal": "`id`"}, query_context_type=QueryContextType.DataFrame, fragment="col", ) class DataFrameQueryContextTests(DataFrameQueryContextTestsMixin, ReusedSQLTestCase): pass if __name__ == "__main__": from pyspark.testing import main main()