ci: shard the Iceberg extensions test task

.github/workflows/iceberg_spark_test_reusable.yml

@@ -206,8 +206,11 @@ jobs:
 
   iceberg-spark-extensions:
     needs: build-native
-    name: iceberg-spark-extensions/iceberg-${{ inputs.iceberg-full }}/spark-${{ inputs.spark-full }}/scala-${{ inputs.scala }}/java-${{ inputs.java }}
+    name: iceberg-spark-extensions/iceberg-${{ inputs.iceberg-full }}/spark-${{ inputs.spark-full }}/scala-${{ inputs.scala }}/java-${{ inputs.java }}/shard-${{ matrix.shard }}
     runs-on: ubuntu-24.04
+    strategy:
+      fail-fast: false
+      matrix: ${{ fromJSON(needs.build-native.outputs.shard-matrix) }}
     container:
       image: amd64/rust
     env:
@@ -238,21 +241,52 @@ jobs:
           ENABLE_COMET=true ENABLE_COMET_ONHEAP=true COMET_ICEBERG_DATA_FILE_CONCURRENCY_LIMIT=1 COMET_ICEBERG_WRITE_REPORT_DIR="$PWD/build/comet-iceberg-writes" \
             ./gradlew -DsparkVersions=${{ inputs.spark-short }} -DscalaVersion=${{ inputs.scala }} -DflinkVersions= -DkafkaVersions= \
             :iceberg-spark:iceberg-spark-extensions-${{ inputs.spark-short }}_${{ inputs.scala }}:test \
+            --init-script ../dev/ci/iceberg-test-shards.gradle \
+            -PcometShardTask=:iceberg-spark:iceberg-spark-extensions-${{ inputs.spark-short }}_${{ inputs.scala }}:test \
+            -PcometShardIndex=${{ matrix.shard }} -PcometShardCount=${{ needs.build-native.outputs.shard-count }} \
             -Pquick=true -x javadoc
       - name: Summarize Iceberg writes
         if: ${{ !cancelled() }}
         run: |
-          python3 dev/ci/summarize-iceberg-writes.py --title "iceberg-spark-extensions" \
+          python3 dev/ci/summarize-iceberg-writes.py --title "iceberg-spark-extensions shard ${{ matrix.shard }}" \
             apache-iceberg/build/comet-iceberg-writes
-      - name: Upload Iceberg write report
+      - name: Upload Iceberg shard inventory and test reports
         if: ${{ !cancelled() }}
+        # iceberg-spark-extensions-shard-coverage downloads the inventory, so a
+        # flaky upload here fails that job too. Retry rather than fail on 403.
         uses: ./.github/actions/upload-artifact-retry
         with:
-          name: iceberg-spark-extensions-writes-${{ inputs.iceberg-full }}-spark-${{ inputs.spark-full }}-scala-${{ inputs.scala }}-jdk${{ inputs.java }}-attempt-${{ github.run_attempt }}
-          path: apache-iceberg/build/comet-iceberg-writes/*.jsonl
-          if-no-files-found: ignore
+          name: iceberg-spark-extensions-${{ inputs.iceberg-full }}-spark-${{ inputs.spark-full }}-scala-${{ inputs.scala }}-jdk${{ inputs.java }}-shard-${{ matrix.shard }}-attempt-${{ github.run_attempt }}
+          path: |
+            apache-iceberg/**/build/comet-shards/*.json
+            apache-iceberg/**/build/test-results/test/*.xml
+            apache-iceberg/build/comet-iceberg-writes/*.jsonl
           retention-days: 7
 
+  iceberg-spark-extensions-shard-coverage:
+    needs: iceberg-spark-extensions
+    if: ${{ !cancelled() }}
+    name: iceberg-spark-extensions-shard-coverage/iceberg-${{ inputs.iceberg-full }}/spark-${{ inputs.spark-full }}/scala-${{ inputs.scala }}/java-${{ inputs.java }}
+    runs-on: ubuntu-slim
+    steps:
+      - uses: actions/checkout@v7
+      - name: Download Iceberg extensions shard inventories
+        uses: ./.github/actions/download-artifact-retry
+        with:
+          # Download all attempts: a failed-job rerun retains earlier successful
+          # shards. The checker selects the latest inventory for each index.
+          pattern: iceberg-spark-extensions-${{ inputs.iceberg-full }}-spark-${{ inputs.spark-full }}-scala-${{ inputs.scala }}-jdk${{ inputs.java }}-shard-*-attempt-*
+          path: iceberg-shard-reports
+      - name: Verify complete, disjoint Iceberg extensions candidate coverage
+        run: |
+          python3 dev/ci/check-iceberg-shards.py --manifests iceberg-shard-reports \
+            --task :iceberg-spark:iceberg-spark-extensions-${{ inputs.spark-short }}_${{ inputs.scala }}:test
+      - name: Summarize Iceberg writes across shards
+        if: ${{ !cancelled() }}
+        run: |
+          python3 dev/ci/summarize-iceberg-writes.py --title "iceberg-spark-extensions, all shards" \
+            iceberg-shard-reports
+
   iceberg-spark-runtime:
     needs: build-native
     name: iceberg-spark-runtime/iceberg-${{ inputs.iceberg-full }}/spark-${{ inputs.spark-full }}/scala-${{ inputs.scala }}/java-${{ inputs.java }}

dev/ci/check-ci-config.py

@@ -1595,6 +1595,39 @@ def check_cache_save_scope():
     return not failures
 
 
+def check_iceberg_extensions_shards():
+    """The Iceberg extensions task must be sharded like the core task.
+
+    Its job is the longest unsharded one in the Iceberg workflow. It has to use
+    the shared matrix and the same Gradle init script, and a coverage job must
+    verify the extensions task's candidate inventories, not the core task's.
+    """
+    path = WORKFLOWS / "iceberg_spark_test_reusable.yml"
+    text = path.read_text(encoding="utf-8")
+    failures = []
+
+    job = re.search(r"^  iceberg-spark-extensions:\n(.*?)(?=^  \S)", text, re.M | re.S)
+    block = job.group(1) if job else ""
+    if "fromJSON(needs.build-native.outputs.shard-matrix)" not in block:
+        failures.append("iceberg-spark-extensions does not use the shared shard matrix")
+    if "cometShardTask=:iceberg-spark:iceberg-spark-extensions-" not in block:
+        failures.append("iceberg-spark-extensions does not pass -PcometShardTask for its own task")
+    if "cometShardIndex=" not in block or "cometShardCount=" not in block:
+        failures.append("iceberg-spark-extensions does not pass the shard index/count")
+
+    coverage = re.search(
+        r"^  iceberg-spark-extensions-shard-coverage:\n(.*?)(?=^  \S)", text, re.M | re.S)
+    cblock = coverage.group(1) if coverage else ""
+    if (not coverage or "check-iceberg-shards.py --manifests" not in cblock
+            or "--task :iceberg-spark:iceberg-spark-extensions-" not in cblock):
+        failures.append(
+            "no iceberg-spark-extensions-shard-coverage job verifying the extensions task")
+
+    for failure in failures:
+        print(f"iceberg extensions shards: {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
@@ -1887,6 +1920,7 @@ def check_scheduled_workflows_stay_upstream():
     ok = check_nightly_base_fallback() and ok
     ok = check_cache_save_scope() and ok
     ok = check_local_ci_config() and ok
+    ok = check_iceberg_extensions_shards() and ok
     if not ok:
         sys.exit(1)
     print("CI config checks passed")

dev/ci/check-iceberg-shards.py

@@ -140,6 +140,15 @@ def candidates = new TreeSet<String>()
             file('build/candidates.json').text = JsonOutput.toJson(candidates)
           }
         }
+        tasks.register('iceberg-spark-extensions-fixture', Test) {
+          testClassesDirs = sourceSets.test.output.classesDirs
+          classpath = sourceSets.test.runtimeClasspath
+          useJUnitPlatform { excludeTags 'excluded-tag' }
+          include project.findProperty('extensionFixtureIncludes') ?: '**/Test*.class'
+          exclude '**/TestStructuredStreamingRead*.class', '**/TestExcludedByPattern.class',
+              '**/TestOtherTask.class'
+          exclude { it.name == 'TestExcludedBySpec.class' }
+        }
         tasks.register('otherTest', Test) {
           testClassesDirs = sourceSets.test.output.classesDirs
           classpath = sourceSets.test.runtimeClasspath
@@ -158,6 +167,7 @@ def candidates = new TreeSet<String>()
             @org.junit.jupiter.api.Test void nested() {}
           }
         """,
+        "TestExtensionD": "@org.junit.jupiter.api.Test void extensionShardFour() {}",
         "TestAlpha": "@org.junit.jupiter.api.Test void alpha() {}",
         "TestBeta": "@org.junit.jupiter.api.Test void beta() {}",
         "TestDelta": "@org.junit.jupiter.api.Test void delta() {}",
@@ -239,7 +249,7 @@ def shard_args(index, count=SHARD_COUNT, task=":test"):
 
     baseline = run("baseline")
     inventory = set(json.loads((root / "build/candidates.json").read_text()))
-    assert sum(baseline.values()) == 12, baseline
+    assert sum(baseline.values()) == 13, baseline
     assert all(state == "passed" for _, _, state in baseline), baseline
     assert any(cls == "fixture.TestNewlyAdded" for cls, _, _ in baseline)
     assert any("$Nested" in cls for cls, _, _ in baseline)
@@ -255,6 +265,8 @@ def shard_args(index, count=SHARD_COUNT, task=":test"):
         assert candidates == manifest["candidates"]
         assert set(manifest["unshardedCandidates"]) == inventory
         assert cases, f"empty fixture shard {index}"
+        assert all(("TestStructuredStreamingRead" in name) == (index == 1)
+                   for name in candidates), (index, candidates)
         combined_cases.update(cases)
         combined_candidates.update(candidates)
         if any(cls == "fixture.TestFailurePropagation" for cls, _, _ in cases):
@@ -265,6 +277,46 @@ def shard_args(index, count=SHARD_COUNT, task=":test"):
     assert combined_cases == baseline, (combined_cases, baseline)
     verify_manifests(results, ":test")
     assert run("single-shard", shard_args(1, count=1)) == baseline
+    # Extensions have no streaming family. Exercise the actual shared init
+    # script against that candidate set, rather than a Python allocation model.
+    extensions_task = "iceberg-spark-extensions-fixture"
+    extensions_baseline = run("extensions-baseline", task=extensions_task)
+    assert sum(extensions_baseline.values()) == 9, extensions_baseline
+    extensions_cases = Counter()
+    extensions_candidates = Counter()
+    extensions_results = results / "extensions"
+    extensions_results.mkdir()
+    sizes = []
+    for index in range(1, SHARD_COUNT + 1):
+        cases = run(f"extensions-shard-{index}",
+                    shard_args(index, task=f":{extensions_task}"), task=extensions_task)
+        assert cases, f"empty extensions fixture shard {index}"
+        manifest = json.loads((root / f"build/comet-shards/{extensions_task}-{index}.json").read_text())
+        assert all("StructuredStreamingRead" not in name for name in manifest["unshardedCandidates"])
+        extensions_cases.update(cases)
+        extensions_candidates.update(manifest["candidates"])
+        sizes.append(len(manifest["candidates"]))
+        (extensions_results / f"shard-{index}.json").write_text(json.dumps(manifest, indent=2) + "\n")
+    assert extensions_cases == extensions_baseline, (extensions_cases, extensions_baseline)
+    assert set(extensions_candidates.values()) == {1}, extensions_candidates
+    verify_manifests(extensions_results, f":{extensions_task}")
+    assert run("extensions-single-shard", shard_args(1, count=1, task=f":{extensions_task}"),
+               task=extensions_task) == extensions_baseline
+    print(f"Extensions fixture candidate counts: {sizes}", flush=True)
+    # Hash partitions may legitimately be empty when fewer classes than shards
+    # exist. They must still execute the small inventory exactly once in total.
+    small = ["-PextensionFixtureIncludes=**/TestAlpha.class"]
+    small_baseline = run("extensions-small-baseline", small, task=extensions_task)
+    assert sum(small_baseline.values()) == 1, small_baseline
+    small_combined = Counter()
+    for index in range(1, SHARD_COUNT + 1):
+        small_combined.update(run(f"extensions-small-shard-{index}",
+                                 shard_args(index, task=f":{extensions_task}") + small,
+                                 task=extensions_task))
+    assert small_combined == small_baseline, (small_combined, small_baseline)
+    empty = ["-PextensionFixtureIncludes=**/NoTests.class"]
+    assert not run("extensions-empty", shard_args(1, task=f":{extensions_task}") + empty,
+                   task=extensions_task)
 
     other_baseline = run("other-baseline", task="otherTest")
     assert sum(other_baseline.values()) == 1

dev/ci/iceberg-test-shards.gradle

@@ -25,13 +25,17 @@ import org.gradle.api.tasks.testing.Test
 // its includes, replacing its exclusions, or changing JUnit's discovery/parameters.
 // All classes remain on the test classpath, including inherited test fixtures.
 class CometIcebergTestShard {
-  static int owner(String classFileName, int count) {
+  static int owner(String classFileName, int count, boolean extensions) {
     if (count == 1) {
       return 1
     }
     // Keep JUnit @Nested classes with their enclosing class.
     String outer = classFileName.substring(0, classFileName.length() - '.class'.length())
         .split('\\$', 2)[0]
+    // Extensions have no streaming family, so use every shard.
+    if (extensions) {
+      return 1 + Math.floorMod(outer.hashCode(), count)
+    }
     // This family exists in all supported Iceberg versions. In the 1.11 run used
     // to size these shards it took ~19 of ~67 test minutes. The remaining classes
     // hash to ~18.5 / 16.4 / 13.2 minutes; new classes need no allowlist update.
@@ -69,10 +73,11 @@ gradle.projectsEvaluated {
   int runAttempt = Integer.parseInt(System.getenv('GITHUB_RUN_ATTEMPT') ?: '1')
   testTask.inputs.property('cometShardRunAttempt', runAttempt)
   testTask.inputs.files(gradle.startParameter.initScripts).withPropertyName('cometShardScripts')
+  boolean extensions = properties.cometShardTask.contains(':iceberg-spark-extensions-')
   boolean applyShardFilter = true
   testTask.exclude { FileTreeElement entry ->
     applyShardFilter && !entry.directory && entry.name.endsWith('.class') &&
-        CometIcebergTestShard.owner(entry.name, shardCount) != shardIndex
+        CometIcebergTestShard.owner(entry.name, shardCount, extensions) != shardIndex
   }
 
   // Keep a machine-readable inventory next to the normal JUnit reports. This is

docs/source/contributor-guide/ci.md

@@ -442,7 +442,8 @@ The release process does this before tagging each release candidate; see
 ## Reproducing a suite failure locally
 
 `dev/local-ci.sh` builds the same sandbox a runner builds and runs the Spark SQL or Iceberg
-workflow:
+targets. Iceberg extensions are an exception to job parity: the local runner runs them unsharded,
+while CI uses four extensions shards:
 
 ```sh
 dev/local-ci.sh spark                # everything the Spark job runs
@@ -455,7 +456,9 @@ dev/local-ci.sh --print-config       # what it read from ci.yml and dev/ci
 
 The version defaults to the one the merge queue gates on. That, the matrix rows and the shard count
 are read from the workflow files and from `dev/ci/`, so a local shard runs what the CI shard of the
-same name runs. It prepares first (native build, Comet install, patched Spark or Iceberg source
+same name runs for Spark SQL and Iceberg core tests. To reproduce an extensions shard, use the
+[direct Gradle recipe](iceberg-spark-tests.md#reproducing-an-extensions-shard). The local runner
+prepares first (native build, Comet install, patched Spark or Iceberg source
 under `$COMET_LOCAL_CI_HOME`, default `/tmp/comet-local-ci`), then runs the tests, compiling only
 the sbt projects the selected rows need.