This is an automated email from the ASF dual-hosted git repository.
gaogaotiantian pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new da9ea7d293b6 [SPARK-58393][PYTHON][INFRA] Smart select tests for
PySpark test framework
da9ea7d293b6 is described below
commit da9ea7d293b6ae994fad7215d357a8ae8ccfda4d
Author: Tian Gao <[email protected]>
AuthorDate: Thu Aug 13 21:59:02 2026 -0700
[SPARK-58393][PYTHON][INFRA] Smart select tests for PySpark test framework
### What changes were proposed in this pull request?
Add smart test selection for pyspark test framework. The approach is:
1. In `run-tests.py`, check changed files. Rule out the known irrelevant
files (that belongs to `dev_tool` module). If the rest only contains `pyspark`
files, and we are running CI for post-merge of apache/spark, trigger this
mechanism by
a. get the relevant changed file list
b. write that list to a temp file
c. pass the path to the temp file to `python/run-tests.py`
2. In `python/run-tests.py`, if changed file path is received, add that to
the env var so all the subprocesses get it.
3. For all PySpark tests, in `setUpClass`, try to read that env var (which
is super cheap). If the env var exists, get the list, build a pyspark graph
(sub-second cost per module) and analyze if the module of the test class is
impacted by the changed files.
4. If the test is not impacted, skip it.
### Why are the changes needed?
We will only run impacted tests on python-only changes by accurately
analyzing the changed files. This will save us plenty of CI time.
### Does this PR introduce _any_ user-facing change?
In theory, `pyspark/testing` is public to users. However, this mechanism is
for pyspark tests and is gated by the env var. It should not impact any users
that do not want it.
### How was this patch tested?
A feature test file is added for multiple cases. Some manual tests are done
locally. However, whether it works for E2E scenarios, we need to observe the
post-merge CIs for python-only changes.
The changed code is also reviewed by Claude Code (Opus 4.8)
### Was this patch authored or co-authored using generative AI tooling?
Only the test file is generated by LLM (with some tuning so it's better
organized).
Closes #57587 from gaogaotiantian/changed-files-aware.
Authored-by: Tian Gao <[email protected]>
Signed-off-by: Tian Gao <[email protected]>
(cherry picked from commit eb7a51efc062474c7916042418814d4998059aad)
Signed-off-by: Tian Gao <[email protected]>
---
dev/run-tests.py | 30 ++++++-
dev/sparktestsupport/modules.py | 1 +
pyproject.toml | 2 +
python/pyspark/testing/tests/test_changed_files.py | 98 ++++++++++++++++++++++
python/pyspark/testing/utils.py | 52 ++++++++++++
python/run-tests.py | 12 +++
6 files changed, 193 insertions(+), 2 deletions(-)
diff --git a/dev/run-tests.py b/dev/run-tests.py
index d50307513bb0..a278fe77a7be 100755
--- a/dev/run-tests.py
+++ b/dev/run-tests.py
@@ -23,6 +23,7 @@ import os
import re
import sys
import subprocess
+import tempfile
from contextlib import contextmanager
from sparktestsupport import SPARK_HOME, USER_HOME, ERROR_CODES
@@ -392,7 +393,9 @@ def run_scala_tests(build_tool, extra_profiles,
test_modules, excluded_tags, inc
run_scala_tests_sbt(test_modules, test_profiles)
-def run_python_tests(test_modules, test_pythons, parallelism,
with_coverage=False):
+def run_python_tests(
+ test_modules, test_pythons, parallelism, changed_files=None,
with_coverage=False
+):
set_title_and_block("Running PySpark tests", "BLOCK_PYSPARK_UNIT_TESTS")
if with_coverage:
@@ -408,7 +411,14 @@ def run_python_tests(test_modules, test_pythons,
parallelism, with_coverage=Fals
command.append("--modules=%s" % ",".join(m.name for m in test_modules))
command.append("--parallelism=%i" % parallelism)
command.append("--python-executables=%s" % test_pythons)
- run_cmd(command)
+ if changed_files:
+ with tempfile.NamedTemporaryFile("w") as f:
+ f.write("\n".join(changed_files))
+ f.flush()
+ command.append("--changed-files=%s" % f.name)
+ run_cmd(command)
+ else:
+ run_cmd(command)
def run_python_packaging_tests():
@@ -662,10 +672,26 @@ def main():
modules_with_python_tests = [m for m in test_modules if
m.python_test_goals]
if modules_with_python_tests and not os.environ.get("SKIP_PYTHON"):
+ relevant_changed_files = None
+ # We only do smart test selection on push action of apache/spark
+ # If APACHE_SPARK_REF is set, we are in a forked repository.
+ # Otherwise if we have a list of changed files, we must be in
post-merge CI.
+ if not os.environ.get("APACHE_SPARK_REF", "") and changed_files:
+ # Filter out all dev_tools files because they are not relevant
+ relevant_changed_files = [
+ f for f in changed_files if not
modules.dev_tools.contains_file(f)
+ ]
+ # If there are relevant changed files that are not pyspark, we
don't do smart test
+ if any(
+ not (f.endswith(".py") and f.startswith("python/pyspark/"))
+ for f in relevant_changed_files
+ ):
+ relevant_changed_files = None
run_python_tests(
modules_with_python_tests,
opts.python_executables,
opts.parallelism,
+ changed_files=relevant_changed_files,
with_coverage=os.environ.get("PYSPARK_CODECOV", "false") == "true",
)
run_python_packaging_tests()
diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py
index f5edb21c09fb..e70a5983d931 100644
--- a/dev/sparktestsupport/modules.py
+++ b/dev/sparktestsupport/modules.py
@@ -675,6 +675,7 @@ pyspark_testing = Module(
"pyspark.testing.utils",
"pyspark.testing.pandasutils",
# unittests
+ "pyspark.testing.tests.test_changed_files",
"pyspark.testing.tests.test_fail",
"pyspark.testing.tests.test_fail_in_set_up_class",
"pyspark.testing.tests.test_no_tests",
diff --git a/pyproject.toml b/pyproject.toml
index d761759bb199..02d723547a05 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -162,6 +162,8 @@ internal_lint = [
internal_test = [
# Dependency for running tests
"unittest-xml-reporting",
+ # Import graph support for smart test selection
+ "grimp",
# Coverage tool
"coverage",
# 3rd party libraries that unittests use.
diff --git a/python/pyspark/testing/tests/test_changed_files.py
b/python/pyspark/testing/tests/test_changed_files.py
new file mode 100644
index 000000000000..c142c6afc917
--- /dev/null
+++ b/python/pyspark/testing/tests/test_changed_files.py
@@ -0,0 +1,98 @@
+#
+# 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.
+#
+import os
+import tempfile
+import unittest
+
+from pyspark.testing.utils import (
+ PySparkBaseTestCase,
+ have_grimp,
+ grimp_requirement_message,
+)
+
+# A test module and modules it does / does not import (transitively). These
relationships are
+# stable parts of the pyspark import graph.
+_TEST_MODULE = "pyspark.sql.tests.test_functions"
+_RELEVANT_FILE = "python/pyspark/sql/functions/builtin.py" # imported by
_TEST_MODULE
+_IRRELEVANT_FILE = "python/pyspark/ml/classification.py" # not imported by
_TEST_MODULE
+
+
[email protected](not have_grimp, grimp_requirement_message)
+class ChangedFilesSelectionTests(unittest.TestCase):
+ """Tests for the "smart" test selection driven by PYSPARK_CHANGED_FILES.
+
+ ``PySparkBaseTestCase.skip_if_changed_files_irrelevant`` skips a test
class when none of the
+ changed files' modules are reachable from the test's own module in the
pyspark import graph.
+ These exercise that logic directly against real pyspark modules with known
relationships.
+ """
+
+ def setUp(self):
+ # The relevance check is memoized with functools.cache; clear it so
each case starts fresh.
+ PySparkBaseTestCase._is_module_relevant_to_changed_files.cache_clear()
+
+ def _is_relevant(self, module, files):
+ with tempfile.NamedTemporaryFile("w") as f:
+ f.write("\n".join(files))
+ f.flush()
+ return
PySparkBaseTestCase._is_module_relevant_to_changed_files(module, f.name)
+
+ def test_relevance(self):
+ own_file = "python/" + _TEST_MODULE.replace(".", os.path.sep) + ".py"
+ cases = [
+ ("imported module is relevant", [_RELEVANT_FILE], True),
+ ("unimported module is irrelevant", [_IRRELEVANT_FILE], False),
+ ("relevant among irrelevant is relevant", [_IRRELEVANT_FILE,
_RELEVANT_FILE], True),
+ ("the test's own file is relevant via the self-module short
circuit", [own_file], True),
+ (
+ "non-pyspark files are conservatively relevant",
+ ["sql/core/src/main/scala/Foo.scala"],
+ True,
+ ),
+ (
+ "package __init__ with no graph node is conservatively
relevant",
+ ["python/pyspark/sql/__init__.py"],
+ True,
+ ),
+ ]
+ for desc, files, expected in cases:
+ with self.subTest(desc):
+ self.assertEqual(self._is_relevant(_TEST_MODULE, files),
expected)
+
+ def test_skip_if_changed_files_irrelevant(self):
+ class _Dummy(PySparkBaseTestCase):
+ pass
+
+ _Dummy.__module__ = _TEST_MODULE
+
+ # Irrelevant changes raise SkipTest.
+ with tempfile.NamedTemporaryFile("w") as f:
+ f.write(_IRRELEVANT_FILE)
+ f.flush()
+ with self.assertRaises(unittest.SkipTest):
+ _Dummy.skip_if_changed_files_irrelevant(f.name)
+
+ # Relevant changes do not.
+ with tempfile.NamedTemporaryFile("w") as f:
+ f.write(_RELEVANT_FILE)
+ f.flush()
+ _Dummy.skip_if_changed_files_irrelevant(f.name)
+
+
+if __name__ == "__main__":
+ from pyspark.testing import main
+
+ main()
diff --git a/python/pyspark/testing/utils.py b/python/pyspark/testing/utils.py
index cb549af75f3c..84870585c5c6 100644
--- a/python/pyspark/testing/utils.py
+++ b/python/pyspark/testing/utils.py
@@ -85,6 +85,9 @@ graphviz_requirement_message = "" if have_graphviz else "No
module named 'graphv
have_flameprof = have_package("flameprof")
flameprof_requirement_message = "" if have_flameprof else "No module named
'flameprof'"
+have_grimp = have_package("grimp")
+grimp_requirement_message = "" if have_grimp else "No module named 'grimp'"
+
have_jinja2 = have_package("jinja2")
jinja2_requirement_message = "" if have_jinja2 else "No module named 'jinja2'"
@@ -294,9 +297,58 @@ class QuietTest:
class PySparkBaseTestCase(unittest.TestCase):
@classmethod
def setUpClass(cls):
+ if have_grimp and (path := os.environ.get("PYSPARK_CHANGED_FILES")):
+ # PYSPARK_CHANGED_FILES should only be used when ONLY pyspark
files are changed.
+ # If other files (JVM for example) are changed, do NOT set this.
+ cls.skip_if_changed_files_irrelevant(path)
+
if os.environ.get("PYSPARK_TEST_TIMEOUT"):
faulthandler.register(signal.SIGTERM, file=sys.__stderr__,
all_threads=True)
+ @classmethod
+ def skip_if_changed_files_irrelevant(cls, path: str) -> None:
+ module = cls.__module__
+ if module == "__main__":
+ mod = sys.modules["__main__"]
+ if mod.__spec__ and mod.__spec__.name:
+ module = mod.__spec__.name
+ else:
+ return
+
+ if not cls._is_module_relevant_to_changed_files(module, path):
+ raise unittest.SkipTest("Skipping test because changed files are
irrelevant")
+
+ @staticmethod
+ @functools.cache
+ def _is_module_relevant_to_changed_files(module: str, path: str) -> bool:
+ import grimp
+
+ with open(path, "r") as f:
+ changed_files = f.read().strip().splitlines()
+
+ if not all(f.startswith("python/pyspark/") and f.endswith(".py") for f
in changed_files):
+ # We have a wrong list of files, just run the test.
+ return True
+
+ changed_modules = [
+ f.removeprefix("python/").rsplit(".", 1)[0].replace(os.path.sep,
".")
+ for f in changed_files
+ ]
+
+ graph = grimp.build_graph("pyspark")
+
+ for changed_module in changed_modules:
+ if changed_module == module:
+ return True
+ try:
+ if graph.chain_exists(module, changed_module):
+ return True
+ except Exception:
+ # Any exception, we just be conservative and run the test.
+ return True
+
+ return False
+
@classmethod
def tearDownClass(cls):
if os.environ.get("PYSPARK_TEST_TIMEOUT"):
diff --git a/python/run-tests.py b/python/run-tests.py
index df9af0d4f3fe..b52ae8f5b5b4 100755
--- a/python/run-tests.py
+++ b/python/run-tests.py
@@ -420,6 +420,12 @@ def parse_opts():
default=4,
help="The number of suites to test in parallel (default %(default)d)",
)
+ parser.add_argument(
+ "--changed-files",
+ type=str,
+ default=None,
+ help="A file containing a list of changed files (default:
%(default)s)",
+ )
parser.add_argument("--verbose", action="store_true", help="Enable
additional debug logging")
group = parser.add_argument_group("Developer Options")
@@ -485,6 +491,12 @@ def main():
python_execs = opts.python_executables.split(",")
LOGGER.info("Will test against the following Python executables: %s",
python_execs)
+ if opts.changed_files:
+ with open(opts.changed_files, "r") as f:
+ changed_files = f.read().splitlines()
+ os.environ["PYSPARK_CHANGED_FILES"] = opts.changed_files
+ LOGGER.info("Will select tests based on the following changed files:
%s", changed_files)
+
if should_test_modules:
modules_to_test = []
for module_name in opts.modules.split(","):
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
