From 81ce6b63e110e2878dae10b437ca5462ad770ac5 Mon Sep 17 00:00:00 2001 From: Daniel Tenedorio Date: Fri, 31 Jul 2026 13:43:09 -0700 Subject: [PATCH] commit --- dev/.rat-excludes | 6 - dev/sparktestsupport/modules.py | 2 +- pyproject.toml | 3 - .../pyspark/sql/tests/df_golden/df_golden.py | 551 +++++++++++++----- .../pyspark/sql/tests/df_golden/group_by.test | 184 ++---- .../pyspark/sql/tests/df_golden/regenerate.sh | 27 +- .../scripts/group_by/agg_alias_filter.py | 5 - .../group_by/agg_counts_no_grouping.py | 3 - .../scripts/group_by/agg_stats_with_nulls.py | 27 - .../scripts/group_by/bool_agg_all_nulls.py | 9 - .../scripts/group_by/bool_agg_empty_table.py | 9 - .../scripts/group_by/bool_agg_group_by.py | 9 - .../scripts/group_by/bool_agg_having_false.py | 11 - .../scripts/group_by/bool_agg_having_null.py | 9 - .../group_by/bool_agg_null_filtering.py | 9 - .../group_by/bool_and_decimal_error.py | 6 - .../scripts/group_by/bool_and_window.py | 6 - .../scripts/group_by/bool_or_double_error.py | 5 - .../scripts/group_by/bool_or_int_error.py | 5 - .../scripts/group_by/bool_or_window.py | 6 - .../group_by/bool_or_window_repeated.py | 6 - .../scripts/group_by/count_if_group_by.py | 3 - .../scripts/group_by/empty_input_agg.py | 5 - .../group_by/empty_input_agg_select.py | 3 - .../scripts/group_by/empty_input_group_by.py | 5 - .../scripts/group_by/every_int_error.py | 5 - .../scripts/group_by/every_string.py | 5 - .../scripts/group_by/every_window.py | 8 - .../scripts/group_by/group_by_complex_expr.py | 5 - .../scripts/group_by/group_by_count.py | 5 - .../scripts/group_by/group_by_literal.py | 5 - .../group_by/group_by_literal_hash_agg.py | 10 - .../group_by/group_by_literal_sort_agg.py | 5 - .../group_by/group_by_missing_aggregation.py | 3 - .../group_by/group_by_multiple_counts.py | 3 - .../scripts/group_by/group_by_struct.py | 5 - .../group_by/grouping_exprs_not_optimized.py | 10 - .../scripts/group_by/having_agg_count.py | 10 - .../scripts/group_by/having_max_v.py | 9 - .../group_by/pull_out_grouping_exprs.py | 10 - .../group_by/select_non_grouping_column.py | 5 - .../scripts/group_by/setup_test_agg.py | 8 - .../scripts/group_by/setup_test_data.py | 6 - .../scripts/group_by/some_int_error.py | 5 - .../df_golden/scripts/group_by/some_window.py | 6 - .../sql/tests/df_golden/test_df_golden.py | 143 ----- .../df_golden/test_df_golden_framework.py | 292 +++++----- .../sql/tests/df_golden/test_group_by.py | 324 ++++++++++ 48 files changed, 931 insertions(+), 870 deletions(-) delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/agg_alias_filter.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/agg_counts_no_grouping.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/agg_stats_with_nulls.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_all_nulls.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_empty_table.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_group_by.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_having_false.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_having_null.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_null_filtering.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/bool_and_decimal_error.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/bool_and_window.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/bool_or_double_error.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/bool_or_int_error.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/bool_or_window.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/bool_or_window_repeated.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/count_if_group_by.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/empty_input_agg.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/empty_input_agg_select.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/empty_input_group_by.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/every_int_error.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/every_string.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/every_window.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_complex_expr.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_count.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_literal.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_literal_hash_agg.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_literal_sort_agg.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_missing_aggregation.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_multiple_counts.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_struct.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/grouping_exprs_not_optimized.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/having_agg_count.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/having_max_v.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/pull_out_grouping_exprs.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/select_non_grouping_column.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/setup_test_agg.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/setup_test_data.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/some_int_error.py delete mode 100644 python/pyspark/sql/tests/df_golden/scripts/group_by/some_window.py delete mode 100644 python/pyspark/sql/tests/df_golden/test_df_golden.py create mode 100644 python/pyspark/sql/tests/df_golden/test_group_by.py diff --git a/dev/.rat-excludes b/dev/.rat-excludes index 4aee210e85b98..d987558f6b638 100644 --- a/dev/.rat-excludes +++ b/dev/.rat-excludes @@ -145,9 +145,3 @@ testCommitLog 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 898a3831f484b..e1d97e10644e0 100644 --- a/dev/sparktestsupport/modules.py +++ b/dev/sparktestsupport/modules.py @@ -639,8 +639,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", + "pyspark.sql.tests.df_golden.test_group_by", ], ) diff --git a/pyproject.toml b/pyproject.toml index 35de692cf5b0d..85b3a5c189a10 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,9 +59,6 @@ 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"] [dependency-groups] _py4j = ["py4j>=0.10.9.9"] diff --git a/python/pyspark/sql/tests/df_golden/df_golden.py b/python/pyspark/sql/tests/df_golden/df_golden.py index 69b9c77fc225c..53f3b946bee15 100644 --- a/python/pyspark/sql/tests/df_golden/df_golden.py +++ b/python/pyspark/sql/tests/df_golden/df_golden.py @@ -18,23 +18,48 @@ """ 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``). +A golden test is a pair of files: a Python test module declaring the cases, and +a ``.test`` file next to it holding their expected outputs. Everything a case +needs in order to run -- the DataFrame program, its name, its tags -- lives in +the Python class, so cases are ordinary code that can be read, imported and +stepped through; the ``.test`` file is purely generated output, rewritten in +place when golden files are regenerated (``SPARK_GENERATE_GOLDEN_FILES=1``). + +Test module:: + + class GroupByGoldenTests(DFGoldenTestMixin, ReusedConnectTestCase): + golden_file = "group_by.test" + + @classmethod + def setup_session(cls, spark): + spark.sql("CREATE OR REPLACE TEMPORARY VIEW testData AS ...") + + @unordered + def _test_group_by_count(self, spark): + \"\"\"Aggregate with non-empty GroupBy expressions.\"\"\" + return spark.table("testData").groupBy(col("a")).agg(count(col("b"))) + +Each ``_test_`` method builds and returns the DataFrame of one case, and +:class:`DFGoldenTestMixin` registers a ``test_`` method for it. Cases are +therefore ordinary unittest tests: they are reported individually and can be run +one at a time:: + + python/run-tests --testnames \\ + "pyspark.sql.tests.df_golden.test_group_by GroupByGoldenTests.test_group_by_count" ``.test`` file format:: --! name __file_metadata__ --! source - df_golden/group_by + pyspark.sql.tests.df_golden.test_group_by.GroupByGoldenTests !-- end --! name - range + select + filter + order - --! script - scripts/group_by/range_select.py + group_by_count + --! tags + unordered --! expected_analysis_output Sort [k#x ASC NULLS FIRST], true +- ... @@ -53,22 +78,20 @@ !-- 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. +The first block is named ``__file_metadata__``; its remaining sections are +file-level metadata, matching the convention used by the Scala +``SqlHiFiTestRunner`` framework. ``source`` records the class the file is +generated from, so the ``.test`` file always points back at the code that +produces it. 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). +- ``name``: the case method name without its ``_test_`` prefix, which is what + ties a block to the method that produced it (required). +- ``tags``: whitespace/comma separated, written from the case method's + decorators. Row order is asserted by default; ``@unordered`` sorts result + rows before comparison, for cases whose result has no deterministic order + (aggregate/join/distinct/... without a global sort). - ``expected_analysis_output``: the analyzed logical plan. - ``expected_optimized_output``: the optimized logical plan. - ``expected_output_schema``: ``df.schema.simpleString()``. @@ -84,9 +107,33 @@ 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. + +Running the tests:: + + python/run-tests --testnames pyspark.sql.tests.df_golden.test_group_by + +Regenerating golden files +------------------------- +Set ``SPARK_GENERATE_GOLDEN_FILES=1`` before running the tests, or use the +wrapper script, which covers every golden test module in this directory:: + + python/pyspark/sql/tests/df_golden/regenerate.sh [--verify] + +With ``--verify`` the wrapper re-runs the tests afterwards against the +regenerated files. Regeneration rewrites a whole file at once, so it has to +run a whole test class: a filtered run refuses to write rather than drop the +cases it skipped. + +Adding golden tests +------------------- +Add a case method to an existing class, or add a ``test_.py`` module +with a new class (registering it in ``dev/sparktestsupport/modules.py``, as for +any PySpark test module). Then regenerate: the ``.test`` file is created or +extended from the class, so it needs no hand-editing. """ import hashlib +import inspect import math import os import re @@ -101,7 +148,6 @@ _CASE_SECTION_ORDER = [ "name", "tags", - "script", "expected_analysis_output", "expected_optimized_output", "expected_output_schema", @@ -115,21 +161,24 @@ _KNOWN_HEADER_SECTIONS = {"source"} _KNOWN_TAGS = {"unordered"} +# Methods with this prefix declare golden cases; the mixin registers a real +# ``test_`` method for each one. The prefix keeps them out of unittest's +# own collection, which only picks up names starting with ``test``. +_CASE_METHOD_PREFIX = "_test_" + # --------------------------------------------------------------------------- # .test file parsing / serialization # --------------------------------------------------------------------------- -def parse_test_file(filepath, require_terminated=False): +def parse_test_file(filepath): """ 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. + A file whose last case is missing its ``!-- end`` terminator is rejected: + 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. Returns ------- @@ -188,17 +237,11 @@ def flush(): # 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. + # A case still open here never hit "!-- end", i.e. the file was truncated. flush() - if current: - assert not require_terminated, ( - "{}: file does not end with `{}` (last case is unterminated)".format( - filepath, _CASE_END - ) - ) - cases.append(current) + assert not current, "{}: file does not end with `{}` (last case is unterminated)".format( + filepath, _CASE_END + ) header = {} if cases and cases[0].get("name") == _FILE_METADATA_NAME: @@ -238,6 +281,26 @@ def parse_tags(case): return {tag for tag in re.split(r"[,\s]+", case.get("tags", "")) if tag} +def unordered(method): + """ + Mark a case method whose result rows have no deterministic order. + + Row order is asserted by default; a case whose result has no deterministic + order (aggregate/join/distinct/... without a global sort) opts out with this + decorator, which sorts the rows before they are compared or written. + 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. + """ + method.df_golden_unordered = True + return method + + +def is_unordered(method): + """Whether *method* is decorated with :func:`unordered`.""" + return getattr(method, "df_golden_unordered", False) + + # --------------------------------------------------------------------------- # Output normalisation (mirrors SQLQueryTestHelper.replaceNotIncludedMsg) # --------------------------------------------------------------------------- @@ -326,7 +389,7 @@ 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``). + Connect only - the framework runs over connect (see ``DFGoldenTestMixin``). Triggers analysis, so analysis errors surface here. """ explain = df._explain_string(mode="extended") @@ -540,39 +603,36 @@ def hash_result_rows(rows): # --------------------------------------------------------------------------- -def run_script(spark, script_path): - """ - Execute the test case script and return the DataFrame it assigns to ``df``. +def _build_dataframe(spark, build_df): + """Call a case method and return the DataFrame it built.""" + from pyspark.sql import DataFrame - 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: + df = build_df(spark) + if not isinstance(df, DataFrame): raise AssertionError( - "Test script {} must assign a DataFrame to a variable named `df`".format(script_path) + "case method {} must return the DataFrame under test, got: {!r}".format( + getattr(build_df, "__name__", build_df), df + ) ) - return namespace["df"] + return df -def compute_case_outputs(spark, case, base_dir): +def compute_case_outputs(spark, build_df, sort_rows=False): """ Run a single test case and return a dict of actual ``expected_*`` sections. + + *build_df* is the case method: it takes the session and returns the + DataFrame under test. *sort_rows* sorts the result rows before they are + rendered (see :func:`unordered`). """ 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; + # (NameError, TypeError, ... from a buggy case method) 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) + df = _build_dataframe(spark, build_df) analyzed, optimized = get_plan_strings(df) schema = df.schema.simpleString() except PySparkException as e: @@ -593,24 +653,22 @@ def compute_case_outputs(spark, case, base_dir): 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: + if sort_rows: 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): +def validate_test_file(test_file, header, cases): """ 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. + + Only the golden files being checked go through here: regeneration writes a + file from scratch and never reads the old one, so a file this rejects can + always be repaired by regenerating it. """ unknown_header = set(header) - _KNOWN_HEADER_SECTIONS assert not unknown_header, "{}: unknown header sections: {}".format( @@ -620,113 +678,63 @@ def _validate_test_file(test_file, header, cases, regenerate): 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 = set(case) - set(_CASE_SECTION_ORDER) + assert not unknown, "{}: case `{}` has unknown sections: {}".format( + test_file, name, ", ".join(sorted(unknown)) + ) 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) + # A case with no expected_* section would pass vacuously. + 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 ) - # ``_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) + assert not conflicting, ( + "{}: error case `{}` must carry only `expected_error`, not also: {}".format( + test_file, name, ", ".join(conflicting) ) - - -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) + # 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 _regenerate_case(old_case, actual): +def build_golden_case(name, sort_rows, actual): """ - Build the regenerated form of *old_case* from this run's *actual* outputs. + Build the ``.test`` block for a 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. + Only the outputs are golden: the case's identity (*name*) and its ordering + guard (*sort_rows*) are declared by the case method, so regeneration writes + them from the class rather than carrying them over from the old file. """ - carried = {key: old_case.get(key) for key in ("name", "tags", "script")} - carried.update(actual) - return carried + case = {"name": name} + if sort_rows: + case["tags"] = "unordered" + case.update(actual) + return case -def _compare_case(test_case, case, actual): +def compare_case(test_case, case, actual): """Compare the ``expected_*`` sections of *case* against *actual*.""" name = case["name"] for key in _RESULT_SECTIONS: @@ -744,3 +752,230 @@ def _compare_case(test_case, case, actual): got.strip("\n"), "[{}] mismatch in `{}`".format(name, key), ) + + +def is_generating_golden(): + """Whether this run regenerates the golden files instead of checking them.""" + return os.environ.get("SPARK_GENERATE_GOLDEN_FILES") is not None + + +def check_cases_in_sync(test_file, declared, golden): + """ + Check that the golden file describes exactly the declared cases, in order. + + Each test asserts against its own block, so a case the golden file has never + heard of (or one it still remembers after the method was deleted or renamed) + would otherwise go unnoticed. + """ + if declared == golden: + return + missing = [name for name in declared if name not in golden] + extra = [name for name in golden if name not in declared] + if missing: + detail = "cases with no block in the golden file: " + ", ".join(missing) + elif extra: + detail = "blocks in the golden file with no case method: " + ", ".join(extra) + else: + detail = "the golden file lists the cases in a different order" + raise AssertionError( + "{}: golden file is out of sync with the test class ({}); " + "regenerate the golden files".format(test_file, detail) + ) + + +# --------------------------------------------------------------------------- +# Test class integration +# --------------------------------------------------------------------------- + + +class DFGoldenTestMixin: + """ + Mixin turning a class of case methods into DataFrame golden file tests. + + Mix into a session-providing test case, listing this class first so its + ``setUpClass`` runs once the session exists:: + + class GroupByGoldenTests(DFGoldenTestMixin, ReusedConnectTestCase): + golden_file = "group_by.test" + + def _test_group_by_count(self, spark): + return spark.table("testData").groupBy(col("a")).agg(count(col("b"))) + + Every ``_test_`` method declares one case and returns the DataFrame + under test; a ``test_`` method is registered for each, so cases run, + report and can be selected individually like any other unittest test. + + The cases of a class share one Spark Connect session (``newSession()`` off + the session the test class provides, the Connect counterpart of + ``SQLQueryTestSuite``'s per-file ``newSession()``), prepared once by + :meth:`setup_session`. State created there -- temp views, UDFs, session + confs -- is discarded with the session and cannot leak into other classes. + """ + + #: Name of the ``.test`` golden file, resolved next to the test module. + golden_file = None + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + # This mixin's setUpClass must run after the session-providing class has + # created the session, i.e. its setUpClass must be the outer one, which + # is only true when the mixin is listed first. + for base in cls.__mro__: + if base is DFGoldenTestMixin: + break + if base is not cls and "setUpClass" in vars(base): + raise TypeError( + "{} has incorrect inheritance order: DFGoldenTestMixin must be " + "listed before {}. Use: class {}(DFGoldenTestMixin, {}, ...)".format( + cls.__name__, base.__name__, cls.__name__, base.__name__ + ) + ) + for name in cls.case_names(): + setattr(cls, "test_" + name, _make_case_test(name, cls.case_method(name))) + + @classmethod + def case_names(cls): + """The declared case names, in declaration order.""" + names = [] + for klass in reversed(cls.__mro__): + for attr, value in vars(klass).items(): + if not attr.startswith(_CASE_METHOD_PREFIX) or not callable(value): + continue + name = attr[len(_CASE_METHOD_PREFIX) :] + if name not in names: + names.append(name) + return names + + @classmethod + def case_method(cls, name): + """The case method declaring the case *name*.""" + return getattr(cls, _CASE_METHOD_PREFIX + name) + + @classmethod + def golden_file_path(cls): + """Absolute path of the golden file, resolved next to the test module.""" + assert cls.golden_file, "{}: set `golden_file` to its .test file".format(cls.__name__) + module_file = os.path.abspath(inspect.getfile(cls)) + return os.path.join(os.path.dirname(module_file), cls.golden_file) + + @classmethod + def setup_session(cls, spark): + """ + Hook: prepare the session shared by this class's cases. + + Override to create the temp views and other session state the cases + build on. + """ + + @classmethod + def setUpClass(cls): + cls._golden_regenerating = is_generating_golden() + cls._golden_actual = {} + cls._golden_cases = {} + cls._golden_session = None + case_names = cls.case_names() + + # Read the golden file before the session is started: a malformed or + # stale file is worth failing on right away, and failing here leaves + # nothing to clean up (unittest skips tearDownClass when setUpClass + # raises). Regeneration writes the file from the class alone, so it + # neither reads nor requires an existing one. + if case_names: + cls._golden_path = cls.golden_file_path() + if not cls._golden_regenerating: + header, cases = parse_test_file(cls._golden_path) + validate_test_file(cls._golden_path, header, cases) + golden_names = [case["name"] for case in cases] + check_cases_in_sync(cls._golden_path, case_names, golden_names) + cls._golden_cases = {case["name"]: case for case in cases} + + super().setUpClass() + + if not case_names: + return + try: + cls._golden_session = cls.spark.newSession() + # Golden files are generated with ANSI mode on, matching the SQL + # golden tests. The session is discarded with the class, so there + # is nothing to restore. + cls._golden_session.conf.set("spark.sql.ansi.enabled", "true") + cls.setup_session(cls._golden_session) + except BaseException: + # tearDownClass does not run when setUpClass raises, and leaving the + # session behind would hang the test process at exit. + cls._close_golden_session() + super().tearDownClass() + raise + + @classmethod + def tearDownClass(cls): + try: + if cls._golden_regenerating and cls.case_names(): + cls._write_golden_file() + finally: + try: + cls._close_golden_session() + finally: + super().tearDownClass() + + @classmethod + def _close_golden_session(cls): + session = cls._golden_session + cls._golden_session = None + if session is None: + return + # 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 the session-providing class'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 + + @classmethod + def _write_golden_file(cls): + names = cls.case_names() + missing = [name for name in names if name not in cls._golden_actual] + # Writing what a partial run produced would silently drop the cases that + # did not run, so refuse rather than truncate the golden file. + assert not missing, ( + "{}: cannot regenerate, these cases did not run (a failing case, or a run " + "filtered to a subset): {}".format(cls._golden_path, ", ".join(missing)) + ) + header = {"source": "{}.{}".format(cls.__module__, cls.__qualname__)} + cases = [ + build_golden_case(name, is_unordered(cls.case_method(name)), cls._golden_actual[name]) + for name in names + ] + write_test_file(cls._golden_path, header, cases) + + def _run_golden_case(self, name): + cls = type(self) + build_df = getattr(self, _CASE_METHOD_PREFIX + name) + actual = compute_case_outputs( + cls._golden_session, build_df, sort_rows=is_unordered(build_df) + ) + if cls._golden_regenerating: + cls._golden_actual[name] = actual + else: + compare_case(self, cls._golden_cases[name], actual) + + +def _make_case_test(name, case_method): + """Build the unittest method running the case *name*.""" + + def test_case(self): + self._run_golden_case(name) + + test_case.__name__ = "test_" + name + # Carry the case method's docstring over so verbose runs describe the case. + test_case.__doc__ = case_method.__doc__ + return test_case diff --git a/python/pyspark/sql/tests/df_golden/group_by.test b/python/pyspark/sql/tests/df_golden/group_by.test index 4fe132fd5ebf7..82543fe3c91d6 100644 --- a/python/pyspark/sql/tests/df_golden/group_by.test +++ b/python/pyspark/sql/tests/df_golden/group_by.test @@ -1,40 +1,19 @@ --! name __file_metadata__ --! source -df_golden/group_by +pyspark.sql.tests.df_golden.test_group_by.GroupByGoldenTests !-- 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 +select_non_grouping_column --! 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 +agg_counts_no_grouping --! expected_analysis_output Aggregate [count(a#x) AS count(a)#xL, count(b#x) AS count(b)#xL] +- SubqueryAlias testdata @@ -61,11 +40,9 @@ f68f701c7611c45f2e32de9394884c63089bd8174e53a863a1b61168442dbaac --! name -aggregate with non-empty GroupBy expressions +group_by_count --! 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 @@ -95,9 +72,7 @@ printed all 4 rows. --! name -non-aggregate column not in GROUP BY (error) ---! script -scripts/group_by/group_by_missing_aggregation.py +group_by_missing_aggregation --! 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 @@ -105,11 +80,9 @@ Add the columns or the expression to the GROUP BY, aggregate the expression, or --! name -group by with multiple aggregates +group_by_multiple_counts --! 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 @@ -139,9 +112,7 @@ printed all 4 rows. --! name -aggregate grouped by literal ---! script -scripts/group_by/group_by_literal.py +group_by_literal --! expected_analysis_output Aggregate [foo], [foo AS foo#x, count(a#x) AS count(a)#xL] +- SubqueryAlias testdata @@ -168,9 +139,7 @@ printed all 1 rows. --! name -aggregate grouped by literal (hash aggregate) ---! script -scripts/group_by/group_by_literal_hash_agg.py +group_by_literal_hash_agg --! 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) @@ -196,9 +165,7 @@ e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 --! name -aggregate grouped by literal (sort aggregate) ---! script -scripts/group_by/group_by_literal_sort_agg.py +group_by_literal_sort_agg --! expected_analysis_output Aggregate [foo], [foo AS foo#x, max(struct(a, a#x)) AS max(struct(a))#x] +- Filter (a#x = 0) @@ -224,11 +191,9 @@ e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 --! name -aggregate with complex GroupBy expression +group_by_complex_expr --! 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 @@ -259,11 +224,9 @@ e5a84602fefae835a3f9378c2cb5e6c8a653f258181aca8e9e88f1653b75823a --! name -struct() in group by +group_by_struct --! 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 @@ -293,9 +256,7 @@ printed all 4 rows. --! name -aggregate with nulls ---! script -scripts/group_by/agg_stats_with_nulls.py +agg_stats_with_nulls --! 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 @@ -322,9 +283,7 @@ printed all 1 rows. --! name -aggregate with empty input and non-empty GroupBy expressions ---! script -scripts/group_by/empty_input_group_by.py +empty_input_group_by --! expected_analysis_output Aggregate [a#x], [a#x, count(1) AS count(1)#xL] +- Filter false @@ -350,9 +309,7 @@ e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 --! name -aggregate with empty input and empty GroupBy expressions ---! script -scripts/group_by/empty_input_agg.py +empty_input_agg --! expected_analysis_output Aggregate [count(1) AS count(1)#xL] +- Filter false @@ -380,9 +337,7 @@ printed all 1 rows. --! name -aggregate with empty input and constant projection ---! script -scripts/group_by/empty_input_agg_select.py +empty_input_agg_select --! expected_analysis_output Project [1 AS 1#x] +- Aggregate [count(1) AS count(1)#xL] @@ -411,26 +366,7 @@ printed all 1 rows. --! 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 +bool_agg_empty_table --! 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) @@ -458,9 +394,7 @@ ad29e33ee1f3f4cfe263272e81d7c89ea08bdfb7aa45758a6adfc7752d80ab7e --! name -bool aggregates over all null values ---! script -scripts/group_by/bool_agg_all_nulls.py +bool_agg_all_nulls --! 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) @@ -488,9 +422,7 @@ ad29e33ee1f3f4cfe263272e81d7c89ea08bdfb7aa45758a6adfc7752d80ab7e --! name -bool aggregates null filtering ---! script -scripts/group_by/bool_agg_null_filtering.py +bool_agg_null_filtering --! 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) @@ -518,11 +450,9 @@ printed all 1 rows. --! name -bool aggregates with group by +bool_agg_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 @@ -553,11 +483,9 @@ printed all 5 rows. --! name -bool aggregates with having false +bool_agg_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) @@ -589,9 +517,7 @@ b38b91ef7aec43d0957f73891494f8c09519ffe4731fb729a0569c7c2d970118 --! name -bool aggregates with having null ---! script -scripts/group_by/bool_agg_having_null.py +bool_agg_having_null --! expected_analysis_output Project [k#x, every_v#x] +- Filter isnull(every_v#x) @@ -621,36 +547,28 @@ printed all 1 rows. --! name -every input type checking: int (error) ---! script -scripts/group_by/every_int_error.py +every_int_error --! 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 +some_int_error --! 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 +bool_or_int_error --! 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 +every_string --! expected_analysis_output Aggregate [every(cast(true as boolean)) AS every(true)#x] +- SubqueryAlias test_agg @@ -677,29 +595,23 @@ b5bea41b6c623f7c09f1bf24dcae58ebab3c0cdd90ad966bc43a45b44867e12b --! name -bool_and input type checking: decimal (error) ---! script -scripts/group_by/bool_and_decimal_error.py +bool_and_decimal_error --! 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 +bool_or_double_error --! 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 +every_window --! 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] @@ -738,11 +650,9 @@ printed all 10 rows. --! name -some as window expression +some_window --! 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] @@ -781,11 +691,9 @@ printed all 10 rows. --! name -bool_or as window expression +bool_or_window --! 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] @@ -824,11 +732,9 @@ printed all 10 rows. --! name -bool_and as window expression +bool_and_window --! 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] @@ -867,11 +773,9 @@ printed all 10 rows. --! name -bool_or as window expression (repeated) +bool_or_window_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] @@ -910,9 +814,7 @@ printed all 10 rows. --! name -having referencing aggregate expression ---! script -scripts/group_by/having_agg_count.py +having_agg_count --! expected_analysis_output Project [cnt#xL] +- Filter (cnt#xL > cast(1 as bigint)) @@ -942,11 +844,9 @@ printed all 1 rows. --! name -having on aliased max +having_max_v --! tags unordered ---! script -scripts/group_by/having_max_v.py --! expected_analysis_output Project [k#x, max_v#x] +- Filter (max_v#x = true) @@ -978,9 +878,7 @@ d6212a5d216fffb25342120cad425c6ebe96fc8485c6b082be8624faf6a531c4 --! name -aggregate expression referenced through alias ---! script -scripts/group_by/agg_alias_filter.py +agg_alias_filter --! expected_analysis_output Filter (cnt#xL > cast(1 as bigint)) +- Aggregate [count(k#x) AS cnt#xL] @@ -1009,11 +907,9 @@ printed all 1 rows. --! name -SPARK-34581: do not optimize out grouping expressions +grouping_exprs_not_optimized --! 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 @@ -1041,11 +937,9 @@ c5421419b0d389431648774e2aa0156b8472da926b1dd3cd50c94821001e7bcd --! name -PullOutGroupingExpressions pulls out grouping expressions +pull_out_grouping_exprs --! 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 @@ -1073,11 +967,9 @@ ce61427fcddb697480c2921f7450c541ac0bf1fe284b7e88bd9d784a84a2120a --! name -count_if with grouping expression reference +count_if_group_by --! 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 diff --git a/python/pyspark/sql/tests/df_golden/regenerate.sh b/python/pyspark/sql/tests/df_golden/regenerate.sh index a185abaebbf18..d3870b4d8ffa7 100755 --- a/python/pyspark/sql/tests/df_golden/regenerate.sh +++ b/python/pyspark/sql/tests/df_golden/regenerate.sh @@ -16,7 +16,7 @@ # 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. +# re-running their test modules with SPARK_GENERATE_GOLDEN_FILES=1. # # Usage: # python/pyspark/sql/tests/df_golden/regenerate.sh [--verify] @@ -33,21 +33,36 @@ for arg in "$@"; do esac done -REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && git rev-parse --show-toplevel)" +THIS_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$THIS_DIR" && 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" +# Every module declaring golden cases, i.e. every test module in this directory +# with a top-level class mixing in DFGoldenTestMixin. Adding one needs no change +# here. +MODULES=() +for file in "$THIS_DIR"/test_*.py; do + if grep -q "^class .*(DFGoldenTestMixin" "$file"; then + MODULES+=("pyspark.sql.tests.df_golden.$(basename "$file" .py)") + fi +done +if [[ ${#MODULES[@]} -eq 0 ]]; then + echo "ERROR: no golden test modules found in $REL_DIR" >&2 + exit 1 +fi +TESTNAMES="$(IFS=,; echo "${MODULES[*]}")" + +echo ">>> Regenerating golden files for: $TESTNAMES" +SPARK_GENERATE_GOLDEN_FILES=1 python/run-tests --testnames "$TESTNAMES" 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" + python/run-tests --testnames "$TESTNAMES" 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 deleted file mode 100644 index 1db2ee02d52f3..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/agg_alias_filter.py +++ /dev/null @@ -1,5 +0,0 @@ -# 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 deleted file mode 100644 index 969c5f69d8d15..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/agg_counts_no_grouping.py +++ /dev/null @@ -1,3 +0,0 @@ -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 deleted file mode 100644 index db0ecf242d791..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/agg_stats_with_nulls.py +++ /dev/null @@ -1,27 +0,0 @@ -# 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 deleted file mode 100644 index 2115205c1a38d..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_all_nulls.py +++ /dev/null @@ -1,9 +0,0 @@ -# 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 deleted file mode 100644 index 727000377f6c6..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_empty_table.py +++ /dev/null @@ -1,9 +0,0 @@ -# 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 deleted file mode 100644 index 6bc7ebf478cea..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_group_by.py +++ /dev/null @@ -1,9 +0,0 @@ -# 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 deleted file mode 100644 index 1d8837200bca9..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_having_false.py +++ /dev/null @@ -1,11 +0,0 @@ -# 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 deleted file mode 100644 index 99cf57b5f1541..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_having_null.py +++ /dev/null @@ -1,9 +0,0 @@ -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 deleted file mode 100644 index 81c19c4db4a17..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_agg_null_filtering.py +++ /dev/null @@ -1,9 +0,0 @@ -# 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 deleted file mode 100644 index c212f96113b21..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_and_decimal_error.py +++ /dev/null @@ -1,6 +0,0 @@ -# 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 deleted file mode 100644 index 1fae3d538036a..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_and_window.py +++ /dev/null @@ -1,6 +0,0 @@ -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 deleted file mode 100644 index 6176a2ce894d0..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_or_double_error.py +++ /dev/null @@ -1,5 +0,0 @@ -# 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 deleted file mode 100644 index 5ad8187a2998d..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_or_int_error.py +++ /dev/null @@ -1,5 +0,0 @@ -# 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 deleted file mode 100644 index a7e4dc0d281f6..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_or_window.py +++ /dev/null @@ -1,6 +0,0 @@ -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 deleted file mode 100644 index a7e4dc0d281f6..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_or_window_repeated.py +++ /dev/null @@ -1,6 +0,0 @@ -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 deleted file mode 100644 index 35443d59fb2fc..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/count_if_group_by.py +++ /dev/null @@ -1,3 +0,0 @@ -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 deleted file mode 100644 index a5efcab8a0897..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/empty_input_agg.py +++ /dev/null @@ -1,5 +0,0 @@ -# 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 deleted file mode 100644 index e7fd87eaa97f3..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/empty_input_agg_select.py +++ /dev/null @@ -1,3 +0,0 @@ -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 deleted file mode 100644 index cb64aaf418c09..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/empty_input_group_by.py +++ /dev/null @@ -1,5 +0,0 @@ -# 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 deleted file mode 100644 index 0e484981c9c48..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/every_int_error.py +++ /dev/null @@ -1,5 +0,0 @@ -# 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 deleted file mode 100644 index 5a69c9075e999..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/every_string.py +++ /dev/null @@ -1,5 +0,0 @@ -# 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 deleted file mode 100644 index 95c1241a917b1..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/every_window.py +++ /dev/null @@ -1,8 +0,0 @@ -# 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 deleted file mode 100644 index a3f362b2cf405..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_complex_expr.py +++ /dev/null @@ -1,5 +0,0 @@ -# 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 deleted file mode 100644 index 483ed988a46ea..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_count.py +++ /dev/null @@ -1,5 +0,0 @@ -# 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 deleted file mode 100644 index c4f5d6b216731..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_literal.py +++ /dev/null @@ -1,5 +0,0 @@ -# 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 deleted file mode 100644 index b4d42eaf0baf6..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_literal_hash_agg.py +++ /dev/null @@ -1,10 +0,0 @@ -# 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 deleted file mode 100644 index cb83ecdf9ee0e..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_literal_sort_agg.py +++ /dev/null @@ -1,5 +0,0 @@ -# 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 deleted file mode 100644 index ff822ed8f9ca5..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_missing_aggregation.py +++ /dev/null @@ -1,3 +0,0 @@ -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 deleted file mode 100644 index 8e1a2d5ea055c..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_multiple_counts.py +++ /dev/null @@ -1,3 +0,0 @@ -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 deleted file mode 100644 index ff6e728367afc..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/group_by_struct.py +++ /dev/null @@ -1,5 +0,0 @@ -# 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 deleted file mode 100644 index 91086696a432a..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/grouping_exprs_not_optimized.py +++ /dev/null @@ -1,10 +0,0 @@ -# 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 deleted file mode 100644 index 879f498da0e8a..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/having_agg_count.py +++ /dev/null @@ -1,10 +0,0 @@ -# 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 deleted file mode 100644 index 3303169b549fd..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/having_max_v.py +++ /dev/null @@ -1,9 +0,0 @@ -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 deleted file mode 100644 index 8d1f6807f21a7..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/pull_out_grouping_exprs.py +++ /dev/null @@ -1,10 +0,0 @@ -# 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 deleted file mode 100644 index 466a65a153d79..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/select_non_grouping_column.py +++ /dev/null @@ -1,5 +0,0 @@ -# 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 deleted file mode 100644 index 8e1734fe1ac3e..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/setup_test_agg.py +++ /dev/null @@ -1,8 +0,0 @@ -# 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 deleted file mode 100644 index 4a83505a3d79d..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/setup_test_data.py +++ /dev/null @@ -1,6 +0,0 @@ -# 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 deleted file mode 100644 index 3bd9dbba84172..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/some_int_error.py +++ /dev/null @@ -1,5 +0,0 @@ -# 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 deleted file mode 100644 index 21c84d14babea..0000000000000 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/some_window.py +++ /dev/null @@ -1,6 +0,0 @@ -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 deleted file mode 100644 index 311491c5d7ed2..0000000000000 --- a/python/pyspark/sql/tests/df_golden/test_df_golden.py +++ /dev/null @@ -1,143 +0,0 @@ -# -# 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.sql.tests.df_golden.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.sql.tests.df_golden.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 index 5949093e24efe..0335c2f2891a1 100644 --- a/python/pyspark/sql/tests/df_golden/test_df_golden_framework.py +++ b/python/pyspark/sql/tests/df_golden/test_df_golden_framework.py @@ -19,9 +19,10 @@ Unit tests for the ``pyspark.sql.tests.df_golden.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``. +validation, output normalization, result rendering and case registration -- +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 the ``test_.py`` +modules next to this one. """ import os @@ -29,16 +30,20 @@ import unittest from pyspark.sql.tests.df_golden.df_golden import ( - _compare_case, - _regenerate_case, - _validate_test_file, + DFGoldenTestMixin, + build_golden_case, + check_cases_in_sync, + compare_case, format_double, format_error, hash_result_rows, + is_unordered, parse_tags, parse_test_file, render_result_table, replace_not_included, + unordered, + validate_test_file, write_test_file, ) @@ -58,9 +63,7 @@ def _write(self, text): def test_parse_basic_case(self): path = self._write( "--! name\n" - "my case\n" - "--! script\n" - "scripts/x.py\n" + "my_case\n" "--! expected_output_schema\n" "struct\n" "!-- end\n" @@ -68,8 +71,7 @@ def test_parse_basic_case(self): 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]["name"], "my_case") self.assertEqual(cases[0]["expected_output_schema"], "struct") def test_parse_extracts_file_metadata_header(self): @@ -77,18 +79,21 @@ def test_parse_extracts_file_metadata_header(self): "--! name\n" "__file_metadata__\n" "--! source\n" - "df_golden/group_by\n" + "pyspark.sql.tests.df_golden.test_group_by.GroupByGoldenTests\n" "!-- end\n" "\n\n" "--! name\n" "c1\n" - "--! script\n" - "scripts/a.py\n" + "--! expected_output_schema\n" + "struct\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( + header, + {"source": "pyspark.sql.tests.df_golden.test_group_by.GroupByGoldenTests"}, + ) self.assertEqual(len(cases), 1) self.assertEqual(cases[0]["name"], "c1") @@ -96,8 +101,6 @@ 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" @@ -111,20 +114,18 @@ def test_parse_preserves_multiline_section_body(self): ) def test_round_trip_parse_write_parse(self): - header = {"source": "df_golden/group_by"} + header = {"source": "pyspark.sql.tests.df_golden.test_group_by.GroupByGoldenTests"} cases = [ { - "name": "ordered case", - "script": "scripts/a.py", + "name": "ordered_case", "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", + "name": "error_case", "tags": "unordered", - "script": "scripts/b.py", "expected_error": "[SOME_ERROR] boom", }, ] @@ -140,13 +141,13 @@ def test_write_only_emits_known_sections_in_order(self): write_test_file( path, {}, - [{"script": "scripts/a.py", "name": "c", "junk": "ignored"}], + [{"expected_output_schema": "struct", "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")) + # name precedes the outputs in the canonical order regardless of dict order. + self.assertLess(body.index("--! name"), body.index("--! expected_output_schema")) # -- tags ------------------------------------------------------------- @@ -155,12 +156,22 @@ def test_parse_tags_splits_on_whitespace_and_commas(self): self.assertEqual(parse_tags({}), set()) self.assertEqual(parse_tags({"tags": ""}), set()) + def test_unordered_decorator_marks_the_case_method(self): + def plain(self, spark): + pass + + @unordered + def marked(self, spark): + pass + + self.assertFalse(is_unordered(plain)) + self.assertTrue(is_unordered(marked)) + # -- validation ------------------------------------------------------- def _valid_case(self, **overrides): case = { "name": "c", - "script": "scripts/a.py", "expected_output_schema": "struct", } case.update(overrides) @@ -168,91 +179,67 @@ def _valid_case(self, **overrides): def test_validate_accepts_well_formed_file(self): # Should not raise. - _validate_test_file("f.test", {"source": "x"}, [self._valid_case()], regenerate=False) + validate_test_file("f.test", {"source": "x"}, [self._valid_case()]) 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) + validate_test_file("f.test", {"bogus": "x"}, [self._valid_case()]) def test_validate_rejects_no_cases(self): with self.assertRaisesRegex(AssertionError, "no test cases found"): - _validate_test_file("f.test", {}, [], regenerate=False) + validate_test_file("f.test", {}, []) 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, - ) + validate_test_file("f.test", {}, [{"expected_output_schema": "x"}]) 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 - ) + validate_test_file("f.test", {}, [self._valid_case(expected_bogus="x")]) 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) + validate_test_file("f.test", {}, [self._valid_case(tags="wat")]) def test_validate_accepts_known_unordered_tag(self): - _validate_test_file("f.test", {}, [self._valid_case(tags="unordered")], regenerate=False) + validate_test_file("f.test", {}, [self._valid_case(tags="unordered")]) - def test_validate_rejects_vacuous_case_in_verify_mode(self): + def test_validate_rejects_vacuous_case(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 - ) + validate_test_file("f.test", {}, [{"name": "c"}]) 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, - ) + validate_test_file("f.test", {}, [{"name": "c", "expected_error": "[ERR] boom"}]) - 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"): + def test_validate_rejects_legacy_sections(self): + # Sections dropped from the format (the analysis/execution error split + # collapsed into ``expected_error``, the script pointer replaced by the + # case method) are now unknown sections. + for legacy in ("expected_analysis_error", "expected_execution_error", "script"): 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, - ) + validate_test_file("f.test", {}, [self._valid_case(**{legacy: "x"})]) - 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) + # -- case / golden file sync ------------------------------------------ + + def test_check_cases_in_sync_accepts_identical_lists(self): + check_cases_in_sync("f.test", ["a", "b"], ["a", "b"]) + + def test_check_cases_in_sync_rejects_case_missing_from_golden(self): + with self.assertRaisesRegex(AssertionError, "no block in the golden file: b"): + check_cases_in_sync("f.test", ["a", "b"], ["a"]) + + def test_check_cases_in_sync_rejects_golden_block_without_case(self): + with self.assertRaisesRegex(AssertionError, "no case method: b"): + check_cases_in_sync("f.test", ["a"], ["a", "b"]) + + def test_check_cases_in_sync_rejects_reordering(self): + # The golden file is written in declaration order, so a different order + # means the file no longer matches the class. + with self.assertRaisesRegex(AssertionError, "different order"): + check_cases_in_sync("f.test", ["a", "b"], ["b", "a"]) # -- output normalization -------------------------------------------- @@ -268,7 +255,7 @@ 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' + '"__truediv__" was called from /abs/path/test_arithmetic.py:7\n' "\n" "JVM stacktrace:\n" "org.apache.spark.SparkArithmeticException: ..." @@ -324,12 +311,12 @@ def test_compare_case_passes_on_match(self): "name": "c", "expected_output_schema": "struct", } - _compare_case(self, case, {"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( + compare_case( self, {"name": "c", "expected_output_schema": "struct"}, { @@ -340,7 +327,7 @@ def test_compare_case_ignores_sections_absent_from_golden(self): def test_compare_case_fails_on_value_mismatch(self): with self.assertRaises(AssertionError): - _compare_case( + compare_case( self, {"name": "c", "expected_output_schema": "struct"}, {"expected_output_schema": "struct"}, @@ -348,7 +335,7 @@ def test_compare_case_fails_on_value_mismatch(self): def test_compare_case_fails_when_expected_section_not_produced(self): with self.assertRaisesRegex(AssertionError, "expected section `expected_result`"): - _compare_case( + compare_case( self, {"name": "c", "expected_result": "printed all 0 rows."}, {"expected_error": "[ERR] boom"}, @@ -357,59 +344,49 @@ def test_compare_case_fails_when_expected_section_not_produced(self): # -- under-assertion guards ------------------------------------------ def test_validate_accepts_result_with_hash(self): - _validate_test_file( + 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, + validate_test_file( + "f.test", {}, [self._valid_case(expected_result="printed all 0 rows.")] ) 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 - ) + validate_test_file("f.test", {}, [self._valid_case(expected_result_hash="h")]) def test_validate_rejects_error_case_mixed_with_result(self): with self.assertRaisesRegex(AssertionError, "must carry only `expected_error`"): - _validate_test_file( + 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( + 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 -------------------------------- @@ -459,42 +436,33 @@ def test_format_value_refuses_float(self): 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" + "--! name\nc\n--! expected_output_schema\nx\n--! expected_output_schema\ny\n!-- end\n" ) - with self.assertRaisesRegex(AssertionError, "duplicate section `script`"): + with self.assertRaisesRegex(AssertionError, "duplicate section `expected_output_schema`"): 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" - ) + path = self._write("--! name\nc\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") + path = self._write("stray text\n--! name\nc\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" - ) + path = self._write("--! name\nc1\n!-- end\n\n\n--! name\nc2\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: + def test_parse_rejects_unterminated_trailing_case(self): + # Last case missing "!-- end", i.e. a truncated file. + path = self._write("--! name\nc\n--! expected_output_schema\nx\n") with self.assertRaisesRegex(AssertionError, "does not end with"): - parse_test_file(path, require_terminated=True) + parse_test_file(path) def test_format_value_refuses_tab_or_newline_in_string(self): try: @@ -510,25 +478,73 @@ def test_format_value_refuses_tab_or_newline_in_string(self): # -- 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", - } + def test_build_golden_case_writes_identity_from_the_case_method(self): 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") + case = build_golden_case("c", False, actual) + self.assertEqual(case["name"], "c") + self.assertNotIn("tags", case) + self.assertEqual(case["expected_optimized_output"], "new base") + self.assertEqual(case["expected_output_schema"], "new schema") + # An unordered case records the tag that explains its sorted rows. + self.assertEqual(build_golden_case("c", True, actual)["tags"], "unordered") + + # -- case registration ------------------------------------------------ + + def test_mixin_registers_a_test_per_case_in_declaration_order(self): + class Cases(DFGoldenTestMixin, unittest.TestCase): + golden_file = "x.test" + + def _test_second(self, spark): + """Second case.""" + + def _test_first(self, spark): + pass + + # Declaration order, not alphabetical: the golden file is written in it. + self.assertEqual(Cases.case_names(), ["second", "first"]) + self.assertTrue(callable(Cases.test_second)) + self.assertTrue(callable(Cases.test_first)) + self.assertEqual(Cases.test_second.__doc__, "Second case.") + + def test_mixin_inherits_cases_from_base_classes(self): + class Base(DFGoldenTestMixin, unittest.TestCase): + golden_file = "x.test" + + def _test_inherited(self, spark): + pass + + class Derived(Base): + def _test_added(self, spark): + pass + + self.assertEqual(Derived.case_names(), ["inherited", "added"]) + + def test_mixin_rejects_wrong_inheritance_order(self): + # The mixin's setUpClass has to run after the session exists, which only + # happens when it precedes the session-providing class. + with self.assertRaisesRegex(TypeError, "must be listed before"): + + class Wrong(unittest.TestCase, DFGoldenTestMixin): + golden_file = "x.test" + + def test_mixin_resolves_golden_file_next_to_the_test_module(self): + class Cases(DFGoldenTestMixin, unittest.TestCase): + golden_file = "group_by.test" + + self.assertEqual( + Cases.golden_file_path(), + os.path.join(os.path.dirname(os.path.abspath(__file__)), "group_by.test"), + ) + + def test_mixin_requires_a_golden_file(self): + class Cases(DFGoldenTestMixin, unittest.TestCase): + pass + + with self.assertRaisesRegex(AssertionError, "set `golden_file`"): + Cases.golden_file_path() if __name__ == "__main__": diff --git a/python/pyspark/sql/tests/df_golden/test_group_by.py b/python/pyspark/sql/tests/df_golden/test_group_by.py new file mode 100644 index 0000000000000..2e2b628a2f65f --- /dev/null +++ b/python/pyspark/sql/tests/df_golden/test_group_by.py @@ -0,0 +1,324 @@ +# +# 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 for GROUP BY and aggregate operations. + +The cases are modelled on the GROUP BY cases of the SQL golden tests +(``sql/core/src/test/resources/sql-tests/inputs/group-by.sql``). Their expected +outputs live in ``group_by.test``; see ``pyspark.sql.tests.df_golden.df_golden`` +for the framework and how to regenerate them. +""" + +from decimal import Decimal + +from pyspark.sql import Window +from pyspark.sql.functions import ( + approx_count_distinct, + avg, + bool_and, + bool_or, + col, + count, + count_if, + every, + first, + kurtosis, + lit, + max, + min, + round, + skewness, + some, + stddev, + struct, + sum, + variance, + when, +) +from pyspark.sql.tests.df_golden.df_golden import DFGoldenTestMixin, unordered +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class GroupByGoldenTests(DFGoldenTestMixin, ReusedConnectTestCase): + """Test the aggregate operator.""" + + golden_file = "group_by.test" + + @classmethod + def setup_session(cls, spark): + 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)""") + 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)""") + + def _test_select_non_grouping_column(self, spark): + """Select a non-grouping column without a GROUP BY (error).""" + return spark.table("testData").select(col("a"), count(col("b"))) + + def _test_agg_counts_no_grouping(self, spark): + """Aggregate with empty GroupBy expressions.""" + return spark.table("testData").agg(count(col("a")), count(col("b"))) + + @unordered + def _test_group_by_count(self, spark): + """Aggregate with non-empty GroupBy expressions.""" + return spark.table("testData").groupBy(col("a")).agg(count(col("b"))) + + def _test_group_by_missing_aggregation(self, spark): + """Non-aggregate column not in GROUP BY (error).""" + return spark.table("testData").groupBy(col("b")).agg(col("a"), count(col("b"))) + + @unordered + def _test_group_by_multiple_counts(self, spark): + """Group by with multiple aggregates.""" + return spark.table("testData").groupBy(col("a")).agg(count(col("a")), count(col("b"))) + + def _test_group_by_literal(self, spark): + """Aggregate grouped by literal.""" + return spark.table("testData").groupBy(lit("foo")).agg(count(col("a"))) + + def _test_group_by_literal_hash_agg(self, spark): + """Aggregate grouped by literal (hash aggregate).""" + return ( + spark.table("testData") + .filter(col("a") == 0) + .groupBy(lit("foo")) + .agg(approx_count_distinct(col("a"))) + ) + + def _test_group_by_literal_sort_agg(self, spark): + """Aggregate grouped by literal (sort aggregate).""" + return ( + spark.table("testData") + .filter(col("a") == 0) + .groupBy(lit("foo")) + .agg(max(struct(col("a")))) + ) + + @unordered + def _test_group_by_complex_expr(self, spark): + """Aggregate with complex GroupBy expression.""" + return spark.table("testData").groupBy(col("a") + col("b")).agg(count(col("b"))) + + @unordered + def _test_group_by_struct(self, spark): + """struct() in group by.""" + return ( + spark.table("testData") + .groupBy(struct((col("a") + 0.1).alias("aa"))) + .agg(count(lit(1))) + ) + + def _test_agg_stats_with_nulls(self, spark): + """Aggregate with nulls.""" + return 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")), + ) + + def _test_empty_input_group_by(self, spark): + """Aggregate with empty input and non-empty GroupBy expressions.""" + return spark.table("testData").filter(lit(False)).groupBy(col("a")).agg(count(lit(1))) + + def _test_empty_input_agg(self, spark): + """Aggregate with empty input and empty GroupBy expressions.""" + return spark.table("testData").filter(lit(False)).agg(count(lit(1))) + + def _test_empty_input_agg_select(self, spark): + """Aggregate with empty input and constant projection.""" + return spark.table("testData").filter(lit(False)).agg(count(lit(1))).select(lit(1)) + + def _test_bool_agg_empty_table(self, spark): + """Bool aggregates over empty table.""" + return self._bool_aggs(spark.table("test_agg").filter(lit(1) == 0)) + + def _test_bool_agg_all_nulls(self, spark): + """Bool aggregates over all null values.""" + return self._bool_aggs(spark.table("test_agg").filter(col("k") == 4)) + + def _test_bool_agg_null_filtering(self, spark): + """Bool aggregates null filtering.""" + return self._bool_aggs(spark.table("test_agg").filter(col("k") == 5)) + + @unordered + def _test_bool_agg_group_by(self, spark): + """Bool aggregates with group by.""" + return self._bool_aggs(spark.table("test_agg").groupBy(col("k"))) + + @unordered + def _test_bool_agg_having_false(self, spark): + """Bool aggregates with having false.""" + return ( + 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")) + ) + + def _test_bool_agg_having_null(self, spark): + """Bool aggregates with having null.""" + return ( + 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")) + ) + + def _test_every_int_error(self, spark): + """every() input type checking: int (error).""" + return spark.table("test_agg").select(every(lit(1))) + + def _test_some_int_error(self, spark): + """some() input type checking: int (error).""" + return spark.table("test_agg").select(some(lit(1))) + + def _test_bool_or_int_error(self, spark): + """bool_or() input type checking: int (error).""" + return spark.table("test_agg").select(bool_or(lit(1))) + + def _test_every_string(self, spark): + """every() input type checking: string.""" + return spark.table("test_agg").select(every(lit("true"))) + + def _test_bool_and_decimal_error(self, spark): + """bool_and() input type checking: decimal (error).""" + return spark.table("test_agg").select(bool_and(lit(Decimal("1.0")))) + + def _test_bool_or_double_error(self, spark): + """bool_or() input type checking: double (error).""" + return spark.table("test_agg").select(bool_or(lit(1.0))) + + @unordered + def _test_every_window(self, spark): + """every() as window expression.""" + return self._bool_agg_window(spark, every) + + @unordered + def _test_some_window(self, spark): + """some() as window expression.""" + return self._bool_agg_window(spark, some) + + @unordered + def _test_bool_or_window(self, spark): + """bool_or() as window expression.""" + return self._bool_agg_window(spark, bool_or) + + @unordered + def _test_bool_and_window(self, spark): + """bool_and() as window expression.""" + return self._bool_agg_window(spark, bool_and) + + @unordered + def _test_bool_or_window_repeated(self, spark): + """bool_or() as window expression (repeated).""" + return self._bool_agg_window(spark, bool_or) + + def _test_having_agg_count(self, spark): + """Having referencing aggregate expression.""" + return ( + spark.table("test_agg") + .agg(count(col("k")).alias("cnt")) + .filter(col("cnt") > 1) + .select(col("cnt")) + ) + + @unordered + def _test_having_max_v(self, spark): + """Having on aliased max.""" + return ( + 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")) + ) + + def _test_agg_alias_filter(self, spark): + """Aggregate expression referenced through alias.""" + return spark.table("test_agg").agg(count(col("k")).alias("cnt")).filter(col("cnt") > 1) + + @unordered + def _test_grouping_exprs_not_optimized(self, spark): + """SPARK-34581: do not optimize out grouping expressions. + + They must survive in aggregate expressions without aggregate function. + """ + return ( + spark.table("testData") + .groupBy(col("a").isNull()) + .agg(~col("a").isNull(), count("*").alias("c")) + ) + + @unordered + def _test_pull_out_grouping_exprs(self, spark): + """PullOutGroupingExpressions pulls out grouping expressions. + + It pulls them out from inside AggregateExpressions too. + """ + return ( + spark.table("testData") + .groupBy(col("a").isNull()) + .agg(when(~col("a").isNull(), 0).otherwise(1), first(col("a").isNull()).alias("c")) + ) + + @unordered + def _test_count_if_group_by(self, spark): + """count_if() with grouping expression reference.""" + return ( + spark.table("testData") + .groupBy(col("a") + 1) + .agg(count_if((col("a") + 1) == col("b"))) + ) + + @staticmethod + def _bool_aggs(relation): + """Aggregate ``v`` with every boolean aggregate, over a DataFrame or GroupedData.""" + return relation.agg( + every(col("v")), + some(col("v")), + bool_or(col("v")), + bool_and(col("v")), + bool_or(col("v")), + ) + + @staticmethod + def _bool_agg_window(spark, bool_agg): + """Apply *bool_agg* to ``v`` as a window expression over ``test_agg``.""" + return spark.table("test_agg").select( + col("k"), col("v"), bool_agg(col("v")).over(Window.partitionBy("k").orderBy("v")) + ) + + +if __name__ == "__main__": + from pyspark.testing import main + + main()