[SPARK-58332][PYTHON][TEST] Move compare_or_generate_golden_matrix into GoldenFileTestMixin ### What changes were proposed in this pull request? `compare_or_generate_golden_matrix` was duplicated verbatim across three PyArrow golden-file test files: - `python/pyspark/tests/upstream/pyarrow/test_pyarrow_array_cast.py` - `python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_default.py` - `python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_non_default.py` This PR moves it into `GoldenFileTestMixin` (`python/pyspark/testing/goldenutils.py`), which all three suites already inherit, and removes the local copies. Imports that became unused after the removal (`inspect`, `os`, `typing.Callable/List/Optional`) are dropped from the test files. This is a follow-up to #57435, where reviewers asked to centralize the duplicated matrix driver into the mixin. ### Why are the changes needed? Removes duplicated test machinery so the golden-file matrix driver has a single implementation, making it easier to maintain and reuse for future golden-file suites. ### Does this PR introduce _any_ user-facing change? No. Test-only, behavior-preserving refactor. ### How was this patch tested? Existing suites pass in compare mode (no golden files regenerated): ``` python -m pytest \ python/pyspark/tests/upstream/pyarrow/test_pyarrow_array_cast.py \ python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_default.py \ python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_non_default.py ``` ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 4.8) Closes #57510 from Spenserrrr/centralize-golden-matrix-helper. Authored-by: Spenser Sun <haotian.sun@databricks.com> Signed-off-by: Ruifeng Zheng <ruifengz@apache.org>
diff --git a/python/pyspark/testing/goldenutils.py b/python/pyspark/testing/goldenutils.py index b1e4af5..d0a191e 100644 --- a/python/pyspark/testing/goldenutils.py +++ b/python/pyspark/testing/goldenutils.py
@@ -15,7 +15,8 @@ # limitations under the License. # -from typing import Any, Optional +from typing import Any, Callable, List, Optional +import inspect import os import time @@ -194,6 +195,85 @@ "Install 'tabulate' package to generate markdown files." ) + def compare_or_generate_golden_matrix( + self, + row_names: List[str], + col_names: List[str], + compute_cell: Callable[[str, str], str], + golden_file_prefix: str, + index_name: str = "source \\ target", + overrides: Optional[dict[tuple[str, str], str]] = None, + ) -> None: + """ + Run a matrix of computations and compare against (or generate) a golden file. + + 1. If SPARK_GENERATE_GOLDEN_FILES=1, compute every cell, build a + DataFrame, and save it as the new golden CSV / Markdown file. + 2. Otherwise, load the existing golden file and assert that every cell + matches the freshly computed value. + + Parameters + ---------- + row_names : list[str] + Ordered row labels (becomes the DataFrame index). + col_names : list[str] + Ordered column labels. + compute_cell : (row_name, col_name) -> str + Function that computes the string result for one cell. + golden_file_prefix : str + Prefix for the golden CSV/MD files (without extension). + Files are placed in the same directory as the concrete test file. + index_name : str, default "source \\ target" + Name for the index column in the golden file. + overrides : dict[(row, col) -> str], optional + Version-specific expected values that take precedence over the golden + file. Use this to document known behavioral differences across + library versions (e.g. PyArrow 18 vs 22) directly in the test code, + so that the same golden file works for multiple versions. + """ + generating = self.is_generating_golden() + + test_dir = os.path.dirname(inspect.getfile(type(self))) + golden_csv = os.path.join(test_dir, f"{golden_file_prefix}.csv") + golden_md = os.path.join(test_dir, f"{golden_file_prefix}.md") + + golden = None + if not generating: + golden = self.load_golden_csv(golden_csv) + + errors = [] + results = {} + + for row_name in row_names: + for col_name in col_names: + result = compute_cell(row_name, col_name) + results[(row_name, col_name)] = result + + if not generating: + if overrides and (row_name, col_name) in overrides: + expected = overrides[(row_name, col_name)] + else: + expected = golden.loc[row_name, col_name] + if expected != result: + errors.append( + f"{row_name} -> {col_name}: expected '{expected}', got '{result}'" + ) + + if generating: + import pandas as pd + + index = pd.Index(row_names, name=index_name) + df = pd.DataFrame(index=index) + for col_name in col_names: + df[col_name] = [results[(row, col_name)] for row in row_names] + self.save_golden(df, golden_csv, golden_md) + else: + self.assertEqual( + len(errors), + 0, + f"\n{len(errors)} golden file mismatches:\n" + "\n".join(errors), + ) + @staticmethod def repr_type(t: Any) -> str: """
diff --git a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_array_cast.py b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_array_cast.py index 6c48e4e..128557f 100644 --- a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_array_cast.py +++ b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_array_cast.py
@@ -55,12 +55,9 @@ | pa.array(floats, pa.float16()) natively | requires numpy | requires numpy | native | """ -import inspect -import os import platform import unittest from decimal import Decimal -from typing import Callable, List, Optional from pyspark.loose_version import LooseVersion from pyspark.testing.utils import ( @@ -134,85 +131,6 @@ except Exception as e: return f"ERR@{type(e).__name__}" - def compare_or_generate_golden_matrix( - self, - row_names: List[str], - col_names: List[str], - compute_cell: Callable[[str, str], str], - golden_file_prefix: str, - index_name: str = "source \\ target", - overrides: Optional[dict[tuple[str, str], str]] = None, - ) -> None: - """ - Run a matrix of computations and compare against (or generate) a golden file. - - 1. If SPARK_GENERATE_GOLDEN_FILES=1, compute every cell, build a - DataFrame, and save it as the new golden CSV / Markdown file. - 2. Otherwise, load the existing golden file and assert that every cell - matches the freshly computed value. - - Parameters - ---------- - row_names : list[str] - Ordered row labels (becomes the DataFrame index). - col_names : list[str] - Ordered column labels. - compute_cell : (row_name, col_name) -> str - Function that computes the string result for one cell. - golden_file_prefix : str - Prefix for the golden CSV/MD files (without extension). - Files are placed in the same directory as the concrete test file. - index_name : str, default "source \\ target" - Name for the index column in the golden file. - overrides : dict[(row, col) -> str], optional - Version-specific expected values that take precedence over the golden - file. Use this to document known behavioral differences across - library versions (e.g. PyArrow 18 vs 22) directly in the test code, - so that the same golden file works for multiple versions. - """ - generating = self.is_generating_golden() - - test_dir = os.path.dirname(inspect.getfile(type(self))) - golden_csv = os.path.join(test_dir, f"{golden_file_prefix}.csv") - golden_md = os.path.join(test_dir, f"{golden_file_prefix}.md") - - golden = None - if not generating: - golden = self.load_golden_csv(golden_csv) - - errors = [] - results = {} - - for row_name in row_names: - for col_name in col_names: - result = compute_cell(row_name, col_name) - results[(row_name, col_name)] = result - - if not generating: - if overrides and (row_name, col_name) in overrides: - expected = overrides[(row_name, col_name)] - else: - expected = golden.loc[row_name, col_name] - if expected != result: - errors.append( - f"{row_name} -> {col_name}: expected '{expected}', got '{result}'" - ) - - if generating: - import pandas as pd - - index = pd.Index(row_names, name=index_name) - df = pd.DataFrame(index=index) - for col_name in col_names: - df[col_name] = [results[(row, col_name)] for row in row_names] - self.save_golden(df, golden_csv, golden_md) - else: - self.assertEqual( - len(errors), - 0, - f"\n{len(errors)} golden file mismatches:\n" + "\n".join(errors), - ) - # ============================================================ # Scalar Type Cast Tests
diff --git a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_default.py b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_default.py index 566d10a..e3a9a0c 100644 --- a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_default.py +++ b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_default.py
@@ -51,11 +51,8 @@ """ import datetime -import inspect -import os import unittest from decimal import Decimal -from typing import Callable, List, Optional from pyspark.loose_version import LooseVersion from pyspark.testing.utils import ( @@ -87,66 +84,6 @@ Each type is tested both without and with null values. """ - def compare_or_generate_golden_matrix( - self, - row_names: List[str], - col_names: List[str], - compute_cell: Callable[[str, str], str], - golden_file_prefix: str, - index_name: str = "source \\ target", - overrides: Optional[dict[tuple[str, str], str]] = None, - ) -> None: - """ - Run a matrix of computations and compare against (or generate) a golden file. - - 1. If SPARK_GENERATE_GOLDEN_FILES=1, compute every cell, build a - DataFrame, and save it as the new golden CSV / Markdown file. - 2. Otherwise, load the existing golden file and assert that every cell - matches the freshly computed value. - """ - generating = self.is_generating_golden() - - test_dir = os.path.dirname(inspect.getfile(type(self))) - golden_csv = os.path.join(test_dir, f"{golden_file_prefix}.csv") - golden_md = os.path.join(test_dir, f"{golden_file_prefix}.md") - - golden = None - if not generating: - golden = self.load_golden_csv(golden_csv) - - errors = [] - results = {} - - for row_name in row_names: - for col_name in col_names: - result = compute_cell(row_name, col_name) - results[(row_name, col_name)] = result - - if not generating: - if overrides and (row_name, col_name) in overrides: - expected = overrides[(row_name, col_name)] - else: - expected = golden.loc[row_name, col_name] - if expected != result: - errors.append( - f"{row_name} -> {col_name}: expected '{expected}', got '{result}'" - ) - - if generating: - import pandas as pd - - index = pd.Index(row_names, name=index_name) - df = pd.DataFrame(index=index) - for col_name in col_names: - df[col_name] = [results[(row, col_name)] for row in row_names] - self.save_golden(df, golden_csv, golden_md) - else: - self.assertEqual( - len(errors), - 0, - f"\n{len(errors)} golden file mismatches:\n" + "\n".join(errors), - ) - def _build_source_arrays(self): """Build an ordered dict of named source PyArrow arrays for testing.""" import pyarrow as pa
diff --git a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_non_default.py b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_non_default.py index 6587f5d..ddfc8f7 100644 --- a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_non_default.py +++ b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_non_default.py
@@ -60,10 +60,7 @@ """ import datetime -import inspect -import os import unittest -from typing import Callable, List, Optional from pyspark.loose_version import LooseVersion from pyspark.testing.utils import ( @@ -91,66 +88,6 @@ This base defines no ``test_*`` methods, so it contributes no tests itself. """ - def compare_or_generate_golden_matrix( - self, - row_names: List[str], - col_names: List[str], - compute_cell: Callable[[str, str], str], - golden_file_prefix: str, - index_name: str = "source \\ target", - overrides: Optional[dict[tuple[str, str], str]] = None, - ) -> None: - """ - Run a matrix of computations and compare against (or generate) a golden file. - - 1. If SPARK_GENERATE_GOLDEN_FILES=1, compute every cell, build a - DataFrame, and save it as the new golden CSV / Markdown file. - 2. Otherwise, load the existing golden file and assert that every cell - matches the freshly computed value. - """ - generating = self.is_generating_golden() - - test_dir = os.path.dirname(inspect.getfile(type(self))) - golden_csv = os.path.join(test_dir, f"{golden_file_prefix}.csv") - golden_md = os.path.join(test_dir, f"{golden_file_prefix}.md") - - golden = None - if not generating: - golden = self.load_golden_csv(golden_csv) - - errors = [] - results = {} - - for row_name in row_names: - for col_name in col_names: - result = compute_cell(row_name, col_name) - results[(row_name, col_name)] = result - - if not generating: - if overrides and (row_name, col_name) in overrides: - expected = overrides[(row_name, col_name)] - else: - expected = golden.loc[row_name, col_name] - if expected != result: - errors.append( - f"{row_name} -> {col_name}: expected '{expected}', got '{result}'" - ) - - if generating: - import pandas as pd - - index = pd.Index(row_names, name=index_name) - df = pd.DataFrame(index=index) - for col_name in col_names: - df[col_name] = [results[(row, col_name)] for row in row_names] - self.save_golden(df, golden_csv, golden_md) - else: - self.assertEqual( - len(errors), - 0, - f"\n{len(errors)} golden file mismatches:\n" + "\n".join(errors), - ) - def _to_pandas_cell(self, arr, **to_pandas_kwargs) -> str: """ Convert ``arr`` via ``to_pandas(**to_pandas_kwargs)`` and format the