[SPARK-58480][PYTHON][TEST] Declare DataFrame golden test cases as methods of a test class by dtenedor · Pull Request #57691 · apache/spark

dtenedor · 2026-07-31T20:55:29Z

What changes were proposed in this pull request?

This PR follows up on #57122 and #57639 by replacing the per-case Python snippet files of the DataFrame golden test framework with a single test class per golden file.

The files under python/pyspark/sql/tests/df_golden/scripts/group_by/ are replaced by one module, python/pyspark/sql/tests/df_golden/test_group_by.py, holding one class whose _test_<case> methods each build and return the DataFrame under test:

        class GroupByGoldenTests(DFGoldenTestMixin, ReusedConnectTestCase):
            golden_file = "group_by.test"
            @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")))

This is ordinary Python: imports are at the top of the module, cases can share helpers, and spark is a parameter rather than a name injected into an exec() namespace.

Cases are ordinary unit tests: DFGoldenTestMixin (a mixin, following the existing pyspark.testing.goldenutils.GoldenFileTestMixin pattern) registers a real test_<case> method for every _test_<case> method, carrying its docstring over. Each case is reported individually and can be run on its own:

        python/run-tests --testnames \
          "pyspark.sql.tests.df_golden.test_group_by GroupByGoldenTests.test_group_by_count"

Why are the changes needed?

The framework as merged put each case's DataFrame program in its own file, which drew several objections in the review of #57122:

File count. One file per case does not scale.
The snippets were not real Python. They were exec()'d with spark injected into their namespace, which is why they needed an F821 lint exemption and a RAT exclusion, and why they could not be imported or run outside the framework.
Reproducing a failure was hard A whole .test file used to run as one unittest method, so a failure named a case but gave no way to run just that case. Now every case is a test method that can be selected by name, and the case body is an importable method.
Structure. Organizing cases as methods of a class is the conventional Python shape, and it lets cases share helpers.

Does this PR introduce any user-facing change?

No. This is test-only: the framework and its test corpus.

How was this patch tested?

This PR is test-only. The migration was verified as follows.

Golden outputs are unchanged.
The regeneration writer reproduces the file byte for byte.
The class and the golden file are in sync.
Framework unit tests: pyspark.sql.tests.df_golden.test_df_golden_framework passes (51 tests).

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Cursor (Opus 5)

uros-b · 2026-07-31T21:28:04Z

Thank you @dtenedor! cc @zhengruifeng

gaogaotiantian

Generally, I believe this is the correct direction. We should have a relatively complicated mixin and many easy to read/reproduce test cases. I'm glad that we can get rid of a lot of files.

I have some comments about the implementation. It's a framework and I hope the code is cleaner and easier for extension in the future.

I also have some suggestions for next steps, but they are not for this PR.

gaogaotiantian · 2026-07-31T21:50:54Z

-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):


I think this function is a bit too simple now to be a function. Let's just inline this to the only callsite.

gaogaotiantian · 2026-07-31T22:02:16Z

+    return os.environ.get("SPARK_GENERATE_GOLDEN_FILES") is not None
+
+
+def check_cases_in_sync(test_file, declared, golden):


If we are doing an assertion here, let's change the function name to assert_cases_in_sync.

gaogaotiantian · 2026-07-31T22:08:09Z

+        for base in cls.__mro__:
+            if base is DFGoldenTestMixin:
+                break
+            if base is not cls and "setUpClass" in vars(base):


Is this correct? A subclass can have a setUpClass method that works with the system right? So A -> B(setUpClass) -> DFGoldenTestMixin -> PySparkTestBase should be a valid inherit chain but this will blocks it.

I don't believe we should use setUpClass to guarantee this. I don't think we have to stop this when the subclass is declared. It's fine to catch this on run-time as it will be a one time thing. We won't keep losing time because we are not able to catch it on declaration.

gaogaotiantian · 2026-07-31T22:26:31Z

+                    )
+                )
+        for name in cls.case_names():
+            setattr(cls, "test_" + name, _make_case_test(name, cls.case_method(name)))


A few things to make this _test -> test.

This line is technically not 100% correct. This line is converting all the subclass _test methods to its own test methods. It probably would just work because how unittest discover methods, but we can probably make this technically correct.

case_method is actually only used here. I think we should make _make_case_test a classmethod of DFGoldenTestMixin. The natually the function has the cls argument and can add test method to it. We can also get rid of case_method.

gaogaotiantian · 2026-07-31T22:27:32Z

+        """The declared case names, in declaration order."""
+        names = []
+        for klass in reversed(cls.__mro__):
+            for attr, value in vars(klass).items():


This might be better:

for name, func in inspect.getmembers(B, predicate=inspect.isfunction):

gaogaotiantian · 2026-07-31T22:40:35Z

+                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)


It seems like header is only used for validation, so we don't even need to get it from parsing I guess? How about we directly get the dict for cls._golden_cases and check whether it's sync?

gaogaotiantian · 2026-07-31T22:43:15Z

+
+    @classmethod
+    def _close_golden_session(cls):
+        session = cls._golden_session


This logic is a bit weird if we do not consider any racing issues.

It's just

if cls._golden_session is None: return client = cls._golden_session.client cls.golden_session = None

right?

gaogaotiantian · 2026-07-31T22:46:36Z

+        )
+        header = {"source": "{}.{}".format(cls.__module__, cls.__qualname__)}
+        cases = [
+            build_golden_case(name, is_unordered(cls.case_method(name)), cls._golden_actual[name])


I don't think is_unordered should be specially treated. It should be a "tag" of a case - so whether it's actual or golden, it's just part of the data. When you compare, you check that tag.

We probably have other tags in the future, we don't want to add a new argument for every tag. It's just data.

gaogaotiantian · 2026-07-31T22:47:35Z

+        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)


Similar here. build_df already has the information whether it's unordered. It does not need to be part of the arguments. compute_case_outputs can deal with the special attributes of build_df inside it and convert everything to tags.

gaogaotiantian · 2026-07-31T22:49:37Z

 __file_metadata__
 --! source
-df_golden/group_by
+pyspark.sql.tests.df_golden.test_group_by.GroupByGoldenTests


I hope eventually we can get rid of our own file format and use something more generic like toml - which are also super readable by both humans and machines. We won't need to maintain our own parser for each section anymore. But this is not part of this PR.

commit

81ce6b6

dtenedor requested a review from gaogaotiantian July 31, 2026 20:55

dtenedor changed the title ~~[SPARK-58480][PYTHON][TEST][FOLLOWUP] Declare DataFrame golden test cases as methods of a test class~~ [SPARK-58480][PYTHON][TEST] Declare DataFrame golden test cases as methods of a test class Jul 31, 2026

uros-b approved these changes Jul 31, 2026

View reviewed changes

gaogaotiantian reviewed Jul 31, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[SPARK-58480][PYTHON][TEST] Declare DataFrame golden test cases as methods of a test class - #57691

[SPARK-58480][PYTHON][TEST] Declare DataFrame golden test cases as methods of a test class#57691
dtenedor wants to merge 1 commit into
apache:masterfrom
dtenedor:combine-golden-files

dtenedor commented Jul 31, 2026

Uh oh!

uros-b commented Jul 31, 2026

Uh oh!

gaogaotiantian left a comment

Uh oh!

gaogaotiantian Jul 31, 2026

Uh oh!

gaogaotiantian Jul 31, 2026

Uh oh!

gaogaotiantian Jul 31, 2026

Uh oh!

gaogaotiantian Jul 31, 2026

Uh oh!

gaogaotiantian Jul 31, 2026

Uh oh!

gaogaotiantian Jul 31, 2026

Uh oh!

gaogaotiantian Jul 31, 2026

Uh oh!

gaogaotiantian Jul 31, 2026

Uh oh!

gaogaotiantian Jul 31, 2026

Uh oh!

gaogaotiantian Jul 31, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		return os.environ.get("SPARK_GENERATE_GOLDEN_FILES") is not None


		def check_cases_in_sync(test_file, declared, golden):

Uh oh!

Conversation

dtenedor commented Jul 31, 2026

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

Was this patch authored or co-authored using generative AI tooling?

Uh oh!

uros-b commented Jul 31, 2026

Uh oh!

gaogaotiantian left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants