diff --git a/dev/.rat-excludes b/dev/.rat-excludes index 81f05cd011754..4aee210e85b98 100644 --- a/dev/.rat-excludes +++ b/dev/.rat-excludes @@ -141,6 +141,13 @@ ui-test/package-lock.json core/src/main/resources/org/apache/spark/ui/static/package.json testCommitLog .*\.har +.*\.test spark-logo.svg spark-logo-rev.svg .nojekyll +# PySpark DataFrame golden test fixture scripts +# (python/pyspark/sql/tests/df_golden/scripts//). These are executable +# test fixtures, not source, so they are excluded like the SQL golden inputs +# (see .*\.sql above). Add one directory name per golden group below. RAT matches +# path components, so each name must be unique across the whole repository. +group_by diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py index 33049885102ce..f8c1e0a9b0d18 100644 --- a/dev/sparktestsupport/modules.py +++ b/dev/sparktestsupport/modules.py @@ -637,6 +637,8 @@ def __hash__(self): "pyspark.sql.tests.coercion.test_python_udf_input_type", "pyspark.sql.tests.coercion.test_pandas_udf_return_type", "pyspark.sql.tests.coercion.test_python_udf_return_type", + "pyspark.sql.tests.df_golden.test_df_golden", + "pyspark.sql.tests.df_golden.test_df_golden_framework", ], ) diff --git a/pyproject.toml b/pyproject.toml index a28533c3539b1..c5dc0281383e1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,6 +59,9 @@ ignore = [ "python/pyspark/errors/error_classes.py" = ["E501"] # Examples contain some unused variables. "examples/src/main/python/sql/datasource.py" = ["F841"] + # DataFrame golden test scripts are executed by the framework with ``spark`` + # injected into their namespace, so it is undefined at lint time. + "python/pyspark/sql/tests/df_golden/scripts/**/*.py" = ["F821"] [tool.black] # When changing the version, we have to update diff --git a/python/MANIFEST.in b/python/MANIFEST.in index 82979a0344c3c..4901fe9379668 100644 --- a/python/MANIFEST.in +++ b/python/MANIFEST.in @@ -17,6 +17,7 @@ # Reference: https://setuptools.pypa.io/en/latest/userguide/miscellaneous.html recursive-include pyspark *.pyi py.typed *.json +recursive-include pyspark/sql/tests/df_golden *.test *.py recursive-include deps/jars *.jar graft deps/bin recursive-include deps/sbin spark-config.sh spark-daemon.sh start-history-server.sh stop-history-server.sh diff --git a/python/pyspark/sql/tests/df_golden/__init__.py b/python/pyspark/sql/tests/df_golden/__init__.py new file mode 100644 index 0000000000000..cce3acad34a49 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/__init__.py @@ -0,0 +1,16 @@ +# +# 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. +# diff --git a/python/pyspark/sql/tests/df_golden/group_by.test b/python/pyspark/sql/tests/df_golden/group_by.test new file mode 100644 index 0000000000000..4fe132fd5ebf7 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/group_by.test @@ -0,0 +1,1106 @@ +--! name +__file_metadata__ +--! source +df_golden/group_by +!-- end + + +--! name +create testData view +--! script +scripts/group_by/setup_test_data.py +--! expected_analysis_output +LocalRelation +--! expected_optimized_output +LocalRelation +--! expected_output_schema +struct<> +--! expected_result +printed all 0 rows. +--! expected_result_hash +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 +!-- end + + +--! name +select non-grouping column without GROUP BY (error) +--! script +scripts/group_by/select_non_grouping_column.py +--! expected_error +[MISSING_GROUP_BY] The query does not include a GROUP BY clause. Add GROUP BY or turn it into the window functions using OVER clauses. SQLSTATE: 42803 +!-- end + + +--! name +aggregate with empty GroupBy expressions +--! script +scripts/group_by/agg_counts_no_grouping.py +--! expected_analysis_output +Aggregate [count(a#x) AS count(a)#xL, count(b#x) AS count(b)#xL] ++- SubqueryAlias testdata + +- View (`testData`, [a#x, b#x]) + +- Project [cast(a#x as int) AS a#x, cast(b#x as int) AS b#x] + +- Project [a#x, b#x] + +- SubqueryAlias testData + +- LocalRelation [a#x, b#x] +--! expected_optimized_output +Aggregate [count(a#x) AS count(a)#xL, count(b#x) AS count(b)#xL] ++- LocalRelation [a#x, b#x] +--! expected_output_schema +struct +--! expected_result ++----------+----------+ +| count(a) | count(b) | ++----------+----------+ +| 7 | 7 | ++----------+----------+ +printed all 1 rows. +--! expected_result_hash +f68f701c7611c45f2e32de9394884c63089bd8174e53a863a1b61168442dbaac +!-- end + + +--! name +aggregate with non-empty GroupBy expressions +--! tags +unordered +--! script +scripts/group_by/group_by_count.py +--! expected_analysis_output +Aggregate [a#x], [a#x, count(b#x) AS count(b)#xL] ++- SubqueryAlias testdata + +- View (`testData`, [a#x, b#x]) + +- Project [cast(a#x as int) AS a#x, cast(b#x as int) AS b#x] + +- Project [a#x, b#x] + +- SubqueryAlias testData + +- LocalRelation [a#x, b#x] +--! expected_optimized_output +Aggregate [a#x], [a#x, count(b#x) AS count(b)#xL] ++- LocalRelation [a#x, b#x] +--! expected_output_schema +struct +--! expected_result ++------+----------+ +| a | count(b) | ++------+----------+ +| 1 | 2 | +| 2 | 2 | +| 3 | 2 | +| NULL | 1 | ++------+----------+ +printed all 4 rows. +--! expected_result_hash +18c7a56d93a03d040c2ccb08307e9d9b38d789a0cd654823cbe7659cb182fcf9 +!-- end + + +--! name +non-aggregate column not in GROUP BY (error) +--! script +scripts/group_by/group_by_missing_aggregation.py +--! expected_error +[MISSING_AGGREGATION] The non-aggregating expression "a" is based on columns which are not participating in the GROUP BY clause. +Add the columns or the expression to the GROUP BY, aggregate the expression, or use "any_value(a)" if you do not care which of the values within a group is returned. SQLSTATE: 42803 +!-- end + + +--! name +group by with multiple aggregates +--! tags +unordered +--! script +scripts/group_by/group_by_multiple_counts.py +--! expected_analysis_output +Aggregate [a#x], [a#x, count(a#x) AS count(a)#xL, count(b#x) AS count(b)#xL] ++- SubqueryAlias testdata + +- View (`testData`, [a#x, b#x]) + +- Project [cast(a#x as int) AS a#x, cast(b#x as int) AS b#x] + +- Project [a#x, b#x] + +- SubqueryAlias testData + +- LocalRelation [a#x, b#x] +--! expected_optimized_output +Aggregate [a#x], [a#x, count(a#x) AS count(a)#xL, count(b#x) AS count(b)#xL] ++- LocalRelation [a#x, b#x] +--! expected_output_schema +struct +--! expected_result ++------+----------+----------+ +| a | count(a) | count(b) | ++------+----------+----------+ +| 1 | 2 | 2 | +| 2 | 2 | 2 | +| 3 | 3 | 2 | +| NULL | 0 | 1 | ++------+----------+----------+ +printed all 4 rows. +--! expected_result_hash +978157032220537c5982e8c4ceda91f71f8cc1159b71bf8023c71622b591dc9e +!-- end + + +--! name +aggregate grouped by literal +--! script +scripts/group_by/group_by_literal.py +--! expected_analysis_output +Aggregate [foo], [foo AS foo#x, count(a#x) AS count(a)#xL] ++- SubqueryAlias testdata + +- View (`testData`, [a#x, b#x]) + +- Project [cast(a#x as int) AS a#x, cast(b#x as int) AS b#x] + +- Project [a#x, b#x] + +- SubqueryAlias testData + +- LocalRelation [a#x, b#x] +--! expected_optimized_output +Aggregate [0], [foo AS foo#x, count(a#x) AS count(a)#xL] ++- LocalRelation [a#x] +--! expected_output_schema +struct +--! expected_result ++-----+----------+ +| foo | count(a) | ++-----+----------+ +| foo | 7 | ++-----+----------+ +printed all 1 rows. +--! expected_result_hash +56f7c8b96bd39b7def5a4b948f8551ccb63c76b3adbd7b7e2f3ea1e174312343 +!-- end + + +--! name +aggregate grouped by literal (hash aggregate) +--! script +scripts/group_by/group_by_literal_hash_agg.py +--! expected_analysis_output +Aggregate [foo], [foo AS foo#x, approx_count_distinct(a#x, 0.05, 0, 0) AS approx_count_distinct(a)#xL] ++- Filter (a#x = 0) + +- SubqueryAlias testdata + +- View (`testData`, [a#x, b#x]) + +- Project [cast(a#x as int) AS a#x, cast(b#x as int) AS b#x] + +- Project [a#x, b#x] + +- SubqueryAlias testData + +- LocalRelation [a#x, b#x] +--! expected_optimized_output +LocalRelation , [foo#x, approx_count_distinct(a)#xL] +--! expected_output_schema +struct +--! expected_result ++-----+--------------------------+ +| foo | approx_count_distinct(a) | ++-----+--------------------------+ ++-----+--------------------------+ +printed all 0 rows. +--! expected_result_hash +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 +!-- end + + +--! name +aggregate grouped by literal (sort aggregate) +--! script +scripts/group_by/group_by_literal_sort_agg.py +--! expected_analysis_output +Aggregate [foo], [foo AS foo#x, max(struct(a, a#x)) AS max(struct(a))#x] ++- Filter (a#x = 0) + +- SubqueryAlias testdata + +- View (`testData`, [a#x, b#x]) + +- Project [cast(a#x as int) AS a#x, cast(b#x as int) AS b#x] + +- Project [a#x, b#x] + +- SubqueryAlias testData + +- LocalRelation [a#x, b#x] +--! expected_optimized_output +LocalRelation , [foo#x, max(struct(a))#x] +--! expected_output_schema +struct> +--! expected_result ++-----+----------------+ +| foo | max(struct(a)) | ++-----+----------------+ ++-----+----------------+ +printed all 0 rows. +--! expected_result_hash +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 +!-- end + + +--! name +aggregate with complex GroupBy expression +--! tags +unordered +--! script +scripts/group_by/group_by_complex_expr.py +--! expected_analysis_output +Aggregate [(a#x + b#x)], [(a#x + b#x) AS (a + b)#x, count(b#x) AS count(b)#xL] ++- SubqueryAlias testdata + +- View (`testData`, [a#x, b#x]) + +- Project [cast(a#x as int) AS a#x, cast(b#x as int) AS b#x] + +- Project [a#x, b#x] + +- SubqueryAlias testData + +- LocalRelation [a#x, b#x] +--! expected_optimized_output +Aggregate [_groupingexpression#x], [_groupingexpression#x AS (a + b)#x, count(b#x) AS count(b)#xL] ++- LocalRelation [b#x, _groupingexpression#x] +--! expected_output_schema +struct<(a + b):int,count(b):bigint> +--! expected_result ++---------+----------+ +| (a + b) | count(b) | ++---------+----------+ +| 2 | 1 | +| 3 | 2 | +| 4 | 2 | +| 5 | 1 | +| NULL | 1 | ++---------+----------+ +printed all 5 rows. +--! expected_result_hash +e5a84602fefae835a3f9378c2cb5e6c8a653f258181aca8e9e88f1653b75823a +!-- end + + +--! name +struct() in group by +--! tags +unordered +--! script +scripts/group_by/group_by_struct.py +--! expected_analysis_output +Aggregate [struct(aa, (cast(a#x as double) + 0.1))], [struct(aa, (cast(a#x as double) + 0.1)) AS struct((a + 0.1) AS aa)#x, count(1) AS count(1)#xL] ++- SubqueryAlias testdata + +- View (`testData`, [a#x, b#x]) + +- Project [cast(a#x as int) AS a#x, cast(b#x as int) AS b#x] + +- Project [a#x, b#x] + +- SubqueryAlias testData + +- LocalRelation [a#x, b#x] +--! expected_optimized_output +Aggregate [_groupingexpression#x], [_groupingexpression#x AS struct((a + 0.1) AS aa)#x, count(1) AS count(1)#xL] ++- LocalRelation [_groupingexpression#x] +--! expected_output_schema +struct,count(1):bigint> +--! expected_result ++-------------------------+----------+ +| struct((a + 0.1) AS aa) | count(1) | ++-------------------------+----------+ +| {"aa":1.1} | 2 | +| {"aa":2.1} | 2 | +| {"aa":3.1} | 3 | +| {"aa":null} | 2 | ++-------------------------+----------+ +printed all 4 rows. +--! expected_result_hash +109ec364586f8c9f18e50baa49085b25107bad900acfe53183feff0968f4cc33 +!-- end + + +--! name +aggregate with nulls +--! script +scripts/group_by/agg_stats_with_nulls.py +--! expected_analysis_output +Aggregate [round(skewness(cast(a#x as double)), 12) AS round(skewness(a), 12)#x, round(kurtosis(cast(a#x as double)), 12) AS round(kurtosis(a), 12)#x, min(a#x) AS min(a)#x, max(a#x) AS max(a)#x, round(avg(a#x), 12) AS round(avg(a), 12)#x, round(variance(cast(a#x as double)), 12) AS round(variance(a), 12)#x, round(stddev(cast(a#x as double)), 12) AS round(stddev(a), 12)#x, sum(a#x) AS sum(a)#xL, count(a#x) AS count(a)#xL] ++- SubqueryAlias testdata + +- View (`testData`, [a#x, b#x]) + +- Project [cast(a#x as int) AS a#x, cast(b#x as int) AS b#x] + +- Project [a#x, b#x] + +- SubqueryAlias testData + +- LocalRelation [a#x, b#x] +--! expected_optimized_output +Aggregate [round(skewness(cast(a#x as double)), 12) AS round(skewness(a), 12)#x, round(kurtosis(cast(a#x as double)), 12) AS round(kurtosis(a), 12)#x, min(a#x) AS min(a)#x, max(a#x) AS max(a)#x, round(avg(a#x), 12) AS round(avg(a), 12)#x, round(variance(cast(a#x as double)), 12) AS round(variance(a), 12)#x, round(stddev(cast(a#x as double)), 12) AS round(stddev(a), 12)#x, sum(a#x) AS sum(a)#xL, count(a#x) AS count(a)#xL] ++- LocalRelation [a#x] +--! expected_output_schema +struct +--! expected_result ++------------------------+------------------------+--------+--------+-------------------+------------------------+----------------------+--------+----------+ +| round(skewness(a), 12) | round(kurtosis(a), 12) | min(a) | max(a) | round(avg(a), 12) | round(variance(a), 12) | round(stddev(a), 12) | sum(a) | count(a) | ++------------------------+------------------------+--------+--------+-------------------+------------------------+----------------------+--------+----------+ +| -0.272380105815 | -1.506920415225 | 1 | 3 | 2.142857142857 | 0.809523809524 | 0.899735410842 | 15 | 7 | ++------------------------+------------------------+--------+--------+-------------------+------------------------+----------------------+--------+----------+ +printed all 1 rows. +--! expected_result_hash +5b9fa3c2109f35b123ebd7b633649e02c09e73ed40574556e19a05e532838ab1 +!-- end + + +--! name +aggregate with empty input and non-empty GroupBy expressions +--! script +scripts/group_by/empty_input_group_by.py +--! expected_analysis_output +Aggregate [a#x], [a#x, count(1) AS count(1)#xL] ++- Filter false + +- SubqueryAlias testdata + +- View (`testData`, [a#x, b#x]) + +- Project [cast(a#x as int) AS a#x, cast(b#x as int) AS b#x] + +- Project [a#x, b#x] + +- SubqueryAlias testData + +- LocalRelation [a#x, b#x] +--! expected_optimized_output +LocalRelation , [a#x, count(1)#xL] +--! expected_output_schema +struct +--! expected_result ++---+----------+ +| a | count(1) | ++---+----------+ ++---+----------+ +printed all 0 rows. +--! expected_result_hash +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 +!-- end + + +--! name +aggregate with empty input and empty GroupBy expressions +--! script +scripts/group_by/empty_input_agg.py +--! expected_analysis_output +Aggregate [count(1) AS count(1)#xL] ++- Filter false + +- SubqueryAlias testdata + +- View (`testData`, [a#x, b#x]) + +- Project [cast(a#x as int) AS a#x, cast(b#x as int) AS b#x] + +- Project [a#x, b#x] + +- SubqueryAlias testData + +- LocalRelation [a#x, b#x] +--! expected_optimized_output +Aggregate [count(1) AS count(1)#xL] ++- LocalRelation +--! expected_output_schema +struct +--! expected_result ++----------+ +| count(1) | ++----------+ +| 0 | ++----------+ +printed all 1 rows. +--! expected_result_hash +5feceb66ffc86f38d952786c6d696c79c2dbc239dd4e91b46729d73a27fb57e9 +!-- end + + +--! name +aggregate with empty input and constant projection +--! script +scripts/group_by/empty_input_agg_select.py +--! expected_analysis_output +Project [1 AS 1#x] ++- Aggregate [count(1) AS count(1)#xL] + +- Filter false + +- SubqueryAlias testdata + +- View (`testData`, [a#x, b#x]) + +- Project [cast(a#x as int) AS a#x, cast(b#x as int) AS b#x] + +- Project [a#x, b#x] + +- SubqueryAlias testData + +- LocalRelation [a#x, b#x] +--! expected_optimized_output +Aggregate [1 AS 1#x] ++- LocalRelation +--! expected_output_schema +struct<1:int> +--! expected_result ++---+ +| 1 | ++---+ +| 1 | ++---+ +printed all 1 rows. +--! expected_result_hash +6b86b273ff34fce19d6b804eff5a3f5747ada4eaa22f1d49c01e52ddb7875b4b +!-- end + + +--! name +create test_agg view +--! script +scripts/group_by/setup_test_agg.py +--! expected_analysis_output +LocalRelation +--! expected_optimized_output +LocalRelation +--! expected_output_schema +struct<> +--! expected_result +printed all 0 rows. +--! expected_result_hash +e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 +!-- end + + +--! name +bool aggregates over empty table +--! script +scripts/group_by/bool_agg_empty_table.py +--! expected_analysis_output +Aggregate [every(v#x) AS every(v)#x, some(v#x) AS some(v)#x, bool_or(v#x) AS bool_or(v)#x, bool_and(v#x) AS bool_and(v)#x, bool_or(v#x) AS bool_or(v)#x] ++- Filter (1 = 0) + +- SubqueryAlias test_agg + +- View (`test_agg`, [k#x, v#x]) + +- Project [cast(k#x as int) AS k#x, cast(v#x as boolean) AS v#x] + +- Project [k#x, v#x] + +- SubqueryAlias test_agg + +- LocalRelation [k#x, v#x] +--! expected_optimized_output +Aggregate [min(v#x) AS every(v)#x, max(v#x) AS some(v)#x, max(v#x) AS bool_or(v)#x, min(v#x) AS bool_and(v)#x, max(v#x) AS bool_or(v)#x] ++- LocalRelation , [v#x] +--! expected_output_schema +struct +--! expected_result ++----------+---------+------------+-------------+------------+ +| every(v) | some(v) | bool_or(v) | bool_and(v) | bool_or(v) | ++----------+---------+------------+-------------+------------+ +| NULL | NULL | NULL | NULL | NULL | ++----------+---------+------------+-------------+------------+ +printed all 1 rows. +--! expected_result_hash +ad29e33ee1f3f4cfe263272e81d7c89ea08bdfb7aa45758a6adfc7752d80ab7e +!-- end + + +--! name +bool aggregates over all null values +--! script +scripts/group_by/bool_agg_all_nulls.py +--! expected_analysis_output +Aggregate [every(v#x) AS every(v)#x, some(v#x) AS some(v)#x, bool_or(v#x) AS bool_or(v)#x, bool_and(v#x) AS bool_and(v)#x, bool_or(v#x) AS bool_or(v)#x] ++- Filter (k#x = 4) + +- SubqueryAlias test_agg + +- View (`test_agg`, [k#x, v#x]) + +- Project [cast(k#x as int) AS k#x, cast(v#x as boolean) AS v#x] + +- Project [k#x, v#x] + +- SubqueryAlias test_agg + +- LocalRelation [k#x, v#x] +--! expected_optimized_output +Aggregate [min(v#x) AS every(v)#x, max(v#x) AS some(v)#x, max(v#x) AS bool_or(v)#x, min(v#x) AS bool_and(v)#x, max(v#x) AS bool_or(v)#x] ++- LocalRelation [v#x] +--! expected_output_schema +struct +--! expected_result ++----------+---------+------------+-------------+------------+ +| every(v) | some(v) | bool_or(v) | bool_and(v) | bool_or(v) | ++----------+---------+------------+-------------+------------+ +| NULL | NULL | NULL | NULL | NULL | ++----------+---------+------------+-------------+------------+ +printed all 1 rows. +--! expected_result_hash +ad29e33ee1f3f4cfe263272e81d7c89ea08bdfb7aa45758a6adfc7752d80ab7e +!-- end + + +--! name +bool aggregates null filtering +--! script +scripts/group_by/bool_agg_null_filtering.py +--! expected_analysis_output +Aggregate [every(v#x) AS every(v)#x, some(v#x) AS some(v)#x, bool_or(v#x) AS bool_or(v)#x, bool_and(v#x) AS bool_and(v)#x, bool_or(v#x) AS bool_or(v)#x] ++- Filter (k#x = 5) + +- SubqueryAlias test_agg + +- View (`test_agg`, [k#x, v#x]) + +- Project [cast(k#x as int) AS k#x, cast(v#x as boolean) AS v#x] + +- Project [k#x, v#x] + +- SubqueryAlias test_agg + +- LocalRelation [k#x, v#x] +--! expected_optimized_output +Aggregate [min(v#x) AS every(v)#x, max(v#x) AS some(v)#x, max(v#x) AS bool_or(v)#x, min(v#x) AS bool_and(v)#x, max(v#x) AS bool_or(v)#x] ++- LocalRelation [v#x] +--! expected_output_schema +struct +--! expected_result ++----------+---------+------------+-------------+------------+ +| every(v) | some(v) | bool_or(v) | bool_and(v) | bool_or(v) | ++----------+---------+------------+-------------+------------+ +| false | true | true | false | true | ++----------+---------+------------+-------------+------------+ +printed all 1 rows. +--! expected_result_hash +3b46bd9f792f9e10cf44f4797d0882661bf2ce33e75e52319637876a858c6eab +!-- end + + +--! name +bool aggregates with group by +--! tags +unordered +--! script +scripts/group_by/bool_agg_group_by.py +--! expected_analysis_output +Aggregate [k#x], [k#x, every(v#x) AS every(v)#x, some(v#x) AS some(v)#x, bool_or(v#x) AS bool_or(v)#x, bool_and(v#x) AS bool_and(v)#x, bool_or(v#x) AS bool_or(v)#x] ++- SubqueryAlias test_agg + +- View (`test_agg`, [k#x, v#x]) + +- Project [cast(k#x as int) AS k#x, cast(v#x as boolean) AS v#x] + +- Project [k#x, v#x] + +- SubqueryAlias test_agg + +- LocalRelation [k#x, v#x] +--! expected_optimized_output +Aggregate [k#x], [k#x, min(v#x) AS every(v)#x, max(v#x) AS some(v)#x, max(v#x) AS bool_or(v)#x, min(v#x) AS bool_and(v)#x, max(v#x) AS bool_or(v)#x] ++- LocalRelation [k#x, v#x] +--! expected_output_schema +struct +--! expected_result ++---+----------+---------+------------+-------------+------------+ +| k | every(v) | some(v) | bool_or(v) | bool_and(v) | bool_or(v) | ++---+----------+---------+------------+-------------+------------+ +| 1 | false | true | true | false | true | +| 2 | true | true | true | true | true | +| 3 | false | false | false | false | false | +| 4 | NULL | NULL | NULL | NULL | NULL | +| 5 | false | true | true | false | true | ++---+----------+---------+------------+-------------+------------+ +printed all 5 rows. +--! expected_result_hash +5dc9ed24b5544acb569b58286aec4265e8f921fe147a7ff12e5d456d4d5ec282 +!-- end + + +--! name +bool aggregates with having false +--! tags +unordered +--! script +scripts/group_by/bool_agg_having_false.py +--! expected_analysis_output +Project [k#x, every_v#x] ++- Filter (every_v#x = false) + +- Aggregate [k#x], [k#x, every(v#x) AS every_v#x] + +- SubqueryAlias test_agg + +- View (`test_agg`, [k#x, v#x]) + +- Project [cast(k#x as int) AS k#x, cast(v#x as boolean) AS v#x] + +- Project [k#x, v#x] + +- SubqueryAlias test_agg + +- LocalRelation [k#x, v#x] +--! expected_optimized_output +Filter (isnotnull(every_v#x) AND NOT every_v#x) ++- Aggregate [k#x], [k#x, min(v#x) AS every_v#x] + +- LocalRelation [k#x, v#x] +--! expected_output_schema +struct +--! expected_result ++---+---------+ +| k | every_v | ++---+---------+ +| 1 | false | +| 3 | false | +| 5 | false | ++---+---------+ +printed all 3 rows. +--! expected_result_hash +b38b91ef7aec43d0957f73891494f8c09519ffe4731fb729a0569c7c2d970118 +!-- end + + +--! name +bool aggregates with having null +--! script +scripts/group_by/bool_agg_having_null.py +--! expected_analysis_output +Project [k#x, every_v#x] ++- Filter isnull(every_v#x) + +- Aggregate [k#x], [k#x, every(v#x) AS every_v#x] + +- SubqueryAlias test_agg + +- View (`test_agg`, [k#x, v#x]) + +- Project [cast(k#x as int) AS k#x, cast(v#x as boolean) AS v#x] + +- Project [k#x, v#x] + +- SubqueryAlias test_agg + +- LocalRelation [k#x, v#x] +--! expected_optimized_output +Filter isnull(every_v#x) ++- Aggregate [k#x], [k#x, min(v#x) AS every_v#x] + +- LocalRelation [k#x, v#x] +--! expected_output_schema +struct +--! expected_result ++---+---------+ +| k | every_v | ++---+---------+ +| 4 | NULL | ++---+---------+ +printed all 1 rows. +--! expected_result_hash +3c3c8c2bd86b08b2f86cf8a0f0796bed2e3989c7d939ee2630e6dbace881b365 +!-- end + + +--! name +every input type checking: int (error) +--! script +scripts/group_by/every_int_error.py +--! expected_error +[DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE] Cannot resolve "every(1)" due to data type mismatch: The first parameter requires the "BOOLEAN" type, however "1" has the type "INT". SQLSTATE: 42K09 +!-- end + + +--! name +some input type checking: int (error) +--! script +scripts/group_by/some_int_error.py +--! expected_error +[DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE] Cannot resolve "some(1)" due to data type mismatch: The first parameter requires the "BOOLEAN" type, however "1" has the type "INT". SQLSTATE: 42K09 +!-- end + + +--! name +bool_or input type checking: int (error) +--! script +scripts/group_by/bool_or_int_error.py +--! expected_error +[DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE] Cannot resolve "bool_or(1)" due to data type mismatch: The first parameter requires the "BOOLEAN" type, however "1" has the type "INT". SQLSTATE: 42K09 +!-- end + + +--! name +every input type checking: string +--! script +scripts/group_by/every_string.py +--! expected_analysis_output +Aggregate [every(cast(true as boolean)) AS every(true)#x] ++- SubqueryAlias test_agg + +- View (`test_agg`, [k#x, v#x]) + +- Project [cast(k#x as int) AS k#x, cast(v#x as boolean) AS v#x] + +- Project [k#x, v#x] + +- SubqueryAlias test_agg + +- LocalRelation [k#x, v#x] +--! expected_optimized_output +Aggregate [min(true) AS every(true)#x] ++- LocalRelation +--! expected_output_schema +struct +--! expected_result ++-------------+ +| every(true) | ++-------------+ +| true | ++-------------+ +printed all 1 rows. +--! expected_result_hash +b5bea41b6c623f7c09f1bf24dcae58ebab3c0cdd90ad966bc43a45b44867e12b +!-- end + + +--! name +bool_and input type checking: decimal (error) +--! script +scripts/group_by/bool_and_decimal_error.py +--! expected_error +[DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE] Cannot resolve "bool_and(1.0)" due to data type mismatch: The first parameter requires the "BOOLEAN" type, however "1.0" has the type "DECIMAL(10,1)". SQLSTATE: 42K09 +!-- end + + +--! name +bool_or input type checking: double (error) +--! script +scripts/group_by/bool_or_double_error.py +--! expected_error +[DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE] Cannot resolve "bool_or(1.0)" due to data type mismatch: The first parameter requires the "BOOLEAN" type, however "1.0" has the type "DOUBLE". SQLSTATE: 42K09 +!-- end + + +--! name +every as window expression +--! tags +unordered +--! script +scripts/group_by/every_window.py +--! expected_analysis_output +Project [k#x, v#x, every(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x] ++- Project [k#x, v#x, every(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x, every(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x] + +- Window [every(v#x) windowspecdefinition(k#x, v#x ASC NULLS FIRST, specifiedwindowframe(RangeFrame, unboundedpreceding$(), currentrow$())) AS every(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x], [k#x], [v#x ASC NULLS FIRST] + +- Project [k#x, v#x] + +- SubqueryAlias test_agg + +- View (`test_agg`, [k#x, v#x]) + +- Project [cast(k#x as int) AS k#x, cast(v#x as boolean) AS v#x] + +- Project [k#x, v#x] + +- SubqueryAlias test_agg + +- LocalRelation [k#x, v#x] +--! expected_optimized_output +Window [min(v#x) windowspecdefinition(k#x, v#x ASC NULLS FIRST, specifiedwindowframe(RangeFrame, unboundedpreceding$(), currentrow$())) AS every(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x], [k#x], [v#x ASC NULLS FIRST] ++- LocalRelation [k#x, v#x] +--! expected_output_schema +struct +--! expected_result ++---+-------+-------------------------------------------------------------------------------------------------------------+ +| k | v | every(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) | ++---+-------+-------------------------------------------------------------------------------------------------------------+ +| 1 | false | false | +| 1 | true | false | +| 2 | true | true | +| 3 | NULL | NULL | +| 3 | false | false | +| 4 | NULL | NULL | +| 4 | NULL | NULL | +| 5 | NULL | NULL | +| 5 | false | false | +| 5 | true | false | ++---+-------+-------------------------------------------------------------------------------------------------------------+ +printed all 10 rows. +--! expected_result_hash +1ad0a462e3a29e3f796db5ad23d8d3970b0d1eb9085231fcc4fcf8e128cf2b2d +!-- end + + +--! name +some as window expression +--! tags +unordered +--! script +scripts/group_by/some_window.py +--! expected_analysis_output +Project [k#x, v#x, some(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x] ++- Project [k#x, v#x, some(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x, some(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x] + +- Window [some(v#x) windowspecdefinition(k#x, v#x ASC NULLS FIRST, specifiedwindowframe(RangeFrame, unboundedpreceding$(), currentrow$())) AS some(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x], [k#x], [v#x ASC NULLS FIRST] + +- Project [k#x, v#x] + +- SubqueryAlias test_agg + +- View (`test_agg`, [k#x, v#x]) + +- Project [cast(k#x as int) AS k#x, cast(v#x as boolean) AS v#x] + +- Project [k#x, v#x] + +- SubqueryAlias test_agg + +- LocalRelation [k#x, v#x] +--! expected_optimized_output +Window [max(v#x) windowspecdefinition(k#x, v#x ASC NULLS FIRST, specifiedwindowframe(RangeFrame, unboundedpreceding$(), currentrow$())) AS some(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x], [k#x], [v#x ASC NULLS FIRST] ++- LocalRelation [k#x, v#x] +--! expected_output_schema +struct +--! expected_result ++---+-------+------------------------------------------------------------------------------------------------------------+ +| k | v | some(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) | ++---+-------+------------------------------------------------------------------------------------------------------------+ +| 1 | false | false | +| 1 | true | true | +| 2 | true | true | +| 3 | NULL | NULL | +| 3 | false | false | +| 4 | NULL | NULL | +| 4 | NULL | NULL | +| 5 | NULL | NULL | +| 5 | false | false | +| 5 | true | true | ++---+-------+------------------------------------------------------------------------------------------------------------+ +printed all 10 rows. +--! expected_result_hash +9a682e38910e199dade7fbe0fd8eb481cf3e924ef37e09ef0c6143b82e9b660e +!-- end + + +--! name +bool_or as window expression +--! tags +unordered +--! script +scripts/group_by/bool_or_window.py +--! expected_analysis_output +Project [k#x, v#x, bool_or(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x] ++- Project [k#x, v#x, bool_or(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x, bool_or(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x] + +- Window [bool_or(v#x) windowspecdefinition(k#x, v#x ASC NULLS FIRST, specifiedwindowframe(RangeFrame, unboundedpreceding$(), currentrow$())) AS bool_or(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x], [k#x], [v#x ASC NULLS FIRST] + +- Project [k#x, v#x] + +- SubqueryAlias test_agg + +- View (`test_agg`, [k#x, v#x]) + +- Project [cast(k#x as int) AS k#x, cast(v#x as boolean) AS v#x] + +- Project [k#x, v#x] + +- SubqueryAlias test_agg + +- LocalRelation [k#x, v#x] +--! expected_optimized_output +Window [max(v#x) windowspecdefinition(k#x, v#x ASC NULLS FIRST, specifiedwindowframe(RangeFrame, unboundedpreceding$(), currentrow$())) AS bool_or(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x], [k#x], [v#x ASC NULLS FIRST] ++- LocalRelation [k#x, v#x] +--! expected_output_schema +struct +--! expected_result ++---+-------+---------------------------------------------------------------------------------------------------------------+ +| k | v | bool_or(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) | ++---+-------+---------------------------------------------------------------------------------------------------------------+ +| 1 | false | false | +| 1 | true | true | +| 2 | true | true | +| 3 | NULL | NULL | +| 3 | false | false | +| 4 | NULL | NULL | +| 4 | NULL | NULL | +| 5 | NULL | NULL | +| 5 | false | false | +| 5 | true | true | ++---+-------+---------------------------------------------------------------------------------------------------------------+ +printed all 10 rows. +--! expected_result_hash +9a682e38910e199dade7fbe0fd8eb481cf3e924ef37e09ef0c6143b82e9b660e +!-- end + + +--! name +bool_and as window expression +--! tags +unordered +--! script +scripts/group_by/bool_and_window.py +--! expected_analysis_output +Project [k#x, v#x, bool_and(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x] ++- Project [k#x, v#x, bool_and(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x, bool_and(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x] + +- Window [bool_and(v#x) windowspecdefinition(k#x, v#x ASC NULLS FIRST, specifiedwindowframe(RangeFrame, unboundedpreceding$(), currentrow$())) AS bool_and(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x], [k#x], [v#x ASC NULLS FIRST] + +- Project [k#x, v#x] + +- SubqueryAlias test_agg + +- View (`test_agg`, [k#x, v#x]) + +- Project [cast(k#x as int) AS k#x, cast(v#x as boolean) AS v#x] + +- Project [k#x, v#x] + +- SubqueryAlias test_agg + +- LocalRelation [k#x, v#x] +--! expected_optimized_output +Window [min(v#x) windowspecdefinition(k#x, v#x ASC NULLS FIRST, specifiedwindowframe(RangeFrame, unboundedpreceding$(), currentrow$())) AS bool_and(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x], [k#x], [v#x ASC NULLS FIRST] ++- LocalRelation [k#x, v#x] +--! expected_output_schema +struct +--! expected_result ++---+-------+----------------------------------------------------------------------------------------------------------------+ +| k | v | bool_and(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) | ++---+-------+----------------------------------------------------------------------------------------------------------------+ +| 1 | false | false | +| 1 | true | false | +| 2 | true | true | +| 3 | NULL | NULL | +| 3 | false | false | +| 4 | NULL | NULL | +| 4 | NULL | NULL | +| 5 | NULL | NULL | +| 5 | false | false | +| 5 | true | false | ++---+-------+----------------------------------------------------------------------------------------------------------------+ +printed all 10 rows. +--! expected_result_hash +1ad0a462e3a29e3f796db5ad23d8d3970b0d1eb9085231fcc4fcf8e128cf2b2d +!-- end + + +--! name +bool_or as window expression (repeated) +--! tags +unordered +--! script +scripts/group_by/bool_or_window_repeated.py +--! expected_analysis_output +Project [k#x, v#x, bool_or(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x] ++- Project [k#x, v#x, bool_or(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x, bool_or(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x] + +- Window [bool_or(v#x) windowspecdefinition(k#x, v#x ASC NULLS FIRST, specifiedwindowframe(RangeFrame, unboundedpreceding$(), currentrow$())) AS bool_or(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x], [k#x], [v#x ASC NULLS FIRST] + +- Project [k#x, v#x] + +- SubqueryAlias test_agg + +- View (`test_agg`, [k#x, v#x]) + +- Project [cast(k#x as int) AS k#x, cast(v#x as boolean) AS v#x] + +- Project [k#x, v#x] + +- SubqueryAlias test_agg + +- LocalRelation [k#x, v#x] +--! expected_optimized_output +Window [max(v#x) windowspecdefinition(k#x, v#x ASC NULLS FIRST, specifiedwindowframe(RangeFrame, unboundedpreceding$(), currentrow$())) AS bool_or(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)#x], [k#x], [v#x ASC NULLS FIRST] ++- LocalRelation [k#x, v#x] +--! expected_output_schema +struct +--! expected_result ++---+-------+---------------------------------------------------------------------------------------------------------------+ +| k | v | bool_or(v) OVER (PARTITION BY k ORDER BY v ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) | ++---+-------+---------------------------------------------------------------------------------------------------------------+ +| 1 | false | false | +| 1 | true | true | +| 2 | true | true | +| 3 | NULL | NULL | +| 3 | false | false | +| 4 | NULL | NULL | +| 4 | NULL | NULL | +| 5 | NULL | NULL | +| 5 | false | false | +| 5 | true | true | ++---+-------+---------------------------------------------------------------------------------------------------------------+ +printed all 10 rows. +--! expected_result_hash +9a682e38910e199dade7fbe0fd8eb481cf3e924ef37e09ef0c6143b82e9b660e +!-- end + + +--! name +having referencing aggregate expression +--! script +scripts/group_by/having_agg_count.py +--! expected_analysis_output +Project [cnt#xL] ++- Filter (cnt#xL > cast(1 as bigint)) + +- Aggregate [count(k#x) AS cnt#xL] + +- SubqueryAlias test_agg + +- View (`test_agg`, [k#x, v#x]) + +- Project [cast(k#x as int) AS k#x, cast(v#x as boolean) AS v#x] + +- Project [k#x, v#x] + +- SubqueryAlias test_agg + +- LocalRelation [k#x, v#x] +--! expected_optimized_output +Filter (cnt#xL > 1) ++- Aggregate [count(1) AS cnt#xL] + +- LocalRelation +--! expected_output_schema +struct +--! expected_result ++-----+ +| cnt | ++-----+ +| 10 | ++-----+ +printed all 1 rows. +--! expected_result_hash +4a44dc15364204a80fe80e9039455cc1608281820fe2b24f1e5233ade6af1dd5 +!-- end + + +--! name +having on aliased max +--! tags +unordered +--! script +scripts/group_by/having_max_v.py +--! expected_analysis_output +Project [k#x, max_v#x] ++- Filter (max_v#x = true) + +- Aggregate [k#x], [k#x, max(v#x) AS max_v#x] + +- SubqueryAlias test_agg + +- View (`test_agg`, [k#x, v#x]) + +- Project [cast(k#x as int) AS k#x, cast(v#x as boolean) AS v#x] + +- Project [k#x, v#x] + +- SubqueryAlias test_agg + +- LocalRelation [k#x, v#x] +--! expected_optimized_output +Filter (isnotnull(max_v#x) AND max_v#x) ++- Aggregate [k#x], [k#x, max(v#x) AS max_v#x] + +- LocalRelation [k#x, v#x] +--! expected_output_schema +struct +--! expected_result ++---+-------+ +| k | max_v | ++---+-------+ +| 1 | true | +| 2 | true | +| 5 | true | ++---+-------+ +printed all 3 rows. +--! expected_result_hash +d6212a5d216fffb25342120cad425c6ebe96fc8485c6b082be8624faf6a531c4 +!-- end + + +--! name +aggregate expression referenced through alias +--! script +scripts/group_by/agg_alias_filter.py +--! expected_analysis_output +Filter (cnt#xL > cast(1 as bigint)) ++- Aggregate [count(k#x) AS cnt#xL] + +- SubqueryAlias test_agg + +- View (`test_agg`, [k#x, v#x]) + +- Project [cast(k#x as int) AS k#x, cast(v#x as boolean) AS v#x] + +- Project [k#x, v#x] + +- SubqueryAlias test_agg + +- LocalRelation [k#x, v#x] +--! expected_optimized_output +Filter (cnt#xL > 1) ++- Aggregate [count(1) AS cnt#xL] + +- LocalRelation +--! expected_output_schema +struct +--! expected_result ++-----+ +| cnt | ++-----+ +| 10 | ++-----+ +printed all 1 rows. +--! expected_result_hash +4a44dc15364204a80fe80e9039455cc1608281820fe2b24f1e5233ade6af1dd5 +!-- end + + +--! name +SPARK-34581: do not optimize out grouping expressions +--! tags +unordered +--! script +scripts/group_by/grouping_exprs_not_optimized.py +--! expected_analysis_output +Aggregate [isnull(a#x)], [isnull(a#x) AS (a IS NULL)#x, NOT isnull(a#x) AS (NOT (a IS NULL))#x, count(1) AS c#xL] ++- SubqueryAlias testdata + +- View (`testData`, [a#x, b#x]) + +- Project [cast(a#x as int) AS a#x, cast(b#x as int) AS b#x] + +- Project [a#x, b#x] + +- SubqueryAlias testData + +- LocalRelation [a#x, b#x] +--! expected_optimized_output +Aggregate [_groupingexpression#x], [_groupingexpression#x AS (a IS NULL)#x, NOT _groupingexpression#x AS (NOT (a IS NULL))#x, count(1) AS c#xL] ++- LocalRelation [_groupingexpression#x] +--! expected_output_schema +struct<(a IS NULL):boolean,(NOT (a IS NULL)):boolean,c:bigint> +--! expected_result ++-------------+-------------------+---+ +| (a IS NULL) | (NOT (a IS NULL)) | c | ++-------------+-------------------+---+ +| false | true | 7 | +| true | false | 2 | ++-------------+-------------------+---+ +printed all 2 rows. +--! expected_result_hash +c5421419b0d389431648774e2aa0156b8472da926b1dd3cd50c94821001e7bcd +!-- end + + +--! name +PullOutGroupingExpressions pulls out grouping expressions +--! tags +unordered +--! script +scripts/group_by/pull_out_grouping_exprs.py +--! expected_analysis_output +Aggregate [isnull(a#x)], [isnull(a#x) AS (a IS NULL)#x, CASE WHEN NOT isnull(a#x) THEN 0 ELSE 1 END AS CASE WHEN (NOT (a IS NULL)) THEN 0 ELSE 1 END#x, first(isnull(a#x), false) AS c#x] ++- SubqueryAlias testdata + +- View (`testData`, [a#x, b#x]) + +- Project [cast(a#x as int) AS a#x, cast(b#x as int) AS b#x] + +- Project [a#x, b#x] + +- SubqueryAlias testData + +- LocalRelation [a#x, b#x] +--! expected_optimized_output +Aggregate [_groupingexpression#x], [_groupingexpression#x AS (a IS NULL)#x, CASE WHEN NOT _groupingexpression#x THEN 0 ELSE 1 END AS CASE WHEN (NOT (a IS NULL)) THEN 0 ELSE 1 END#x, first(isnull(a#x), false) AS c#x] ++- LocalRelation [a#x, _groupingexpression#x] +--! expected_output_schema +struct<(a IS NULL):boolean,CASE WHEN (NOT (a IS NULL)) THEN 0 ELSE 1 END:int,c:boolean> +--! expected_result ++-------------+-----------------------------------------------+-------+ +| (a IS NULL) | CASE WHEN (NOT (a IS NULL)) THEN 0 ELSE 1 END | c | ++-------------+-----------------------------------------------+-------+ +| false | 0 | false | +| true | 1 | true | ++-------------+-----------------------------------------------+-------+ +printed all 2 rows. +--! expected_result_hash +ce61427fcddb697480c2921f7450c541ac0bf1fe284b7e88bd9d784a84a2120a +!-- end + + +--! name +count_if with grouping expression reference +--! tags +unordered +--! script +scripts/group_by/count_if_group_by.py +--! expected_analysis_output +Aggregate [(a#x + 1)], [(a#x + 1) AS (a + 1)#x, count_if(((a#x + 1) = b#x)) AS count_if(((a + 1) = b))#xL] ++- SubqueryAlias testdata + +- View (`testData`, [a#x, b#x]) + +- Project [cast(a#x as int) AS a#x, cast(b#x as int) AS b#x] + +- Project [a#x, b#x] + +- SubqueryAlias testData + +- LocalRelation [a#x, b#x] +--! expected_optimized_output +Aggregate [_groupingexpression#x], [_groupingexpression#x AS (a + 1)#x, count(if (NOT _common_expr_0#x) null else _common_expr_0#x) AS count_if(((a + 1) = b))#xL] ++- LocalRelation [_groupingexpression#x, _common_expr_0#x] +--! expected_output_schema +struct<(a + 1):int,count_if(((a + 1) = b)):bigint> +--! expected_result ++---------+-------------------------+ +| (a + 1) | count_if(((a + 1) = b)) | ++---------+-------------------------+ +| 2 | 1 | +| 3 | 0 | +| 4 | 0 | +| NULL | 0 | ++---------+-------------------------+ +printed all 4 rows. +--! expected_result_hash +a71fb4e2edd0d56bc2b849747cc6dc084d03786bfaed0681f6ee9b564f09f089 +!-- end diff --git a/python/pyspark/sql/tests/df_golden/regenerate.sh b/python/pyspark/sql/tests/df_golden/regenerate.sh new file mode 100755 index 0000000000000..a185abaebbf18 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/regenerate.sh @@ -0,0 +1,53 @@ +#!/usr/bin/env bash +# +# 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. +# +# Regenerates the DataFrame golden test files (*.test) in this directory by +# re-running the test module with SPARK_GENERATE_GOLDEN_FILES=1. +# +# Usage: +# python/pyspark/sql/tests/df_golden/regenerate.sh [--verify] +# +# --verify re-run the tests against the regenerated golden files. + +set -euo pipefail + +VERIFY=false +for arg in "$@"; do + case "$arg" in + --verify) VERIFY=true ;; + *) echo "ERROR: unknown argument: $arg" >&2; exit 1 ;; + esac +done + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && git rev-parse --show-toplevel)" +cd "$REPO_ROOT" + +MODULE="pyspark.sql.tests.df_golden.test_df_golden" +REL_DIR="python/pyspark/sql/tests/df_golden" + +echo ">>> Regenerating golden files..." +SPARK_GENERATE_GOLDEN_FILES=1 python/run-tests --testnames "$MODULE" + +echo ">>> Regenerated. Local changes:" +git status --short -- "$REL_DIR" + +if [[ "$VERIFY" == true ]]; then + echo ">>> Verifying against the regenerated golden files..." + python/run-tests --testnames "$MODULE" +else + echo ">>> Done. Review the diff, then optionally verify with: $0 --verify" +fi diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/agg_alias_filter.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/agg_alias_filter.py new file mode 100644 index 0000000000000..1db2ee02d52f3 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/agg_alias_filter.py @@ -0,0 +1,5 @@ +# Aggregate expressions can be referenced through an alias + +from pyspark.sql.functions import col, count + +df = spark.table("test_agg").agg(count(col("k")).alias("cnt")).filter(col("cnt") > 1) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/agg_counts_no_grouping.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/agg_counts_no_grouping.py new file mode 100644 index 0000000000000..969c5f69d8d15 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/agg_counts_no_grouping.py @@ -0,0 +1,3 @@ +from pyspark.sql.functions import col, count + +df = spark.table("testData").agg(count(col("a")), count(col("b"))) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/agg_stats_with_nulls.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/agg_stats_with_nulls.py new file mode 100644 index 0000000000000..db0ecf242d791 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/agg_stats_with_nulls.py @@ -0,0 +1,27 @@ +# Aggregate with nulls. + +from pyspark.sql.functions import ( + avg, + col, + count, + kurtosis, + max, + min, + round, + skewness, + stddev, + sum, + variance, +) + +df = spark.table("testData").agg( + round(skewness(col("a")), 12), + round(kurtosis(col("a")), 12), + min(col("a")), + max(col("a")), + round(avg(col("a")), 12), + round(variance(col("a")), 12), + round(stddev(col("a")), 12), + sum(col("a")), + count(col("a")), +) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_all_nulls.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_all_nulls.py new file mode 100644 index 0000000000000..2115205c1a38d --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_all_nulls.py @@ -0,0 +1,9 @@ +# all null values + +from pyspark.sql.functions import bool_and, bool_or, col, every, some + +df = ( + spark.table("test_agg") + .filter(col("k") == 4) + .agg(every(col("v")), some(col("v")), bool_or(col("v")), bool_and(col("v")), bool_or(col("v"))) +) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_empty_table.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_empty_table.py new file mode 100644 index 0000000000000..727000377f6c6 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_empty_table.py @@ -0,0 +1,9 @@ +# empty table + +from pyspark.sql.functions import bool_and, bool_or, col, every, lit, some + +df = ( + spark.table("test_agg") + .filter(lit(1) == 0) + .agg(every(col("v")), some(col("v")), bool_or(col("v")), bool_and(col("v")), bool_or(col("v"))) +) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_group_by.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_group_by.py new file mode 100644 index 0000000000000..6bc7ebf478cea --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_group_by.py @@ -0,0 +1,9 @@ +# group by + +from pyspark.sql.functions import bool_and, bool_or, col, every, some + +df = ( + spark.table("test_agg") + .groupBy(col("k")) + .agg(every(col("v")), some(col("v")), bool_or(col("v")), bool_and(col("v")), bool_or(col("v"))) +) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_having_false.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_having_false.py new file mode 100644 index 0000000000000..1d8837200bca9 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_having_false.py @@ -0,0 +1,11 @@ +# having + +from pyspark.sql.functions import col, every + +df = ( + spark.table("test_agg") + .groupBy(col("k")) + .agg(every(col("v")).alias("every_v")) + .filter(col("every_v") == False) # noqa: E712 + .select(col("k"), col("every_v")) +) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_having_null.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_having_null.py new file mode 100644 index 0000000000000..99cf57b5f1541 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_having_null.py @@ -0,0 +1,9 @@ +from pyspark.sql.functions import col, every + +df = ( + spark.table("test_agg") + .groupBy(col("k")) + .agg(every(col("v")).alias("every_v")) + .filter(col("every_v").isNull()) + .select(col("k"), col("every_v")) +) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_null_filtering.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_null_filtering.py new file mode 100644 index 0000000000000..81c19c4db4a17 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_null_filtering.py @@ -0,0 +1,9 @@ +# aggregates are null Filtering + +from pyspark.sql.functions import bool_and, bool_or, col, every, some + +df = ( + spark.table("test_agg") + .filter(col("k") == 5) + .agg(every(col("v")), some(col("v")), bool_or(col("v")), bool_and(col("v")), bool_or(col("v"))) +) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_and_decimal_error.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_and_decimal_error.py new file mode 100644 index 0000000000000..c212f96113b21 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_and_decimal_error.py @@ -0,0 +1,6 @@ +# input type checking Decimal + +from decimal import Decimal +from pyspark.sql.functions import bool_and, lit + +df = spark.table("test_agg").select(bool_and(lit(Decimal("1.0")))) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_and_window.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_and_window.py new file mode 100644 index 0000000000000..1fae3d538036a --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_and_window.py @@ -0,0 +1,6 @@ +from pyspark.sql import Window +from pyspark.sql.functions import bool_and, col + +df = spark.table("test_agg").select( + col("k"), col("v"), bool_and(col("v")).over(Window.partitionBy("k").orderBy("v")) +) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_or_double_error.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_or_double_error.py new file mode 100644 index 0000000000000..6176a2ce894d0 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_or_double_error.py @@ -0,0 +1,5 @@ +# input type checking double + +from pyspark.sql.functions import bool_or, lit + +df = spark.table("test_agg").select(bool_or(lit(1.0))) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_or_int_error.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_or_int_error.py new file mode 100644 index 0000000000000..5ad8187a2998d --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_or_int_error.py @@ -0,0 +1,5 @@ +# input type checking Long + +from pyspark.sql.functions import bool_or, lit + +df = spark.table("test_agg").select(bool_or(lit(1))) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_or_window.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_or_window.py new file mode 100644 index 0000000000000..a7e4dc0d281f6 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_or_window.py @@ -0,0 +1,6 @@ +from pyspark.sql import Window +from pyspark.sql.functions import bool_or, col + +df = spark.table("test_agg").select( + col("k"), col("v"), bool_or(col("v")).over(Window.partitionBy("k").orderBy("v")) +) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_or_window_repeated.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_or_window_repeated.py new file mode 100644 index 0000000000000..a7e4dc0d281f6 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_or_window_repeated.py @@ -0,0 +1,6 @@ +from pyspark.sql import Window +from pyspark.sql.functions import bool_or, col + +df = spark.table("test_agg").select( + col("k"), col("v"), bool_or(col("v")).over(Window.partitionBy("k").orderBy("v")) +) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/count_if_group_by.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/count_if_group_by.py new file mode 100644 index 0000000000000..35443d59fb2fc --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/count_if_group_by.py @@ -0,0 +1,3 @@ +from pyspark.sql.functions import col, count_if + +df = spark.table("testData").groupBy(col("a") + 1).agg(count_if((col("a") + 1) == col("b"))) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/empty_input_agg.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/empty_input_agg.py new file mode 100644 index 0000000000000..a5efcab8a0897 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/empty_input_agg.py @@ -0,0 +1,5 @@ +# Aggregate with empty input and empty GroupBy expressions. + +from pyspark.sql.functions import count, lit + +df = spark.table("testData").filter(lit(False)).agg(count(lit(1))) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/empty_input_agg_select.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/empty_input_agg_select.py new file mode 100644 index 0000000000000..e7fd87eaa97f3 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/empty_input_agg_select.py @@ -0,0 +1,3 @@ +from pyspark.sql.functions import count, lit + +df = spark.table("testData").filter(lit(False)).agg(count(lit(1))).select(lit(1)) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/empty_input_group_by.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/empty_input_group_by.py new file mode 100644 index 0000000000000..cb64aaf418c09 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/empty_input_group_by.py @@ -0,0 +1,5 @@ +# Aggregate with empty input and non-empty GroupBy expressions. + +from pyspark.sql.functions import col, count, lit + +df = spark.table("testData").filter(lit(False)).groupBy(col("a")).agg(count(lit(1))) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/every_int_error.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/every_int_error.py new file mode 100644 index 0000000000000..0e484981c9c48 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/every_int_error.py @@ -0,0 +1,5 @@ +# input type checking Int + +from pyspark.sql.functions import every, lit + +df = spark.table("test_agg").select(every(lit(1))) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/every_string.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/every_string.py new file mode 100644 index 0000000000000..5a69c9075e999 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/every_string.py @@ -0,0 +1,5 @@ +# input type checking String + +from pyspark.sql.functions import every, lit + +df = spark.table("test_agg").select(every(lit("true"))) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/every_window.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/every_window.py new file mode 100644 index 0000000000000..95c1241a917b1 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/every_window.py @@ -0,0 +1,8 @@ +# every/some/any aggregates/bool_and/bool_or are supported as windows expression. + +from pyspark.sql import Window +from pyspark.sql.functions import col, every + +df = spark.table("test_agg").select( + col("k"), col("v"), every(col("v")).over(Window.partitionBy("k").orderBy("v")) +) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_complex_expr.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_complex_expr.py new file mode 100644 index 0000000000000..a3f362b2cf405 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_complex_expr.py @@ -0,0 +1,5 @@ +# Aggregate with complex GroupBy expressions. + +from pyspark.sql.functions import col, count + +df = spark.table("testData").groupBy(col("a") + col("b")).agg(count(col("b"))) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_count.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_count.py new file mode 100644 index 0000000000000..483ed988a46ea --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_count.py @@ -0,0 +1,5 @@ +# Aggregate with non-empty GroupBy expressions. + +from pyspark.sql.functions import col, count + +df = spark.table("testData").groupBy(col("a")).agg(count(col("b"))) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_literal.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_literal.py new file mode 100644 index 0000000000000..c4f5d6b216731 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_literal.py @@ -0,0 +1,5 @@ +# Aggregate grouped by literals. + +from pyspark.sql.functions import col, count, lit + +df = spark.table("testData").groupBy(lit("foo")).agg(count(col("a"))) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_literal_hash_agg.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_literal_hash_agg.py new file mode 100644 index 0000000000000..b4d42eaf0baf6 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_literal_hash_agg.py @@ -0,0 +1,10 @@ +# Aggregate grouped by literals (hash aggregate). + +from pyspark.sql.functions import approx_count_distinct, col, lit + +df = ( + spark.table("testData") + .filter(col("a") == 0) + .groupBy(lit("foo")) + .agg(approx_count_distinct(col("a"))) +) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_literal_sort_agg.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_literal_sort_agg.py new file mode 100644 index 0000000000000..cb83ecdf9ee0e --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_literal_sort_agg.py @@ -0,0 +1,5 @@ +# Aggregate grouped by literals (sort aggregate). + +from pyspark.sql.functions import col, lit, max, struct + +df = spark.table("testData").filter(col("a") == 0).groupBy(lit("foo")).agg(max(struct(col("a")))) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_missing_aggregation.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_missing_aggregation.py new file mode 100644 index 0000000000000..ff822ed8f9ca5 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_missing_aggregation.py @@ -0,0 +1,3 @@ +from pyspark.sql.functions import col, count + +df = spark.table("testData").groupBy(col("b")).agg(col("a"), count(col("b"))) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_multiple_counts.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_multiple_counts.py new file mode 100644 index 0000000000000..8e1a2d5ea055c --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_multiple_counts.py @@ -0,0 +1,3 @@ +from pyspark.sql.functions import col, count + +df = spark.table("testData").groupBy(col("a")).agg(count(col("a")), count(col("b"))) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_struct.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_struct.py new file mode 100644 index 0000000000000..ff6e728367afc --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_struct.py @@ -0,0 +1,5 @@ +# struct() in group by + +from pyspark.sql.functions import col, count, lit, struct + +df = spark.table("testData").groupBy(struct((col("a") + 0.1).alias("aa"))).agg(count(lit(1))) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/grouping_exprs_not_optimized.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/grouping_exprs_not_optimized.py new file mode 100644 index 0000000000000..91086696a432a --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/grouping_exprs_not_optimized.py @@ -0,0 +1,10 @@ +# SPARK-34581: Don't optimize out grouping expressions from aggregate expressions +# without aggregate function + +from pyspark.sql.functions import col, count + +df = ( + spark.table("testData") + .groupBy(col("a").isNull()) + .agg(~col("a").isNull(), count("*").alias("c")) +) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/having_agg_count.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/having_agg_count.py new file mode 100644 index 0000000000000..879f498da0e8a --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/having_agg_count.py @@ -0,0 +1,10 @@ +# Having referencing aggregate expressions is ok. + +from pyspark.sql.functions import col, count + +df = ( + spark.table("test_agg") + .agg(count(col("k")).alias("cnt")) + .filter(col("cnt") > 1) + .select(col("cnt")) +) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/having_max_v.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/having_max_v.py new file mode 100644 index 0000000000000..3303169b549fd --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/having_max_v.py @@ -0,0 +1,9 @@ +from pyspark.sql.functions import col, max + +df = ( + spark.table("test_agg") + .groupBy(col("k")) + .agg(max(col("v")).alias("max_v")) + .filter(col("max_v") == True) # noqa: E712 + .select(col("k"), col("max_v")) +) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/pull_out_grouping_exprs.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/pull_out_grouping_exprs.py new file mode 100644 index 0000000000000..8d1f6807f21a7 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/pull_out_grouping_exprs.py @@ -0,0 +1,10 @@ +# PullOutGroupingExpressions also pulls out grouping expressions from inside +# AggregateExpressions + +from pyspark.sql.functions import col, first, when + +df = ( + spark.table("testData") + .groupBy(col("a").isNull()) + .agg(when(~col("a").isNull(), 0).otherwise(1), first(col("a").isNull()).alias("c")) +) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/select_non_grouping_column.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/select_non_grouping_column.py new file mode 100644 index 0000000000000..466a65a153d79 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/select_non_grouping_column.py @@ -0,0 +1,5 @@ +# Aggregate with empty GroupBy expressions. + +from pyspark.sql.functions import col, count + +df = spark.table("testData").select(col("a"), count(col("b"))) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/setup_test_agg.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/setup_test_agg.py new file mode 100644 index 0000000000000..8e1734fe1ac3e --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/setup_test_agg.py @@ -0,0 +1,8 @@ +# Test data + +df = spark.sql("""CREATE OR REPLACE TEMPORARY VIEW test_agg AS SELECT * FROM VALUES + (1, true), (1, false), + (2, true), + (3, false), (3, null), + (4, null), (4, null), + (5, null), (5, true), (5, false) AS test_agg(k, v)""") diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/setup_test_data.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/setup_test_data.py new file mode 100644 index 0000000000000..4a83505a3d79d --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/setup_test_data.py @@ -0,0 +1,6 @@ +# Test aggregate operator. +# Test data. + +df = spark.sql("""CREATE OR REPLACE TEMPORARY VIEW testData AS SELECT * FROM VALUES +(1, 1), (1, 2), (2, 1), (2, 2), (3, 1), (3, 2), (null, 1), (3, null), (null, null) +AS testData(a, b)""") diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/some_int_error.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/some_int_error.py new file mode 100644 index 0000000000000..3bd9dbba84172 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/some_int_error.py @@ -0,0 +1,5 @@ +# input type checking Short + +from pyspark.sql.functions import lit, some + +df = spark.table("test_agg").select(some(lit(1))) diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/some_window.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/some_window.py new file mode 100644 index 0000000000000..21c84d14babea --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/some_window.py @@ -0,0 +1,6 @@ +from pyspark.sql import Window +from pyspark.sql.functions import col, some + +df = spark.table("test_agg").select( + col("k"), col("v"), some(col("v")).over(Window.partitionBy("k").orderBy("v")) +) diff --git a/python/pyspark/sql/tests/df_golden/test_df_golden.py b/python/pyspark/sql/tests/df_golden/test_df_golden.py new file mode 100644 index 0000000000000..ee9b8a240e281 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/test_df_golden.py @@ -0,0 +1,143 @@ +# +# 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. +# + +""" +DataFrame API golden file tests. + +Analogous to ``SQLQueryTestSuite`` but tests are written in native Python using +the PySpark DataFrame API. Each ``.test`` file in this directory describes a +list of test cases; every case points to a standalone script under +``scripts/`` that builds the DataFrame under test. The expected outputs live +inline in the ``.test`` file, which doubles as the golden file. See +``pyspark.testing.df_golden`` for the file format. + +Running the tests:: + + python/run-tests --testnames pyspark.sql.tests.df_golden.test_df_golden + +Regenerating golden files +------------------------- +Set ``SPARK_GENERATE_GOLDEN_FILES=1`` before running the tests, or use the +wrapper script:: + + python/pyspark/sql/tests/df_golden/regenerate.sh [--verify] + +With ``--verify`` the wrapper re-runs the tests afterwards against the +regenerated files. + +Adding a new test file +---------------------- +Drop a new ``.test`` file in this directory and its case scripts under +``scripts//`` -- the ``.test`` file is auto-discovered (a +``test_`` method is registered here), no code changes needed. Then +regenerate the golden files with the steps above. +""" + +import os + +from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.df_golden import run_golden_test + + +_THIS_DIR = os.path.dirname(os.path.abspath(__file__)) + + +class DFGoldenTestBase(ReusedConnectTestCase): + """ + Base class for DataFrame golden file tests. + + Subclasses just need to exist -- test methods are auto-registered from + the ``.test`` files found in ``test_file_dir``. + + Each ``.test`` file runs in its own Spark Connect session (a + ``spark.newSession()`` call against the shared local Connect server that + :class:`ReusedConnectTestCase` starts in ``setUpClass``), which is the + Connect counterpart of ``SQLQueryTestSuite``'s per-file + ``spark.newSession()``. State created by a file's case scripts -- temp + views, UDFs, session confs -- is discarded with the session and cannot + leak into other files. + """ + + test_file_dir = _THIS_DIR + + def _run_golden_test(self, filename): + session = self.spark.newSession() + try: + run_golden_test(self, session, os.path.join(self.test_file_dir, filename)) + finally: + # Release only this sub-session server-side and close its client + # channel. We must NOT call ``session.stop()``: under + # ``SPARK_LOCAL_REMOTE`` (the test harness) ``stop()`` terminates the + # shared local Connect server, breaking the rest of the suite and + # hanging ``tearDownClass``'s ``spark.stop()`` in release retries + # against the dead server until the test times out. + client = session.client + try: + client.release_session() + except Exception: + pass + try: + client.close() + except Exception: + pass + + +def _register_tests(cls): + """ + Discover ``.test`` files and create a test method for each one on *cls*. + + Finding no ``.test`` files registers a failing test instead of zero + tests: a packaging problem that drops the data files must fail the + suite, not silently turn it into a green no-op. + """ + test_files = [] + if os.path.isdir(cls.test_file_dir): + test_files = sorted(f for f in os.listdir(cls.test_file_dir) if f.endswith(".test")) + + if not test_files: + + def test_discovery_failed(self): + self.fail("no .test files found in {}".format(cls.test_file_dir)) + + cls.test_discovery_failed = test_discovery_failed + return + + for filename in test_files: + + def _make_test(f=filename): + def test_method(self): + self._run_golden_test(f) + + return test_method + + name = filename[: -len(".test")] + setattr(cls, "test_{}".format(name), _make_test()) + + +class DFGoldenTest(DFGoldenTestBase): + """DataFrame golden file tests for all ``.test`` files in this directory.""" + + pass + + +_register_tests(DFGoldenTest) + + +if __name__ == "__main__": + from pyspark.testing import main + + main() diff --git a/python/pyspark/sql/tests/df_golden/test_df_golden_framework.py b/python/pyspark/sql/tests/df_golden/test_df_golden_framework.py new file mode 100644 index 0000000000000..052a52182af26 --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/test_df_golden_framework.py @@ -0,0 +1,537 @@ +# +# 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. +# + +""" +Unit tests for the ``pyspark.testing.df_golden`` write/validation machinery. + +These exercise the pure ``.test`` file plumbing -- parsing, serialization, +validation, output normalization and result rendering -- without a Spark +session, so they are fast and run anywhere. The end-to-end golden runs that +need a Spark Connect server live in ``test_df_golden.py``. +""" + +import os +import tempfile +import unittest + +from pyspark.testing.df_golden import ( + _compare_case, + _regenerate_case, + _validate_test_file, + format_double, + format_error, + hash_result_rows, + parse_tags, + parse_test_file, + render_result_table, + replace_not_included, + write_test_file, +) + + +class DFGoldenFrameworkTests(unittest.TestCase): + # -- parse / serialize ------------------------------------------------ + + def _write(self, text): + """Write *text* to a temp ``.test`` file and return its path.""" + fd, path = tempfile.mkstemp(suffix=".test") + os.close(fd) + with open(path, "w") as f: + f.write(text) + self.addCleanup(os.remove, path) + return path + + def test_parse_basic_case(self): + path = self._write( + "--! name\n" + "my case\n" + "--! script\n" + "scripts/x.py\n" + "--! expected_output_schema\n" + "struct\n" + "!-- end\n" + ) + header, cases = parse_test_file(path) + self.assertEqual(header, {}) + self.assertEqual(len(cases), 1) + self.assertEqual(cases[0]["name"], "my case") + self.assertEqual(cases[0]["script"], "scripts/x.py") + self.assertEqual(cases[0]["expected_output_schema"], "struct") + + def test_parse_extracts_file_metadata_header(self): + path = self._write( + "--! name\n" + "__file_metadata__\n" + "--! source\n" + "df_golden/group_by\n" + "!-- end\n" + "\n\n" + "--! name\n" + "c1\n" + "--! script\n" + "scripts/a.py\n" + "!-- end\n" + ) + header, cases = parse_test_file(path) + # The header block is lifted out and its synthetic name dropped. + self.assertEqual(header, {"source": "df_golden/group_by"}) + self.assertEqual(len(cases), 1) + self.assertEqual(cases[0]["name"], "c1") + + def test_parse_preserves_multiline_section_body(self): + path = self._write( + "--! name\n" + "c\n" + "--! script\n" + "scripts/a.py\n" + "--! expected_analysis_output\n" + "Sort [k#x ASC], true\n" + "+- Project\n" + " +- Range\n" + "!-- end\n" + ) + _, cases = parse_test_file(path) + self.assertEqual( + cases[0]["expected_analysis_output"], + "Sort [k#x ASC], true\n+- Project\n +- Range", + ) + + def test_round_trip_parse_write_parse(self): + header = {"source": "df_golden/group_by"} + cases = [ + { + "name": "ordered case", + "script": "scripts/a.py", + "expected_analysis_output": "Sort [k#x ASC], true\n+- Range", + "expected_output_schema": "struct", + "expected_result": "+---+\n| k |\n+---+\n| 1 |\n+---+\nprinted all 1 rows.", + "expected_result_hash": "abc123", + }, + { + "name": "error case", + "tags": "unordered", + "script": "scripts/b.py", + "expected_error": "[SOME_ERROR] boom", + }, + ] + path = self._write("") + write_test_file(path, header, cases) + header2, cases2 = parse_test_file(path) + self.assertEqual(header2, header) + self.assertEqual(cases2, cases) + + def test_write_only_emits_known_sections_in_order(self): + path = self._write("") + # ``junk`` is not in the canonical section order and must be dropped. + write_test_file( + path, + {}, + [{"script": "scripts/a.py", "name": "c", "junk": "ignored"}], + ) + with open(path) as f: + body = f.read() + self.assertNotIn("junk", body) + # name precedes script in the canonical order regardless of dict order. + self.assertLess(body.index("--! name"), body.index("--! script")) + + # -- tags ------------------------------------------------------------- + + def test_parse_tags_splits_on_whitespace_and_commas(self): + self.assertEqual(parse_tags({"tags": "unordered, foo bar"}), {"unordered", "foo", "bar"}) + self.assertEqual(parse_tags({}), set()) + self.assertEqual(parse_tags({"tags": ""}), set()) + + # -- validation ------------------------------------------------------- + + def _valid_case(self, **overrides): + case = { + "name": "c", + "script": "scripts/a.py", + "expected_output_schema": "struct", + } + case.update(overrides) + return case + + def test_validate_accepts_well_formed_file(self): + # Should not raise. + _validate_test_file("f.test", {"source": "x"}, [self._valid_case()], regenerate=False) + + def test_validate_rejects_unknown_header_section(self): + with self.assertRaisesRegex(AssertionError, "unknown header sections: bogus"): + _validate_test_file("f.test", {"bogus": "x"}, [self._valid_case()], regenerate=False) + + def test_validate_rejects_no_cases(self): + with self.assertRaisesRegex(AssertionError, "no test cases found"): + _validate_test_file("f.test", {}, [], regenerate=False) + + def test_validate_rejects_case_without_name(self): + with self.assertRaisesRegex(AssertionError, "every test case needs a name"): + _validate_test_file("f.test", {}, [{"script": "scripts/a.py"}], regenerate=False) + + def test_validate_rejects_case_without_script(self): + with self.assertRaisesRegex(AssertionError, "needs a script"): + _validate_test_file( + "f.test", + {}, + [{"name": "c", "expected_output_schema": "x"}], + regenerate=False, + ) + + def test_validate_rejects_unknown_section(self): + with self.assertRaisesRegex(AssertionError, "unknown sections: expected_bogus"): + _validate_test_file( + "f.test", {}, [self._valid_case(expected_bogus="x")], regenerate=False + ) + + def test_validate_rejects_unknown_tag(self): + with self.assertRaisesRegex(AssertionError, "unknown tags: wat"): + _validate_test_file("f.test", {}, [self._valid_case(tags="wat")], regenerate=False) + + def test_validate_accepts_known_unordered_tag(self): + _validate_test_file("f.test", {}, [self._valid_case(tags="unordered")], regenerate=False) + + def test_validate_rejects_vacuous_case_in_verify_mode(self): + # A case with no expected_* section asserts nothing. + with self.assertRaisesRegex(AssertionError, "would assert\n?.*nothing"): + _validate_test_file( + "f.test", {}, [{"name": "c", "script": "scripts/a.py"}], regenerate=False + ) + + def test_validate_allows_vacuous_case_in_regenerate_mode(self): + # New cases legitimately have no expected_* sections before regeneration. + _validate_test_file( + "f.test", {}, [{"name": "c", "script": "scripts/a.py"}], regenerate=True + ) + + def test_validate_accepts_error_only_case(self): + # A case carrying only ``expected_error`` is not vacuous: the error is a + # recognized result section, the single output an error case produces. + _validate_test_file( + "f.test", + {}, + [{"name": "c", "script": "scripts/a.py", "expected_error": "[ERR] boom"}], + regenerate=False, + ) + + def test_validate_rejects_legacy_split_error_sections(self): + # The old analysis/execution split was collapsed into ``expected_error``; + # the legacy names are now unknown sections in verify mode. + for legacy in ("expected_analysis_error", "expected_execution_error"): + with self.assertRaisesRegex(AssertionError, "unknown sections: " + legacy): + _validate_test_file( + "f.test", {}, [self._valid_case(**{legacy: "x"})], regenerate=False + ) + + def test_validate_tolerates_unknown_sections_in_regenerate_mode(self): + # Regeneration drops and rewrites unknown sections, so a file still + # carrying a renamed/removed section (e.g. the legacy error split) must + # not be rejected during regeneration - otherwise it could never be + # migrated. + _validate_test_file( + "f.test", + {}, + [self._valid_case(expected_analysis_error="x")], + regenerate=True, + ) + + def test_validate_rejects_unknown_tag_even_in_regenerate_mode(self): + # Tags are preserved verbatim across regeneration, so an unknown tag + # would persist; it is rejected in both modes. + with self.assertRaisesRegex(AssertionError, "unknown tags: wat"): + _validate_test_file("f.test", {}, [self._valid_case(tags="wat")], regenerate=True) + + # -- output normalization -------------------------------------------- + + def test_replace_not_included_normalizes_volatile_ids(self): + self.assertEqual(replace_not_included("k#1234 + v#5"), "k#x + v#x") + self.assertEqual(replace_not_included("plan_id=42"), "plan_id=x") + self.assertEqual( + replace_not_included("CTERelationDef 17, false"), + "CTERelationDef xxxx, false", + ) + + def test_format_error_strips_volatile_trailers(self): + msg = ( + "[DIVIDE_BY_ZERO] Division by zero. SQLSTATE: 22012\n" + "== DataFrame ==\n" + '"__truediv__" was called from /abs/path/script.py:7\n' + "\n" + "JVM stacktrace:\n" + "org.apache.spark.SparkArithmeticException: ..." + ) + self.assertEqual( + format_error(Exception(msg)), + "[DIVIDE_BY_ZERO] Division by zero. SQLSTATE: 22012", + ) + + def test_format_error_strips_trailing_plan_dump(self): + self.assertEqual( + format_error(Exception("[ERR] bad column;\nProject [a#1]\n+- Range")), + "[ERR] bad column", + ) + + def test_format_error_keeps_message_with_internal_semicolon_newline(self): + # ";\n" not followed by a plan root (uppercase / "'") is part of the + # message and must be preserved, not treated as the plan separator. + msg = "[ERR] first clause;\nand the second clause continues" + self.assertEqual(format_error(Exception(msg)), msg) + + # -- result rendering ------------------------------------------------- + + def test_render_result_table_pads_columns(self): + table = render_result_table(["k", "v"], ["1\t10", "200\t3"]) + self.assertEqual( + table, + "\n".join( + [ + "+-----+----+", + "| k | v |", + "+-----+----+", + "| 1 | 10 |", + "| 200 | 3 |", + "+-----+----+", + "printed all 2 rows.", + ] + ), + ) + + def test_render_result_table_no_columns_is_trailer_only(self): + self.assertEqual(render_result_table([], []), "printed all 0 rows.") + + def test_hash_result_rows_is_stable_and_order_sensitive(self): + h1 = hash_result_rows(["a", "b"]) + self.assertEqual(h1, hash_result_rows(["a", "b"])) + self.assertNotEqual(h1, hash_result_rows(["b", "a"])) + + # -- case comparison -------------------------------------------------- + + def test_compare_case_passes_on_match(self): + case = { + "name": "c", + "expected_output_schema": "struct", + } + _compare_case(self, case, {"expected_output_schema": "struct"}) + + def test_compare_case_ignores_sections_absent_from_golden(self): + # Only sections present in the golden file are checked; extras in + # ``actual`` are ignored. + _compare_case( + self, + {"name": "c", "expected_output_schema": "struct"}, + { + "expected_output_schema": "struct", + "expected_optimized_output": "Range", + }, + ) + + def test_compare_case_fails_on_value_mismatch(self): + with self.assertRaises(AssertionError): + _compare_case( + self, + {"name": "c", "expected_output_schema": "struct"}, + {"expected_output_schema": "struct"}, + ) + + def test_compare_case_fails_when_expected_section_not_produced(self): + with self.assertRaisesRegex(AssertionError, "expected section `expected_result`"): + _compare_case( + self, + {"name": "c", "expected_result": "printed all 0 rows."}, + {"expected_error": "[ERR] boom"}, + ) + + # -- under-assertion guards ------------------------------------------ + + def test_validate_accepts_result_with_hash(self): + _validate_test_file( + "f.test", + {}, + [self._valid_case(expected_result="printed all 0 rows.", expected_result_hash="h")], + regenerate=False, + ) + + def test_validate_rejects_result_without_hash(self): + with self.assertRaisesRegex(AssertionError, "or neither"): + _validate_test_file( + "f.test", + {}, + [self._valid_case(expected_result="printed all 0 rows.")], + regenerate=False, + ) + + def test_validate_rejects_hash_without_result(self): + with self.assertRaisesRegex(AssertionError, "or neither"): + _validate_test_file( + "f.test", {}, [self._valid_case(expected_result_hash="h")], regenerate=False + ) + + def test_validate_rejects_error_case_mixed_with_result(self): + with self.assertRaisesRegex(AssertionError, "must carry only `expected_error`"): + _validate_test_file( + "f.test", + {}, + [ + { + "name": "c", + "script": "scripts/a.py", + "expected_error": "[ERR] boom", + "expected_result": "printed all 0 rows.", + "expected_result_hash": "h", + } + ], + regenerate=False, + ) + + def test_validate_rejects_error_case_mixed_with_plan(self): + with self.assertRaisesRegex(AssertionError, "must carry only `expected_error`"): + _validate_test_file( + "f.test", + {}, + [ + { + "name": "c", + "script": "scripts/a.py", + "expected_error": "[ERR] boom", + "expected_analysis_output": "Range", + } + ], + regenerate=False, + ) + + # -- double formatting / float refusal -------------------------------- + + def test_format_double_matches_java_double_to_string(self): + # Hive renders doubles via Java Double.toString + # (HiveResult: ``case (n, _: NumericType) => n.toString``); these are the + # cases where it diverges from Python's str()/repr. + # A list (not a dict): 0.0 and -0.0 compare equal and would collide as + # dict keys. + cases = [ + (1.0, "1.0"), + (100.0, "100.0"), + (0.5, "0.5"), + (0.001, "0.001"), + (2.142857142857, "2.142857142857"), + (-0.272380105815, "-0.272380105815"), + (1e7, "1.0E7"), + (1e-4, "1.0E-4"), + (12345678.0, "1.2345678E7"), + (1234567.0, "1234567.0"), + (1e20, "1.0E20"), + (0.0, "0.0"), + (-0.0, "-0.0"), + (float("nan"), "NaN"), + (float("inf"), "Infinity"), + (float("-inf"), "-Infinity"), + ] + for value, expected in cases: + self.assertEqual(format_double(value), expected, "format_double(%r)" % value) + + def test_format_value_refuses_float(self): + # double is supported via format_double; float still needs float32- + # shortest rendering to match Java Float.toString, so it is refused. + try: + from pyspark.sql.types import DoubleType, FloatType + except Exception: + self.skipTest("pyspark.sql.types unavailable in this environment") + from pyspark.testing.df_golden import _format_value + + with self.assertRaisesRegex(AssertionError, "not supported yet"): + _format_value(0.1, FloatType()) + # double does not raise. + self.assertEqual(_format_value(0.5, DoubleType()), "0.5") + + # -- loud parser failures -------------------------------------------- + + def test_parse_rejects_duplicate_section(self): + path = self._write( + "--! name\nc\n--! script\nscripts/a.py\n--! script\nscripts/b.py\n!-- end\n" + ) + with self.assertRaisesRegex(AssertionError, "duplicate section `script`"): + parse_test_file(path) + + def test_parse_rejects_malformed_marker(self): + # Missing space after "--!": a typo'd marker, not body text. + path = self._write( + "--! name\nc\n--! script\nscripts/a.py\n--!expected_result\nx\n!-- end\n" + ) + with self.assertRaisesRegex(AssertionError, "malformed section marker"): + parse_test_file(path) + + def test_parse_rejects_stray_content_outside_section(self): + path = self._write("stray text\n--! name\nc\n--! script\nscripts/a.py\n!-- end\n") + with self.assertRaisesRegex(AssertionError, "content outside any section"): + parse_test_file(path) + + def test_parse_allows_blank_lines_between_blocks(self): + # Blank separators outside sections are fine (not stray content). + path = self._write( + "--! name\nc1\n--! script\nscripts/a.py\n!-- end\n\n\n" + "--! name\nc2\n--! script\nscripts/b.py\n!-- end\n" + ) + _, cases = parse_test_file(path) + self.assertEqual([c["name"] for c in cases], ["c1", "c2"]) + + def test_parse_unterminated_trailing_case(self): + # Last case missing "!-- end". + path = self._write("--! name\nc\n--! script\nscripts/a.py\n") + # Lenient by default (regeneration rewrites with terminators): + _, cases = parse_test_file(path) + self.assertEqual(len(cases), 1) + # Strict (verify mode) rejects it: + with self.assertRaisesRegex(AssertionError, "does not end with"): + parse_test_file(path, require_terminated=True) + + def test_format_value_refuses_tab_or_newline_in_string(self): + try: + from pyspark.sql.types import StringType + except Exception: + self.skipTest("pyspark.sql.types unavailable in this environment") + from pyspark.testing.df_golden import _format_value + + self.assertEqual(_format_value("ok", StringType()), "ok") + for bad in ("a\tb", "a\nb"): + with self.assertRaisesRegex(AssertionError, "tab or newline"): + _format_value(bad, StringType()) + + # -- regeneration ---------------------------------------------------- + + def test_regenerate_replaces_expected_sections_and_carries_identity(self): + old = { + "name": "c", + "tags": "unordered", + "script": "scripts/a.py", + "expected_optimized_output": "old base", + "expected_output_schema": "old schema", + } + actual = { + "expected_optimized_output": "new base", + "expected_output_schema": "new schema", + } + new = _regenerate_case(old, actual) + # Name / tags / script are carried over; expected_* come from this run. + self.assertEqual(new["name"], "c") + self.assertEqual(new["tags"], "unordered") + self.assertEqual(new["script"], "scripts/a.py") + self.assertEqual(new["expected_optimized_output"], "new base") + self.assertEqual(new["expected_output_schema"], "new schema") + + +if __name__ == "__main__": + from pyspark.testing import main + + main() diff --git a/python/pyspark/testing/df_golden.py b/python/pyspark/testing/df_golden.py new file mode 100644 index 0000000000000..69b9c77fc225c --- /dev/null +++ b/python/pyspark/testing/df_golden.py @@ -0,0 +1,746 @@ +# +# 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. +# + +""" +Framework for DataFrame API golden file tests, analogous to SQLQueryTestSuite for SQL. + +A test is described by a ``.test`` file which doubles as the golden file: the +expected outputs are stored inline and rewritten in place when golden files are +regenerated (``SPARK_GENERATE_GOLDEN_FILES=1``). + +``.test`` file format:: + + --! name + __file_metadata__ + --! source + df_golden/group_by + !-- end + + + --! name + range + select + filter + order + --! script + scripts/group_by/range_select.py + --! expected_analysis_output + Sort [k#x ASC NULLS FIRST], true + +- ... + --! expected_optimized_output + ... + --! expected_output_schema + struct + --! expected_result + +---+ + | k | + +---+ + | 1 | + +---+ + printed all 1 rows. + --! expected_result_hash + + !-- end + +The first block may be named ``__file_metadata__``; its remaining sections +(e.g. ``source``) are file-level metadata, matching the convention used by the +Scala ``SqlHiFiTestRunner`` framework. Each test case references a standalone +Python script (path relative to the ``.test`` file) that is executed with +``spark`` in scope and must assign the DataFrame under test to a variable +named ``df``. Cases run in file order against the same session, so earlier +cases can set up temp views for later ones. + +Sections: + +- ``name``: human-readable test case name (required). +- ``tags``: optional, whitespace/comma separated. Row order is asserted by + default; add the ``unordered`` tag to sort result rows before comparison for + cases whose result has no deterministic order (aggregate/join/distinct/... + without a global sort). +- ``script``: path to the Python script (required). +- ``expected_analysis_output``: the analyzed logical plan. +- ``expected_optimized_output``: the optimized logical plan. +- ``expected_output_schema``: ``df.schema.simpleString()``. +- ``expected_result``: pretty-printed result table plus a ``printed all N + rows.`` trailer. +- ``expected_result_hash``: sha256 over the (normalized, post-sort) result + rows -- a compact checksum of the same rows rendered in ``expected_result`` + (the table is not truncated), co-required with it. +- ``expected_error``: expected error message when analysis or execution fails. + Mutually exclusive with the plan/schema/result sections: as in the SQL golden + suite, any error records only the message and discards the plan and schema. + +At comparison time only the ``expected_*`` sections present in the file are +checked, so optional sections (e.g. ``expected_optimized_output``) may be +omitted. Regeneration writes all sections the case produces. +""" + +import hashlib +import math +import os +import re +from decimal import Decimal + + +_CASE_END = "!-- end" +_SECTION_PREFIX = "--! " +_FILE_METADATA_NAME = "__file_metadata__" + +# Canonical section order used when (re)generating a ``.test`` file. +_CASE_SECTION_ORDER = [ + "name", + "tags", + "script", + "expected_analysis_output", + "expected_optimized_output", + "expected_output_schema", + "expected_result", + "expected_result_hash", + "expected_error", +] + +_RESULT_SECTIONS = [s for s in _CASE_SECTION_ORDER if s.startswith("expected_")] + +_KNOWN_HEADER_SECTIONS = {"source"} +_KNOWN_TAGS = {"unordered"} + + +# --------------------------------------------------------------------------- +# .test file parsing / serialization +# --------------------------------------------------------------------------- + + +def parse_test_file(filepath, require_terminated=False): + """ + Parse a ``.test`` file. + + When *require_terminated* is set (verify mode), a file whose last case is + missing its ``!-- end`` terminator is rejected: in verify mode an unclosed + final case is corruption (e.g. a truncating bad merge) that could otherwise + pass by matching a partial case or silently merging two. Regeneration leaves + it unset and stays lenient, since it rewrites the file with terminators. + + Returns + ------- + header : dict + File-level metadata sections from the ``__file_metadata__`` block + (excluding ``name``), e.g. ``{"source": ...}``. + cases : list[dict] + One dict per test case, mapping section name to content. + """ + with open(filepath, "r") as f: + lines = f.read().split("\n") + + cases = [] + current = None + section_key = None + section_lines = [] + + def flush(): + nonlocal section_key, section_lines + if section_key is not None and current is not None: + # A repeated section name is a copy/paste or merge mistake; last-wins + # would silently discard one of the two, so fail loudly instead. + assert section_key not in current, "{}: duplicate section `{}`".format( + filepath, section_key + ) + current[section_key] = "\n".join(section_lines).strip("\n") + section_key = None + section_lines = [] + + for line in lines: + stripped = line.rstrip() + if stripped == _CASE_END: + flush() + if current: + cases.append(current) + current = None + elif stripped.startswith(_SECTION_PREFIX): + flush() + if current is None: + current = {} + section_key = stripped[len(_SECTION_PREFIX) :].strip() + elif stripped.startswith("--!"): + # Reaches here only because the space after "--!" is missing, i.e. a + # typo'd section marker. Left as body it would silently turn an + # assertion into inert prose, so reject it. + raise AssertionError( + "{}: malformed section marker (expected `{}`): {!r}".format( + filepath, _SECTION_PREFIX, line + ) + ) + elif section_key is not None: + section_lines.append(line) + elif stripped: + # Non-blank content outside any section (before the first marker or + # between cases) is dropped by the original loop; that hides stray + # text, so fail loudly. Blank separator lines are fine. + raise AssertionError("{}: content outside any section: {!r}".format(filepath, line)) + + # A case still open here never hit "!-- end". In verify mode that is + # corruption; under regeneration stay lenient and keep it so the rewrite can + # fix the formatting. + flush() + if current: + assert not require_terminated, ( + "{}: file does not end with `{}` (last case is unterminated)".format( + filepath, _CASE_END + ) + ) + cases.append(current) + + header = {} + if cases and cases[0].get("name") == _FILE_METADATA_NAME: + header = cases.pop(0) + del header["name"] + + return header, cases + + +def write_test_file(filepath, header, cases): + """Serialize *header* and *cases* back into ``.test`` file format.""" + blocks = [] + if header: + header_lines = [_SECTION_PREFIX + "name", _FILE_METADATA_NAME] + for key, value in header.items(): + header_lines.append(_SECTION_PREFIX + key) + header_lines.append(value) + header_lines.append(_CASE_END) + blocks.append("\n".join(header_lines)) + + for case in cases: + case_lines = [] + for key in _CASE_SECTION_ORDER: + value = case.get(key) + if value is not None: + case_lines.append(_SECTION_PREFIX + key) + case_lines.append(value) + case_lines.append(_CASE_END) + blocks.append("\n".join(case_lines)) + + with open(filepath, "w") as f: + f.write("\n\n\n".join(blocks) + "\n") + + +def parse_tags(case): + """Return the set of tags declared on *case*.""" + return {tag for tag in re.split(r"[,\s]+", case.get("tags", "")) if tag} + + +# --------------------------------------------------------------------------- +# Output normalisation (mirrors SQLQueryTestHelper.replaceNotIncludedMsg) +# --------------------------------------------------------------------------- + +# Compiled once for performance. +_NORMALIZATION_RULES = [ + (re.compile(r"#\d+"), "#x"), + (re.compile(r"plan_id=\d+"), "plan_id=x"), + (re.compile(r"joinId=\d+"), "joinId=x"), + (re.compile(r"repartitionId=\d+"), "repartitionId=x"), + (re.compile(r"uuid\(Some\(-?\d+\)\)"), "uuid(Some(x))"), + (re.compile(r"CTERelationDef \d+,"), "CTERelationDef xxxx,"), + (re.compile(r"CTERelationRef \d+,"), "CTERelationRef xxxx,"), + (re.compile(r"cterelationdef \d+,"), "cterelationdef xxxx,"), + (re.compile(r"cterelationref \d+,"), "cterelationref xxxx,"), + (re.compile(r"UnionLoop \d+"), "UnionLoop xxxx"), + (re.compile(r"UnionLoopRef \d+,"), "UnionLoopRef xxxx,"), + (re.compile(r"Loop id: \d+"), "Loop id: xxxx"), + (re.compile(r"@\w*,"), "@xxxxxxxx,"), + (re.compile(r"\*\(\d+\) "), "*"), +] + + +def replace_not_included(text): + """Normalise environment-dependent fragments in *text*.""" + for pattern, repl in _NORMALIZATION_RULES: + text = pattern.sub(repl, text) + return text + + +def format_error(e): + """ + Format an exception message for golden file comparison. + + Uses ``str(e)``, which for connect exceptions is the server-side message + (``[ERROR_CLASS] message SQLSTATE: xxxxx``). Stripped to keep the output + deterministic: + + - the appended JVM stacktrace; + - the ``== DataFrame ==`` query context block, which embeds the absolute + script path and line number of the DataFrame call (editing a script + comment must not break golden files); + - the trailing logical plan dump (a ``;\\n`` followed by the plan tree); + + and expression ids are normalized. + """ + msg = str(e) + msg = msg.split("\n\nJVM stacktrace:")[0] + msg = msg.split("\n== DataFrame ==")[0] + # Drop a trailing logical-plan dump: Spark appends it as ";\n" followed by + # the plan tree, whose root line starts with an (optionally "'"-prefixed) + # uppercase operator name. Anchoring on that lookahead avoids truncating a + # message body that merely contains ";\n" (splitting on the first ";\n" + # unconditionally would lose the remainder of such a message). + msg = re.split(r";\n(?=['A-Z])", msg, maxsplit=1)[0] + return replace_not_included(msg).strip() + + +# --------------------------------------------------------------------------- +# Plan extraction +# --------------------------------------------------------------------------- + + +_EXPLAIN_HEADER = re.compile(r"^== .+ ==$", re.MULTILINE) + + +def _extract_explain_section(explain, marker): + """ + Return the body of the *marker* section of an extended explain output, + ending at the next ``== ... ==`` header. + """ + start = explain.find(marker) + if start < 0: + return None + start = explain.find("\n", start + len(marker)) + if start < 0: + return None + start += 1 + match = _EXPLAIN_HEADER.search(explain, start) + end = match.start() if match else len(explain) + return explain[start:end].strip("\n") + + +def get_plan_strings(df): + """ + Return ``(analyzed, optimized)`` normalized logical plan strings. + + Uses ``df._explain_string(mode="extended")``, which exists on Spark + Connect only - the framework runs over connect (see ``DFGoldenTestBase``). + Triggers analysis, so analysis errors surface here. + """ + explain = df._explain_string(mode="extended") + analyzed = _extract_explain_section(explain, "== Analyzed Logical Plan ==") + optimized = _extract_explain_section(explain, "== Optimized Logical Plan ==") + if analyzed is None: + raise AssertionError("explain output has no analyzed plan section:\n" + explain) + + # When the output schema is non-empty, the analyzed section starts with a + # schema header line (possibly truncated by spark.sql.debug.maxToStringFields). + # The schema has its own golden section, so drop the header by position. + if df.schema.fields: + analyzed = "\n".join(analyzed.split("\n")[1:]).strip("\n") + + optimized = replace_not_included(optimized) if optimized is not None else None + return replace_not_included(analyzed), optimized + + +# --------------------------------------------------------------------------- +# Result formatting +# --------------------------------------------------------------------------- + + +def format_double(value): + """ + Render *value* (a Python ``float`` from a ``double`` column) exactly as Java + ``Double.toString`` does, which is what Hive output uses for numeric types + (``HiveResult.toHiveStringDefault``: ``case (n, _: NumericType) => + n.toString``). Matching it keeps double results comparable with the SQL + ``.sql.out`` goldens; Python's own ``str``/``repr`` differs for special + values (``nan``/``inf``) and for the scientific-notation regime. + + Java's rules: ``NaN``/``Infinity``/``-Infinity`` spelled out; a signed + ``0.0``; plain decimal (always with a fractional digit) when + ``1e-3 <= |x| < 1e7``; otherwise ``d.ddddEexp`` scientific notation with a + single leading digit. The shortest round-tripping digits come from Python's + ``repr`` (normalized to drop the artificial trailing zero of values like + ``1e7`` -> ``10000000.0``); only their placement is reformatted. + """ + if math.isnan(value): + return "NaN" + if math.isinf(value): + return "Infinity" if value > 0 else "-Infinity" + if value == 0.0: + return "-0.0" if math.copysign(1.0, value) < 0 else "0.0" + + sign = "-" if value < 0 else "" + digit_tuple, exp = Decimal(repr(abs(value))).normalize().as_tuple()[1:] + digits = "".join(map(str, digit_tuple)) + nd = len(digits) + # Power of ten of the leading significant digit. + leading_exp = exp + nd - 1 + + if -3 <= leading_exp < 7: + if leading_exp >= 0: + if nd <= leading_exp + 1: + body = digits + "0" * (leading_exp + 1 - nd) + ".0" + else: + body = digits[: leading_exp + 1] + "." + digits[leading_exp + 1 :] + else: + body = "0." + "0" * (-leading_exp - 1) + digits + else: + body = digits[0] + "." + (digits[1:] or "0") + "E" + str(leading_exp) + return sign + body + + +def _format_value(value, data_type, nested=False): + """ + Format a single cell value for golden file output, mirroring + ``HiveResult.toHiveStringDefault`` so values line up with the SQL + ``.sql.out`` goldens: structs carry quoted field names, strings are quoted + when nested, and a top-level null (``NULL``) differs from a nested one + (``null``). + """ + from pyspark.sql.types import ( + ArrayType, + BinaryType, + BooleanType, + DateType, + DecimalType, + DoubleType, + FloatType, + MapType, + StringType, + StructType, + TimestampNTZType, + TimestampType, + ) + + if value is None: + return "null" if nested else "NULL" + + if isinstance(data_type, BooleanType): + return "true" if value else "false" + if isinstance(data_type, StringType): + # A tab or newline in a cell would desync the rendered table (cells are + # tab-joined and re-split, the file is newline-delimited) while the hash + # stayed self-consistent, so --verify could not catch the misrender. + # Refuse loudly rather than bake a corrupt golden; add escaping with the + # first case that legitimately needs such a value. + if "\t" in value or "\n" in value: + raise AssertionError( + "df_golden: result string contains a tab or newline, which is " + "not supported yet (would desync the rendered table): {!r}".format(value) + ) + return '"' + value + '"' if nested else value + if isinstance(data_type, DecimalType): + # BigDecimal.toPlainString: never scientific notation, scale preserved. + return format(value, "f") + if isinstance(data_type, DoubleType): + return format_double(value) + if isinstance(data_type, StructType): + parts = [ + '"{}":{}'.format(f.name, _format_value(value[i], f.dataType, nested=True)) + for i, f in enumerate(data_type.fields) + ] + return "{" + ",".join(parts) + "}" + if isinstance(data_type, ArrayType): + parts = [_format_value(v, data_type.elementType, nested=True) for v in value] + return "[" + ",".join(parts) + "]" + if isinstance(data_type, MapType): + parts = [ + _format_value(k, data_type.keyType, nested=True) + + ":" + + _format_value(v, data_type.valueType, nested=True) + for k, v in value.items() + ] + # Hive sorts map entries by their rendered string, not by key. + return "{" + ",".join(sorted(parts)) + "}" + # These types have no faithful ``str()`` rendering and must not fall through + # to the generic branch below: + # - float: Python's repr is the double-precision shortest form, not Java + # ``Float.toString``'s float32-shortest form, so str() would diverge. + # (``double`` is handled above via ``format_double``; ``float`` waits for + # the first float-column case, which needs float32-shortest rendering.) + # - temporal/binary: need a Hive-style formatter and (for LTZ timestamps) a + # pinned session time zone this framework does not set up yet. + # Refuse them loudly rather than silently emit a wrong/non-deterministic + # value; add real formatting together with the first such test case. + if isinstance( + data_type, + (FloatType, DateType, TimestampType, TimestampNTZType, BinaryType), + ): + raise AssertionError( + "df_golden: result column of type {} is not supported yet (needs a " + "Hive-style formatter)".format(data_type.simpleString()) + ) + return str(value) + + +def get_result_rows(df): + """ + Collect *df* and format each row as a tab-separated string matching hive + output conventions (``NULL`` for None, lowercase booleans, etc.). + + Cells are joined with ``\\t`` and later re-split on ``\\t`` by + ``render_result_table``, and the ``.test`` format is newline-delimited, so a + literal tab or newline inside a string value would desync the rendered table + while the hash stayed self-consistent (``--verify`` could not flag it). + ``_format_value`` therefore rejects such strings loudly rather than let a + corrupt golden through. + """ + schema = df.schema + return [ + "\t".join(_format_value(row[i], field.dataType) for i, field in enumerate(schema.fields)) + for row in df.collect() + ] + + +def render_result_table(columns, rows): + """ + Render *rows* (tab-separated strings) as a pretty-printed table:: + + +----+----+ + | c1 | c2 | + +----+----+ + | 1 | 10 | + +----+----+ + printed all 1 rows. + """ + trailer = "printed all {} rows.".format(len(rows)) + if not columns: + return trailer + + cells = [r.split("\t") for r in rows] + widths = [len(c) for c in columns] + for row_cells in cells: + for i, cell in enumerate(row_cells[: len(widths)]): + widths[i] = max(widths[i], len(cell)) + + border = "+" + "+".join("-" * (w + 2) for w in widths) + "+" + + def fmt(values): + padded = [v.ljust(w) for v, w in zip(values, widths)] + return "| " + " | ".join(padded) + " |" + + lines = [border, fmt(columns), border] + lines.extend(fmt(row_cells) for row_cells in cells) + lines.append(border) + lines.append(trailer) + return "\n".join(lines) + + +def hash_result_rows(rows): + """sha256 over the normalized result rows; verifies the full result.""" + return hashlib.sha256("\n".join(rows).encode("utf-8")).hexdigest() + + +# --------------------------------------------------------------------------- +# Test execution engine +# --------------------------------------------------------------------------- + + +def run_script(spark, script_path): + """ + Execute the test case script and return the DataFrame it assigns to ``df``. + + The script runs with ``spark`` in scope and is responsible for its own + imports. + """ + with open(script_path, "r") as f: + code = f.read() + namespace = {"spark": spark} + exec(compile(code, script_path, "exec"), namespace) + if "df" not in namespace: + raise AssertionError( + "Test script {} must assign a DataFrame to a variable named `df`".format(script_path) + ) + return namespace["df"] + + +def compute_case_outputs(spark, case, base_dir): + """ + Run a single test case and return a dict of actual ``expected_*`` sections. + """ + from pyspark.errors import PySparkException + + tags = parse_tags(case) + script_path = os.path.join(base_dir, case["script"]) + + # Only Spark errors are legitimate expected outputs. Anything else + # (NameError, ImportError, ... from a buggy script) must fail the test; + # capturing it would write the Python error into the golden file as the + # expected output on regeneration. + try: + df = run_script(spark, script_path) + analyzed, optimized = get_plan_strings(df) + schema = df.schema.simpleString() + except PySparkException as e: + return {"expected_error": format_error(e)} + + actual = { + "expected_analysis_output": analyzed, + "expected_output_schema": schema, + } + if optimized is not None: + actual["expected_optimized_output"] = optimized + + try: + rows = get_result_rows(df) + except PySparkException as e: + # Match the SQL golden suite: on any error keep only the message and + # discard the analyzed plan / schema captured before execution. + return {"expected_error": format_error(e)} + + rows = [replace_not_included(r) for r in rows] + # Sort the rows only when the case is explicitly tagged ``unordered``. Row + # order is asserted by default; a case whose result has no deterministic + # order (aggregate/join/distinct/... without a global sort) must opt out via + # the tag. Deriving orderedness from the rendered plan text was rejected as + # too loose: it silently sorts genuinely order-sensitive results, hiding real + # ordering regressions from the golden. + if "unordered" in tags: + rows = sorted(rows) + actual["expected_result"] = render_result_table(df.columns, rows) + actual["expected_result_hash"] = hash_result_rows(rows) + return actual + + +def _validate_test_file(test_file, header, cases, regenerate): + """ + Fail loudly on malformed ``.test`` content. A misspelled section or tag + that is silently ignored makes a case assert less than it appears to (or + nothing at all), so unknown names are errors, not noise. + """ + unknown_header = set(header) - _KNOWN_HEADER_SECTIONS + assert not unknown_header, "{}: unknown header sections: {}".format( + test_file, ", ".join(sorted(unknown_header)) + ) + assert cases, "{}: no test cases found".format(test_file) + for case in cases: + assert case.get("name"), "{}: every test case needs a name".format(test_file) + name = case["name"] + assert case.get("script"), "{}: case `{}` needs a script".format(test_file, name) + # Unknown sections are dropped and rewritten by regeneration, so only + # reject them in verify mode. Enforcing this during regeneration would + # block the very migration regeneration exists to perform: a section + # renamed or removed in the framework (e.g. the old + # ``expected_analysis_error``/``expected_execution_error`` split folded + # into ``expected_error``) leaves the on-disk file carrying a name no + # longer in ``_CASE_SECTION_ORDER`` until it is regenerated. + if not regenerate: + unknown = set(case) - set(_CASE_SECTION_ORDER) + assert not unknown, "{}: case `{}` has unknown sections: {}".format( + test_file, name, ", ".join(sorted(unknown)) + ) + # Tags are preserved verbatim across regeneration, so an unknown tag + # would persist; reject it in both modes. + unknown_tags = parse_tags(case) - _KNOWN_TAGS + assert not unknown_tags, "{}: case `{}` has unknown tags: {}".format( + test_file, name, ", ".join(sorted(unknown_tags)) + ) + # In regenerate mode new cases legitimately have no expected_* + # sections yet; in verify mode such a case would pass vacuously. + if not regenerate: + assert any(case.get(key) is not None for key in _RESULT_SECTIONS), ( + "{}: case `{}` has no expected_* sections and would assert " + "nothing; regenerate the golden files".format(test_file, name) + ) + # ``_compare_case`` only checks sections present in the file, so a + # dropped section (merge/manual edit) silently shrinks coverage + # without failing. Pin down what a well-formed case must look like: + has_error = case.get("expected_error") is not None + has_result = case.get("expected_result") is not None + has_hash = case.get("expected_result_hash") is not None + if has_error: + # An error case records only the error (the run discards plan, + # schema and result on failure); anything else is a corrupt file. + conflicting = sorted( + key + for key in _RESULT_SECTIONS + if key != "expected_error" and case.get(key) is not None + ) + assert not conflicting, ( + "{}: error case `{}` must carry only `expected_error`, not also: {}".format( + test_file, name, ", ".join(conflicting) + ) + ) + else: + # The result table and its hash are a pair; dropping one leaves + # the other asserting half the result, so require both or neither. + assert has_result == has_hash, ( + "{}: case `{}` must have both `expected_result` and " + "`expected_result_hash` or neither".format(test_file, name) + ) + + +def run_golden_test(test_case, spark, test_file): + """ + Run all cases of a ``.test`` file. + + Parameters + ---------- + test_case : unittest.TestCase + The test case instance (for assertions). + spark : SparkSession + The session to run against. The caller provides a fresh session per + ``.test`` file (the connect counterpart of ``SQLQueryTestSuite``'s + per-file ``newSession()``), so state created by case scripts - temp + views, UDFs, confs - is discarded with the session and cannot leak + into other files. + test_file : str + Absolute path to the ``.test`` file. + """ + regenerate = os.environ.get("SPARK_GENERATE_GOLDEN_FILES") is not None + base_dir = os.path.dirname(test_file) + + header, cases = parse_test_file(test_file, require_terminated=not regenerate) + _validate_test_file(test_file, header, cases, regenerate) + + # Golden files are generated with ANSI mode on, matching the SQL golden + # tests. The session is discarded after the file, so nothing to restore. + spark.conf.set("spark.sql.ansi.enabled", "true") + + regenerated_cases = [] + for case in cases: + actual = compute_case_outputs(spark, case, base_dir) + if regenerate: + regenerated_cases.append(_regenerate_case(case, actual)) + else: + _compare_case(test_case, case, actual) + + if regenerate: + write_test_file(test_file, header, regenerated_cases) + + +def _regenerate_case(old_case, actual): + """ + Build the regenerated form of *old_case* from this run's *actual* outputs. + + Every populated ``expected_*`` section from this run replaces the on-disk + value; ``name`` / ``tags`` / ``script`` are carried over unchanged so the + case's identity, ordering guard, and script pointer survive regeneration. + """ + carried = {key: old_case.get(key) for key in ("name", "tags", "script")} + carried.update(actual) + return carried + + +def _compare_case(test_case, case, actual): + """Compare the ``expected_*`` sections of *case* against *actual*.""" + name = case["name"] + for key in _RESULT_SECTIONS: + expected = case.get(key) + if expected is None: + continue + got = actual.get(key) + if got is None: + produced = ", ".join(sorted(actual)) or "" + test_case.fail( + "[{}] expected section `{}` but the case produced: {}".format(name, key, produced) + ) + test_case.assertEqual( + expected.strip("\n"), + got.strip("\n"), + "[{}] mismatch in `{}`".format(name, key), + )