ci: key TPC dataset caches on pinned generators

.github/workflows/README.md

@@ -44,8 +44,10 @@ merge. Three jobs are still on `push`: `docs`, because it deploys to
 `pr_build_macos`, because of `actions/cache` scoping. A pull request can only restore caches saved on its
 own branch or on `main`, and the queue runs on a throwaway
 `gh-readonly-queue/*` branch whose caches are deleted with it. Without a push
-run, a `Cargo.lock` or `pom.xml` change would leave the cargo-ci, Maven and
-TPC-H/TPC-DS caches on `main` stale until the next unrelated change.
+run, a `Cargo.lock` or `pom.xml` change would leave the cargo-ci and Maven
+caches on `main` stale until the next unrelated change. The TPC-H and TPC-DS
+datasets do not depend on those files; the push run re-warms them only after
+eviction.
 
 Warming those caches is the only thing the push run is for, so on `push` the
 Linux build runs in **cache-refresh-only** mode: `build-native` and the two
@@ -511,15 +513,18 @@ any of the paths above, and rejects a `save` that is missing the `github.ref`
 guard. `publish_snapshot.yml` is exempt in `CACHE_SAVE_SCOPE_EXEMPT`, because
 it runs from main on a schedule already.
 
-The TPC-H and TPC-DS dataset caches keep the read-write form and are out of
-scope entirely: `./tpch` and `./tpcds-sf-1` are a few hundred MB, they are not
-dependency trees, and they are keyed on this workflow file, so a pull request
-that edits it would regenerate the data on every run rather than once.
+The TPC-H and TPC-DS dataset caches also use separate restore and main-only
+save steps. Their keys include the generator (`GenTPCHData.scala` for TPC-H,
+the pinned `tpcds-kit` commit for TPC-DS), scale factor and partition count.
+The generation arguments, keys and dataset paths use the same job environment
+values, so changing an input invalidates the dataset without unrelated workflow
+edits rotating its key. `dev/ci/check-ci-config.py` guards this coupling.
 
 A job that only ever runs on a pull request or in the queue keeps the guard
 anyway, and so never writes. That is deliberate — it restores from main's
-entry through `restore-keys` and downloads whatever else it needs, which is
-what a cold pull request already did. See the push-tier discussion above for
+entry through `restore-keys` (the TPC dataset restores require an exact key
+instead) and downloads whatever else it needs, which is what a cold pull
+request already did. See the push-tier discussion above for
 which jobs do run on main and therefore do write.
 
 ## Retrying flaky network operations

.github/workflows/pr_build_linux.yml

@@ -768,6 +768,8 @@ jobs:
     container:
       image: amd64/rust
     env:
+      TPC_SCALE_FACTOR: 1
+      TPC_NUM_PARTITIONS: 1
       JAVA_TOOL_OPTIONS: --add-exports=java.base/sun.nio.ch=ALL-UNNAMED --add-exports=java.base/sun.util.calendar=ALL-UNNAMED --add-opens=java.base/java.lang=ALL-UNNAMED --add-opens=java.base/java.lang.invoke=ALL-UNNAMED --add-opens=java.base/java.lang.reflect=ALL-UNNAMED --add-opens=java.base/java.io=ALL-UNNAMED --add-opens=java.base/java.net=ALL-UNNAMED --add-opens=java.base/java.nio=ALL-UNNAMED --add-opens=java.base/java.util=ALL-UNNAMED --add-opens=java.base/java.util.concurrent=ALL-UNNAMED --add-opens=java.base/java.util.concurrent.atomic=ALL-UNNAMED --add-opens=java.base/jdk.internal.ref=ALL-UNNAMED --add-opens=java.base/sun.nio.ch=ALL-UNNAMED --add-opens=java.base/sun.nio.cs=ALL-UNNAMED --add-opens=java.base/sun.security.action=ALL-UNNAMED --add-opens=java.base/sun.util.calendar=ALL-UNNAMED
     steps:
       - uses: actions/checkout@v7
@@ -795,26 +797,37 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-java-maven-
 
-      - name: Cache TPC-H data
+      - name: Restore TPC-H data
         id: cache-tpch
-        uses: actions/cache@v6
+        uses: actions/cache/restore@v6
         with:
           path: ./tpch
-          key: tpch-${{ hashFiles('.github/workflows/pr_build_linux.yml') }}
+          key: tpch-sf${{ env.TPC_SCALE_FACTOR }}-p${{ env.TPC_NUM_PARTITIONS }}-${{ hashFiles('spark/src/test/scala/org/apache/spark/sql/GenTPCHData.scala') }}
 
       - name: Build project
         run: |
           ./mvnw -B -Prelease install -DskipTests
 
-      - name: Generate TPC-H data (SF=1)
+      - name: Generate TPC-H data
         if: steps.cache-tpch.outputs.cache-hit != 'true'
         run: |
-          cd spark && MAVEN_OPTS='-Xmx20g' ../mvnw -B -Prelease exec:java -Dexec.mainClass="org.apache.spark.sql.GenTPCHData" -Dexec.classpathScope="test" -Dexec.cleanupDaemonThreads="false" -Dexec.args="--location `pwd`/.. --scaleFactor 1 --numPartitions 1 --overwrite"
+          cd spark && MAVEN_OPTS='-Xmx20g' ../mvnw -B -Prelease exec:java -Dexec.mainClass="org.apache.spark.sql.GenTPCHData" -Dexec.classpathScope="test" -Dexec.cleanupDaemonThreads="false" -Dexec.args="--location `pwd`/.. --scaleFactor $TPC_SCALE_FACTOR --numPartitions $TPC_NUM_PARTITIONS --overwrite"
 
       - name: Run TPC-H queries
         if: ${{ !inputs.cache-refresh-only }}
         run: |
-          SPARK_HOME=`pwd` SPARK_TPCH_DATA=`pwd`/tpch/sf1_parquet ./mvnw -B -Prelease -Dsuites=org.apache.spark.sql.CometTPCHQuerySuite test
+          SPARK_HOME=`pwd` SPARK_TPCH_DATA=`pwd`/tpch/sf${TPC_SCALE_FACTOR}_parquet ./mvnw -B -Prelease -Dsuites=org.apache.spark.sql.CometTPCHQuerySuite test
+
+      # The dataset key includes the generator, scale factor and partition count.
+      # Main caches it like any other large entry. An entry written from a pull
+      # request or a `gh-readonly-queue/*` branch can never be restored by a
+      # later run and evicts main's from the shared budget.
+      - name: Save TPC-H data
+        if: ${{ github.ref == 'refs/heads/main' && steps.cache-tpch.outputs.cache-hit != 'true' }}
+        uses: actions/cache/save@v6
+        with:
+          path: ./tpch
+          key: tpch-sf${{ env.TPC_SCALE_FACTOR }}-p${{ env.TPC_NUM_PARTITIONS }}-${{ hashFiles('spark/src/test/scala/org/apache/spark/sql/GenTPCHData.scala') }}
 
       # Saved only on main: an entry written from a pull request or a
       # `gh-readonly-queue/*` branch cannot be restored by any later run, and
@@ -842,6 +855,8 @@ jobs:
     container:
       image: amd64/rust
     env:
+      TPC_SCALE_FACTOR: 1
+      TPC_NUM_PARTITIONS: 1
       JAVA_TOOL_OPTIONS: --add-exports=java.base/sun.nio.ch=ALL-UNNAMED --add-exports=java.base/sun.util.calendar=ALL-UNNAMED --add-opens=java.base/java.lang=ALL-UNNAMED --add-opens=java.base/java.lang.invoke=ALL-UNNAMED --add-opens=java.base/java.lang.reflect=ALL-UNNAMED --add-opens=java.base/java.io=ALL-UNNAMED --add-opens=java.base/java.net=ALL-UNNAMED --add-opens=java.base/java.nio=ALL-UNNAMED --add-opens=java.base/java.util=ALL-UNNAMED --add-opens=java.base/java.util.concurrent=ALL-UNNAMED --add-opens=java.base/java.util.concurrent.atomic=ALL-UNNAMED --add-opens=java.base/jdk.internal.ref=ALL-UNNAMED --add-opens=java.base/sun.nio.ch=ALL-UNNAMED --add-opens=java.base/sun.nio.cs=ALL-UNNAMED --add-opens=java.base/sun.security.action=ALL-UNNAMED --add-opens=java.base/sun.util.calendar=ALL-UNNAMED
     steps:
       - uses: actions/checkout@v7
@@ -869,12 +884,12 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-java-maven-
 
-      - name: Cache TPC-DS data
+      - name: Restore TPC-DS data
         id: cache-tpcds
-        uses: actions/cache@v6
+        uses: actions/cache/restore@v6
         with:
-          path: ./tpcds-sf-1
-          key: tpcds-${{ hashFiles('.github/workflows/pr_build_linux.yml') }}
+          path: ./tpcds-sf-${{ env.TPC_SCALE_FACTOR }}
+          key: tpcds-sf${{ env.TPC_SCALE_FACTOR }}-p${{ env.TPC_NUM_PARTITIONS }}-1b7fb7529eda
 
       - name: Build project
         run: |
@@ -885,6 +900,7 @@ jobs:
         uses: actions/checkout@v7
         with:
           repository: databricks/tpcds-kit
+          ref: 1b7fb7529edae091684201fab142d956d6afd881
           path: ./tpcds-kit
 
       - name: Build tpcds-kit
@@ -894,15 +910,15 @@ jobs:
           update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 120 --slave /usr/bin/g++ g++ /usr/bin/g++-12
           cd tpcds-kit/tools && make OS=LINUX
 
-      - name: Generate TPC-DS data (SF=1)
+      - name: Generate TPC-DS data
         if: steps.cache-tpcds.outputs.cache-hit != 'true'
         run: |
-          cd spark && MAVEN_OPTS='-Xmx20g' ../mvnw -B -Prelease exec:java -Dexec.mainClass="org.apache.spark.sql.GenTPCDSData" -Dexec.classpathScope="test" -Dexec.cleanupDaemonThreads="false" -Dexec.args="--dsdgenDir `pwd`/../tpcds-kit/tools --location `pwd`/../tpcds-sf-1 --scaleFactor 1 --numPartitions 1"
+          cd spark && MAVEN_OPTS='-Xmx20g' ../mvnw -B -Prelease exec:java -Dexec.mainClass="org.apache.spark.sql.GenTPCDSData" -Dexec.classpathScope="test" -Dexec.cleanupDaemonThreads="false" -Dexec.args="--dsdgenDir `pwd`/../tpcds-kit/tools --location `pwd`/../tpcds-sf-$TPC_SCALE_FACTOR --scaleFactor $TPC_SCALE_FACTOR --numPartitions $TPC_NUM_PARTITIONS"
 
       - name: Run TPC-DS queries (Sort merge join)
         if: ${{ !inputs.cache-refresh-only }}
         run: |
-          SPARK_HOME=`pwd` SPARK_TPCDS_DATA=`pwd`/tpcds-sf-1 ./mvnw -B -Prelease -Dsuites=org.apache.spark.sql.CometTPCDSQuerySuite test
+          SPARK_HOME=`pwd` SPARK_TPCDS_DATA=`pwd`/tpcds-sf-$TPC_SCALE_FACTOR ./mvnw -B -Prelease -Dsuites=org.apache.spark.sql.CometTPCDSQuerySuite test
         env:
           SPARK_TPCDS_JOIN_CONF: |
             spark.sql.autoBroadcastJoinThreshold=-1
@@ -911,20 +927,31 @@ jobs:
       - name: Run TPC-DS queries (Broadcast hash join)
         if: ${{ !inputs.cache-refresh-only }}
         run: |
-          SPARK_HOME=`pwd` SPARK_TPCDS_DATA=`pwd`/tpcds-sf-1 ./mvnw -B -Prelease -Dsuites=org.apache.spark.sql.CometTPCDSQuerySuite test
+          SPARK_HOME=`pwd` SPARK_TPCDS_DATA=`pwd`/tpcds-sf-$TPC_SCALE_FACTOR ./mvnw -B -Prelease -Dsuites=org.apache.spark.sql.CometTPCDSQuerySuite test
         env:
           SPARK_TPCDS_JOIN_CONF: |
             spark.sql.autoBroadcastJoinThreshold=10485760
 
       - name: Run TPC-DS queries (Shuffled hash join)
         if: ${{ !inputs.cache-refresh-only }}
         run: |
-          SPARK_HOME=`pwd` SPARK_TPCDS_DATA=`pwd`/tpcds-sf-1 ./mvnw -B -Prelease -Dsuites=org.apache.spark.sql.CometTPCDSQuerySuite test
+          SPARK_HOME=`pwd` SPARK_TPCDS_DATA=`pwd`/tpcds-sf-$TPC_SCALE_FACTOR ./mvnw -B -Prelease -Dsuites=org.apache.spark.sql.CometTPCDSQuerySuite test
         env:
           SPARK_TPCDS_JOIN_CONF: |
             spark.sql.autoBroadcastJoinThreshold=-1
             spark.sql.join.forceApplyShuffledHashJoin=true
 
+      # The dataset key includes the generator, scale factor and partition count.
+      # Main caches it like any other large entry. An entry written from a pull
+      # request or a `gh-readonly-queue/*` branch can never be restored by a
+      # later run and evicts main's from the shared budget.
+      - name: Save TPC-DS data
+        if: ${{ github.ref == 'refs/heads/main' && steps.cache-tpcds.outputs.cache-hit != 'true' }}
+        uses: actions/cache/save@v6
+        with:
+          path: ./tpcds-sf-${{ env.TPC_SCALE_FACTOR }}
+          key: tpcds-sf${{ env.TPC_SCALE_FACTOR }}-p${{ env.TPC_NUM_PARTITIONS }}-1b7fb7529eda
+
       # Saved only on main: an entry written from a pull request or a
       # `gh-readonly-queue/*` branch cannot be restored by any later run, and
       # it evicts main's from the shared budget. See "Large caches are written

dev/ci/check-ci-config.py

@@ -1595,6 +1595,123 @@ def check_cache_save_scope():
     return not failures
 
 
+TPC_WORKFLOW = WORKFLOWS / "pr_build_linux.yml"
+TPC_WORKFLOW_HASH = re.compile(
+    r"hashFiles\(\s*['\"]\.github/workflows/pr_build_linux\.yml['\"]\s*\)"
+)
+TPCDS_KIT_CHECKOUT = re.compile(
+    r"repository:[ \t]*databricks/tpcds-kit\n"
+    r"(?:[ \t]+[^\n]*\n)*?"
+    r"[ \t]+ref:[ \t]*([0-9a-f]{40})"
+)
+TPCH_CACHE_KEY = re.compile(
+    r"key:\s*tpch-sf.*-\$\{\{\s*hashFiles\("
+    r"'spark/src/test/scala/org/apache/spark/sql/GenTPCHData\.scala'\)\s*\}\}"
+)
+TPCDS_CACHE_KEY = re.compile(r"key:\s*tpcds-sf.*-([0-9a-f]{12})\b")
+GIT_COMMIT = re.compile(r"\b[0-9a-f]{40}\b")
+GIT_CHECKOUT = re.compile(r"git checkout")
+
+
+def check_tpc_dataset_caches():
+    """The TPC datasets are keyed on a pinned generator, not on the workflow file.
+
+    Both dataset caches used `hashFiles('.github/workflows/pr_build_linux.yml')`
+    as their key, and PR #5973 left them exempt from the main-only save rule
+    because `./tpch` and `./tpcds-sf-<scale factor>` are not dependency trees. The workflow
+    file is one of the most edited files in the repository, so the key rotated
+    with every CI change, orphaning main's entry and regenerating a
+    multi-hundred-megabyte dataset -- and the dataset kept competing with the
+    ~4.1 GiB `Linux-cargo-debug` cache whose eviction costs a ~26 minute cold
+    build. Key them on the generator input and pin the generators themselves, so
+    a dataset is regenerated only when the generator or its arguments change.
+    """
+    failures = []
+    text = TPC_WORKFLOW.read_text(encoding="utf-8")
+    workflow_hash = TPC_WORKFLOW_HASH.search(text)
+    if workflow_hash:
+        failures.append(
+            f"{TPC_WORKFLOW}: `{workflow_hash.group(0)}` keys a TPC dataset on "
+            f"this workflow file, so every CI edit rotates the cache key and "
+            f"regenerates the SF=1 dataset. Key it on the pinned generator"
+        )
+    if not TPCH_CACHE_KEY.search(text):
+        failures.append(
+            f"{TPC_WORKFLOW}: the TPC-H dataset cache key is not keyed on "
+            f"`GenTPCHData.scala`, so a generator change would not rotate it"
+        )
+    kit = TPCDS_KIT_CHECKOUT.search(text)
+    if not kit:
+        failures.append(
+            f"{TPC_WORKFLOW}: the `databricks/tpcds-kit` checkout has no pinned "
+            f"`ref:`, so the TPC-DS dataset can change without the key rotating"
+        )
+    else:
+        key = TPCDS_CACHE_KEY.search(text)
+        if not key or key.group(1) != kit.group(1)[:12]:
+            failures.append(
+                f"{TPC_WORKFLOW}: the TPC-DS dataset cache key does not embed the "
+                f"`tpcds-kit` revision `{kit.group(1)[:12]}`, so a generator bump "
+                f"would not rotate it"
+            )
+    for line_no, uses, step in _cache_steps(text.splitlines()):
+        if not CACHE_RW_USES.match(uses):
+            continue
+        body = "\n".join(step)
+        if "tpch" not in body and "tpcds" not in body:
+            continue
+        failures.append(
+            f"{TPC_WORKFLOW}:{line_no}: a TPC dataset cache uses "
+            f"`actions/cache@vN`, whose save runs in an implicit post step that "
+            f"no `if:` can reach. Split it into `actions/cache/restore` plus a "
+            f"main-only `actions/cache/save`"
+        )
+    for dataset in ("tpch", "tpcds"):
+        job = f"verify-benchmark-results-{dataset}"
+        block = re.search(rf"(?ms)^  {job}:\n(.*?)(?=^  [A-Za-z0-9_-]+:|\Z)", text)
+        if not block:
+            failures.append(f"{TPC_WORKFLOW}: missing dataset job `{job}`")
+            continue
+        body = block.group(1)
+        for option, variable in (
+            ("scaleFactor", "TPC_SCALE_FACTOR"),
+            ("numPartitions", "TPC_NUM_PARTITIONS"),
+        ):
+            if not re.search(rf"(?m)^      {variable}: [1-9][0-9]*$", body):
+                failures.append(f"{job}: define positive `{variable}` in job env")
+            if re.findall(rf"--{option}\s+([^\s\"]+)", body) != [f"${variable}"]:
+                failures.append(f"{job}: generator `--{option}` must use `${variable}`")
+        keys = re.findall(rf"(?m)^          key: ({dataset}-.*)$", body)
+        prefix = f"{dataset}-sf${{{{ env.TPC_SCALE_FACTOR }}}}-p${{{{ env.TPC_NUM_PARTITIONS }}}}-"
+        if len(keys) != 2 or keys[0] != keys[1] or not keys[0].startswith(prefix):
+            failures.append(
+                f"{job}: restore/save keys must match and include actual "
+                f"scale factor and partition count"
+            )
+        restore = re.search(
+            r"(?ms)^      - name: Restore TPC-[A-Z-]+ data\n(.*?)(?=^      - |\Z)", body
+        )
+        if restore and re.search(r"(?m)^          restore-keys:", restore.group(1)):
+            failures.append(
+                f"{job}: dataset restore must not fall back to different generation inputs"
+            )
+        for _, uses, step in _cache_steps(body.splitlines()):
+            step_body = "\n".join(step)
+            if CACHE_SAVE_USES.match(uses) and f"key: {dataset}-" in step_body:
+                if not CACHE_MAIN_GUARD.search(step_body):
+                    failures.append(f"{job}: dataset saves must be main-only")
+    generator = Path("spark/src/test/scala/org/apache/spark/sql/GenTPCHData.scala")
+    generator_text = generator.read_text(encoding="utf-8")
+    if not (GIT_COMMIT.search(generator_text) and GIT_CHECKOUT.search(generator_text)):
+        failures.append(
+            f"{generator}: the tpch-dbgen checkout is not pinned to a commit, so "
+            f"the TPC-H dataset can change without the cache key rotating"
+        )
+    for failure in failures:
+        print(f"tpc dataset cache: {failure}")
+    return not failures
+
+
 # Fork runs. A push to a contributor's fork runs the PR tier there (ci_fork.yml),
 # upstream skips that tier for a fork pull request (compute-changes.py), and
 # fork_ci_status.yml mirrors the fork's run onto the pull request. The pieces
@@ -1886,6 +2003,7 @@ def check_scheduled_workflows_stay_upstream():
     ok = check_nightly_scope() and ok
     ok = check_nightly_base_fallback() and ok
     ok = check_cache_save_scope() and ok
+    ok = check_tpc_dataset_caches() and ok
     ok = check_local_ci_config() and ok
     if not ok:
         sys.exit(1)

spark/src/test/scala/org/apache/spark/sql/GenTPCHData.scala

@@ -119,14 +119,18 @@ object GenTPCHData {
     }
   }
 
-  // Install tpch-dbgen (with the stdout patch)
+  // Install tpch-dbgen, pinned to a commit so the TPC-H dataset cache, which is
+  // keyed on this file, rotates only when the generator changes.
   def installDBGEN(
       baseDir: String,
       url: String = "https://github.com/databricks/tpch-dbgen.git",
       useStdout: Boolean = true)(i: java.lang.Long): Unit = {
-    // Check if we want the revision which makes dbgen output to stdout
+    val dbgenRevision = "6985da461c641fd0d255b214f2d693f1bf08bc33"
+    // The revision whose bm_utils.c makes dbgen output to stdout.
+    val stdoutPatchRevision = "0469309147b42abac8857fa61b4cf69a6d3128a8"
     val checkoutRevision: String =
-      if (useStdout) "git checkout 0469309147b42abac8857fa61b4cf69a6d3128a8 -- bm_utils.c" else ""
+      s"git checkout $dbgenRevision" +
+        (if (useStdout) s" && git checkout $stdoutPatchRevision -- bm_utils.c" else "")
 
     Seq("mkdir", "-p", baseDir).!
     val pw = new PrintWriter(s"${baseDir}/dbgen_$i.sh")