diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8cf28c72dd2..f20eef2ad8a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -191,6 +191,9 @@ jobs: - name: Check Iceberg shard inventory validation run: python3 dev/ci/test-iceberg-shards.py + - name: Check Iceberg write report summary + run: python3 dev/ci/test-summarize-iceberg-writes.py + - name: Check CI config invariants run: python3 dev/ci/check-ci-config.py diff --git a/.github/workflows/iceberg_spark_test_reusable.yml b/.github/workflows/iceberg_spark_test_reusable.yml index ae20b6d7683..9b4a5ae3000 100644 --- a/.github/workflows/iceberg_spark_test_reusable.yml +++ b/.github/workflows/iceberg_spark_test_reusable.yml @@ -154,12 +154,20 @@ jobs: run: | cd apache-iceberg rm -rf /root/.m2/repository/org/apache/parquet # somehow parquet cache requires cleanups - ENABLE_COMET=true ENABLE_COMET_ONHEAP=true ./gradlew -DsparkVersions=${{ inputs.spark-short }} -DscalaVersion=${{ inputs.scala }} -DflinkVersions= -DkafkaVersions= \ + # COMET_ICEBERG_WRITE_REPORT_DIR records which writer ran each Iceberg + # write; see dev/ci/summarize-iceberg-writes.py. + ENABLE_COMET=true ENABLE_COMET_ONHEAP=true COMET_ICEBERG_WRITE_REPORT_DIR="$PWD/build/comet-iceberg-writes" \ + ./gradlew -DsparkVersions=${{ inputs.spark-short }} -DscalaVersion=${{ inputs.scala }} -DflinkVersions= -DkafkaVersions= \ :iceberg-spark:iceberg-spark-${{ inputs.spark-short }}_${{ inputs.scala }}:test \ --init-script ../dev/ci/iceberg-test-shards.gradle \ -PcometShardTask=:iceberg-spark:iceberg-spark-${{ 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 shard ${{ matrix.shard }}" \ + apache-iceberg/build/comet-iceberg-writes - name: Upload Iceberg shard inventory and test reports if: ${{ !cancelled() }} # iceberg-spark-shard-coverage downloads the inventory, so a flaky @@ -170,6 +178,7 @@ jobs: 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-shard-coverage: @@ -190,6 +199,11 @@ jobs: run: | python3 dev/ci/check-iceberg-shards.py --manifests iceberg-shard-reports \ --task :iceberg-spark:iceberg-spark-${{ 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, all shards" \ + iceberg-shard-reports iceberg-spark-extensions: needs: build-native @@ -222,9 +236,23 @@ jobs: run: | cd apache-iceberg rm -rf /root/.m2/repository/org/apache/parquet # somehow parquet cache requires cleanups - ENABLE_COMET=true ENABLE_COMET_ONHEAP=true ./gradlew -DsparkVersions=${{ inputs.spark-short }} -DscalaVersion=${{ inputs.scala }} -DflinkVersions= -DkafkaVersions= \ + ENABLE_COMET=true ENABLE_COMET_ONHEAP=true 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 \ -Pquick=true -x javadoc + - name: Summarize Iceberg writes + if: ${{ !cancelled() }} + run: | + python3 dev/ci/summarize-iceberg-writes.py --title "iceberg-spark-extensions" \ + apache-iceberg/build/comet-iceberg-writes + - name: Upload Iceberg write report + if: ${{ !cancelled() }} + 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 + retention-days: 7 iceberg-spark-runtime: needs: build-native diff --git a/dev/ci/summarize-iceberg-writes.py b/dev/ci/summarize-iceberg-writes.py new file mode 100644 index 00000000000..c94638b4e12 --- /dev/null +++ b/dev/ci/summarize-iceberg-writes.py @@ -0,0 +1,171 @@ +# 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. + +"""Summarize which writer ran the Iceberg writes of an Iceberg Spark test run. + +The Iceberg Spark test jobs set COMET_ICEBERG_WRITE_REPORT_DIR, so Comet's +IcebergWriteReportListener appends one JSON line per Iceberg write to a file +in that directory. This prints how many writes ran on Comet's native writer, +how many Comet's split operator left on Iceberg's JVM writer and why, and how +many Spark planned without Comet's split operator: + + python3 dev/ci/summarize-iceberg-writes.py --title "iceberg-spark shard 1" DIR... + +Files under a directory named like a shard artifact (...-shard-N-attempt-M) +are counted only for the latest attempt of each shard, so a rerun of failed +jobs does not count a shard twice. The latest attempt is the newest such +directory, whether or not it holds any report files, so a rerun that recorded +no writes is reported as missing rather than replaced by an earlier attempt. +The summary is also appended to $GITHUB_STEP_SUMMARY when that is set. It +never fails the job. +""" + +import argparse +from collections import Counter +import json +import os +from pathlib import Path +import re + + +SHARD_ATTEMPT = re.compile(r"-shard-(\d+)-attempt-(\d+)$") +WRITERS = [ + ("native", "Comet native writer"), + ("jvm", "Iceberg JVM writer under Comet's split operator"), + ("spark", "Spark V2 write, not planned by Comet's split operator"), +] +TOP_REASONS = 20 + + +def shard_attempt(path): + """The (shard, attempt) of the shard artifact directory holding path, or None.""" + for part in path.parts: + match = SHARD_ATTEMPT.search(part) + if match: + return int(match.group(1)), int(match.group(2)) + return None + + +def report_files(roots): + """Every report file under roots, keeping only the latest attempt of each shard. + + Returns the files and the (shard, attempt) pairs whose latest attempt holds no report file. + The latest attempt comes from the artifact directories rather than the report files, because + an attempt that recorded no writes still uploads its shard inventory and test reports. + """ + latest = {} + files = [] + for root in map(Path, roots): + for path in [root, *sorted(root.rglob("*"))]: + key = shard_attempt(path) + if key: + latest[key[0]] = max(latest.get(key[0], 0), key[1]) + if path.suffix == ".jsonl" and path.is_file(): + files.append((path, key)) + kept = [(path, key) for path, key in files if key is None or latest[key[0]] == key[1]] + reported = {key for _, key in kept} + missing = sorted(key for key in latest.items() if key not in reported) + return [path for path, _ in kept], missing + + +def load(roots): + files, missing = report_files(roots) + writes = [] + for path in files: + for line in path.read_text(encoding="utf-8").splitlines(): + if line.strip(): + writes.append(json.loads(line)) + return writes, missing + + +def cell(text): + return " ".join(text.split()).replace("|", "\\|") + + +def summarize(title, writes, missing=()): + lines = [f"### Iceberg writes: {title}", ""] + for shard, attempt in missing: + lines += [ + f"Shard {shard} recorded no Iceberg writes in its latest attempt ({attempt}), " + "so none of its writes are counted below.", + "", + ] + if not writes: + lines.append( + "No Iceberg writes were recorded. Either the target ran none or " + "COMET_ICEBERG_WRITE_REPORT_DIR did not reach the test JVMs." + ) + return "\n".join(lines) + "\n" + + total = len(writes) + by_writer = Counter(w["writer"] for w in writes) + lines += ["| Writer | Writes | Share |", "| --- | ---: | ---: |"] + for key, label in WRITERS: + count = by_writer.get(key, 0) + lines.append(f"| {label} | {count} | {100.0 * count / total:.1f}% |") + lines += [f"| Total | {total} | |", ""] + failed = sum(1 for w in writes if w.get("failed")) + if failed: + lines += [f"{failed} of the {total} writes ran in queries that failed.", ""] + + reasons = Counter() + for w in writes: + if w["writer"] == "jvm": + for reason in w.get("reasons") or ["(no reason recorded)"]: + reasons[reason] += 1 + if reasons: + lines += [ + "#### Why the split operator kept the JVM writer", + "", + "A write can have several reasons, so the counts can add up to more than the " + "JVM writes.", + "", + "| Writes | Reason |", + "| ---: | --- |", + ] + for reason, count in reasons.most_common(TOP_REASONS): + lines.append(f"| {count} | {cell(reason)} |") + if len(reasons) > TOP_REASONS: + lines.append(f"| | and {len(reasons) - TOP_REASONS} more reasons |") + lines.append("") + + operators = Counter(w["node"] for w in writes if w["writer"] == "spark") + if operators: + lines += ["#### Spark V2 writes by operator", "", "| Writes | Operator |", "| ---: | --- |"] + for node, count in operators.most_common(): + lines.append(f"| {count} | {cell(node)} |") + lines.append("") + + return "\n".join(lines) + "\n" + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--title", required=True, help="heading for the summary") + parser.add_argument("roots", nargs="+", help="directories holding the report files") + args = parser.parse_args() + + summary = summarize(args.title, *load(args.roots)) + print(summary) + step_summary = os.environ.get("GITHUB_STEP_SUMMARY") + if step_summary: + with open(step_summary, "a", encoding="utf-8") as out: + out.write(summary + "\n") + + +if __name__ == "__main__": + main() diff --git a/dev/ci/test-summarize-iceberg-writes.py b/dev/ci/test-summarize-iceberg-writes.py new file mode 100644 index 00000000000..cce4ada6630 --- /dev/null +++ b/dev/ci/test-summarize-iceberg-writes.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 +# +# 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. + +"""Fast regression tests for the Iceberg write report summary.""" + +import importlib.util +import json +from pathlib import Path +import tempfile +import unittest + + +SPEC = importlib.util.spec_from_file_location( + "summarize_iceberg_writes", Path(__file__).with_name("summarize-iceberg-writes.py")) +SUMMARIZE = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(SUMMARIZE) + +ARTIFACT = "iceberg-spark-1.11.0-spark-4.1.3-scala-2.13-jdk17-shard-{}-attempt-{}" + + +def records(*writers): + return "".join( + json.dumps({"writer": w, "node": "AppendData", "reasons": [], "failed": False}) + "\n" + for w in writers) + + +class SummarizeIcebergWritesTest(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory(prefix="comet-iceberg-writes-test-") + self.addCleanup(self.temp.cleanup) + self.root = Path(self.temp.name) + + def attempt(self, shard, attempt, *writers): + """A shard attempt's artifact as the coverage job downloads it. + + Every attempt uploads its test reports. Only an attempt that recorded writes has a + report file. + """ + artifact = self.root / ARTIFACT.format(shard, attempt) + reports = artifact / "build/test-results/test" + reports.mkdir(parents=True) + (reports / "TEST-org.example.TestFixture.xml").write_text("") + if writers: + writes = artifact / "build/comet-iceberg-writes" + writes.mkdir(parents=True) + (writes / "iceberg-writes-fixture.jsonl").write_text(records(*writers)) + + def load(self): + writes, missing = SUMMARIZE.load([self.root]) + return sorted(w["writer"] for w in writes), missing + + def test_every_shard_counts_once(self): + self.attempt(1, 1, "native") + self.attempt(2, 1, "jvm", "spark") + self.assertEqual(self.load(), (["jvm", "native", "spark"], [])) + + def test_latest_attempt_replaces_an_earlier_one(self): + self.attempt(1, 1, "native", "native") + self.attempt(1, 2, "jvm") + self.assertEqual(self.load(), (["jvm"], [])) + + def test_retry_without_writes_is_reported_missing_instead_of_stale(self): + self.attempt(1, 1, "native") + self.attempt(1, 2) + self.attempt(2, 1, "spark") + self.assertEqual(self.load(), (["spark"], [(1, 2)])) + summary = SUMMARIZE.summarize("fixture", *SUMMARIZE.load([self.root])) + self.assertIn("Shard 1 recorded no Iceberg writes in its latest attempt (2)", summary) + self.assertIn("| Comet native writer | 0 | 0.0% |", summary) + self.assertIn("| Spark V2 write, not planned by Comet's split operator | 1 | 100.0% |", + summary) + + def test_root_that_is_itself_a_shard_artifact(self): + self.attempt(1, 3) + writes, missing = SUMMARIZE.load([self.root / ARTIFACT.format(1, 3)]) + self.assertEqual((writes, missing), ([], [(1, 3)])) + + def test_files_outside_shard_artifacts_always_count(self): + # A shard job and dev/local-ci.sh summarize their own report directory directly. + (self.root / "iceberg-writes-local.jsonl").write_text(records("native", "jvm")) + self.assertEqual(self.load(), (["jvm", "native"], [])) + + def test_no_writes_at_all(self): + summary = SUMMARIZE.summarize("fixture", *SUMMARIZE.load([self.root / "missing"])) + self.assertIn("No Iceberg writes were recorded", summary) + + +if __name__ == "__main__": + unittest.main() diff --git a/dev/local-ci.sh b/dev/local-ci.sh index 8e3267302d9..a6a9d57d83d 100755 --- a/dev/local-ci.sh +++ b/dev/local-ci.sh @@ -389,16 +389,25 @@ run_iceberg() { gradlew ":iceberg-spark:iceberg-spark-runtime-${SPARK}_${SCALA}:integrationTest" ;; esac + case "$target" in + shard-* | extensions) + python3 "$REPO/dev/ci/summarize-iceberg-writes.py" --title "$target" \ + "$dest/build/comet-iceberg-writes/$target" + ;; + esac ok "$target took $(hms $((SECONDS - started)))" done } -# Reads $dest, $spark and $SCALA from run_iceberg. +# Reads $dest, $spark, $SCALA and $target from run_iceberg. gradlew() { ( cd "$dest" # shellcheck disable=SC2031 export SPARK_LOCAL_IP=localhost ENABLE_COMET=true ENABLE_COMET_ONHEAP=true + # One directory per target, emptied first, so a rerun reports only its own writes. + export COMET_ICEBERG_WRITE_REPORT_DIR="$dest/build/comet-iceberg-writes/$target" + rm -rf "$COMET_ICEBERG_WRITE_REPORT_DIR" ./gradlew "-DsparkVersions=$SPARK" "-DscalaVersion=$SCALA" \ -DflinkVersions= -DkafkaVersions= "$@" -Pquick=true -x javadoc ) diff --git a/docs/source/contributor-guide/iceberg-spark-tests.md b/docs/source/contributor-guide/iceberg-spark-tests.md index fb0e262e341..0db54feeffe 100644 --- a/docs/source/contributor-guide/iceberg-spark-tests.md +++ b/docs/source/contributor-guide/iceberg-spark-tests.md @@ -43,7 +43,9 @@ Here is an overview of the changes that the diffs make to Iceberg: `LocalTableScanExec`, the conversion is declined, and the write silently runs on the JVM writer. Many Iceberg suites seed their data that way, so leaving it off hides the native writer from most of the write surface. - Enable fallback logging (`spark.comet.explainFallback.enabled`) so that every operator Comet declines is - reported in the test output together with the reason it was declined. + reported in the test output together with the reason it was declined. The output goes to the JUnit XML + reports rather than the CI job log; see [Which writer ran each Iceberg write](#which-writer-ran-each-iceberg-write) + for how CI reports native write coverage. [#3739]: https://github.com/apache/datafusion-comet/pull/3739 [#5259]: https://github.com/apache/datafusion-comet/issues/5259 @@ -154,3 +156,28 @@ path, reflection code (`org.apache.comet.iceberg.IcebergReflection`), or other l can differ across Iceberg versions. The Comet test suites in the Linux build do not exercise Iceberg's own Spark tests, so without the label the first Iceberg 1.11 verdict is the merge queue's, and the first verdict on the older Iceberg versions is the nightly run's, after the change has landed. + +### Which writer ran each Iceberg write + +A passing Iceberg job does not show that Comet's native writer ran. `CometIcebergNativeWrite` falls back +to Iceberg's JVM writer without failing the write, no upstream test asserts which writer ran, and Gradle +does not copy the fallback warnings into the job log. So the core and extensions jobs set +`COMET_ICEBERG_WRITE_REPORT_DIR`, the environment variable behind the test-only config +`spark.comet.testing.icebergWriteReport.dir`. When it is set, the Comet driver plugin registers +`IcebergWriteReportListener`, which writes one JSON line for each Iceberg write the tests run. Each line +records one of three writers: + +- `native`: Comet's native writer (`CometIcebergWriteExec`). +- `jvm`: Comet's split operator planned the write but kept Iceberg's JVM writer (`IcebergWriteExec`). + The line includes the reasons Comet recorded for not converting it. +- `spark`: Spark's own V2 write operator ran the write, so Comet's split operator never saw it. Examples + are `WriteDelta` for merge-on-read, `WriteToDataSourceV2` for a streaming micro-batch, and on Spark + 3.4 the CTAS and RTAS execs, which write the table themselves. + +`dev/ci/summarize-iceberg-writes.py` turns these records into a table on the job's summary page. It +shows the count and share of each writer, the most common fallback reasons, and the Spark write +operators. Each shard and the extensions job gets its own table. The shard coverage job adds one for +all shards together, counting only the latest attempt of each shard. A shard whose latest attempt +recorded no writes is named above the table rather than counted from an earlier attempt. The raw +records are uploaded with the job's other reports. The summary never fails a job. +`dev/local-ci.sh iceberg` prints the same summary after each shard and after the extensions target. diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 57718ec9ae9..53daeb5710f 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -951,6 +951,19 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(true) + val COMET_ICEBERG_WRITE_REPORT_DIR: ConfigEntry[String] = + conf("spark.comet.testing.icebergWriteReport.dir") + .internal() + .category(CATEGORY_TESTING) + .doc("Test-only. When set, the Comet driver plugin registers a query listener that " + + "records every Iceberg write the application runs, the writer that ran it (Comet's " + + "native writer, the JVM writer behind Comet's split operator, or Spark's own V2 write) " + + "and the reasons Comet recorded for not writing natively, as JSON lines in this " + + "directory. `dev/ci/summarize-iceberg-writes.py` summarizes them. The Iceberg Spark " + + "test jobs set it through the environment variable so the Iceberg diffs need no change.") + .stringConf + .createWithEnvVarOrDefault("COMET_ICEBERG_WRITE_REPORT_DIR", "") + val COMET_SPARK_TO_ARROW_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.sparkToColumnar.enabled") .category(CATEGORY_EXEC) diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteReportListener.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteReportListener.scala new file mode 100644 index 00000000000..d5f3663eb1b --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteReportListener.scala @@ -0,0 +1,135 @@ +/* + * 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. + */ + +package org.apache.comet.iceberg + +import java.io.File +import java.nio.charset.StandardCharsets.UTF_8 +import java.nio.file.{Files, StandardOpenOption} +import java.util.UUID + +import scala.util.control.NonFatal + +import org.json4s.JsonDSL._ +import org.json4s.jackson.JsonMethods._ + +import org.apache.spark.SparkConf +import org.apache.spark.internal.Logging +import org.apache.spark.sql.comet.{CometIcebergWriteExec, IcebergWriteExec} +import org.apache.spark.sql.connector.write.BatchWrite +import org.apache.spark.sql.execution.{CommandResultExec, QueryExecution, SparkPlan} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, QueryStageExec} +import org.apache.spark.sql.execution.datasources.v2.{V2ExistingTableWriteExec, WriteToDataSourceV2Exec} +import org.apache.spark.sql.execution.streaming.sources.MicroBatchWrite +import org.apache.spark.sql.util.QueryExecutionListener + +import org.apache.comet.CometConf.COMET_ICEBERG_WRITE_REPORT_DIR +import org.apache.comet.CometExplainInfo + +/** + * Test-only listener that records which writer ran each Iceberg write, so a CI job running + * Iceberg's own Spark suites can tell a native write from a silent fallback. The Comet driver + * plugin registers it when `spark.comet.testing.icebergWriteReport.dir` is set. Each write is + * appended as one JSON line to a file of its own in that directory, which + * `dev/ci/summarize-iceberg-writes.py` reads. + */ +class IcebergWriteReportListener(conf: SparkConf) extends QueryExecutionListener with Logging { + + private val reportFile: File = { + val dir = new File( + conf + .get(COMET_ICEBERG_WRITE_REPORT_DIR.key, COMET_ICEBERG_WRITE_REPORT_DIR.defaultValue.get)) + dir.mkdirs() + // One file per listener: Gradle runs several test JVMs, each possibly with several sessions. + new File(dir, s"iceberg-writes-${UUID.randomUUID()}.jsonl") + } + + override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = + record(qe, failed = false) + + override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = + record(qe, failed = true) + + private def record(qe: QueryExecution, failed: Boolean): Unit = { + try { + val lines = IcebergWriteReportListener.writes(qe.executedPlan).map { w => + compact( + render(("writer" -> w.writer) ~ ("node" -> w.node) ~ ("reasons" -> w.reasons.toList) ~ + ("failed" -> failed))) + "\n" + } + if (lines.nonEmpty) { + synchronized { + Files.write( + reportFile.toPath, + lines.mkString.getBytes(UTF_8), + StandardOpenOption.CREATE, + StandardOpenOption.APPEND) + } + } + } catch { + // A query that failed during planning has no executed plan; nothing was written. + case NonFatal(e) => logWarning(s"Could not record Iceberg writes for a query: $e") + } + } +} + +object IcebergWriteReportListener { + + /** Comet's native (iceberg-rust) writer ran the write. */ + val Native = "native" + + /** Comet's split operator planned the write, but Iceberg's JVM writer ran it. */ + val Jvm = "jvm" + + /** Spark's own V2 write operator ran the write; Comet's split operator did not plan it. */ + val Spark = "spark" + + case class IcebergWrite(writer: String, node: String, reasons: Seq[String]) + + /** The Iceberg writes in an executed plan, with the reasons Comet did not write natively. */ + def writes(plan: SparkPlan): Seq[IcebergWrite] = plan match { + // A command's writes are reported by the command's own execution, which runs eagerly before + // the query wrapping its result. Descending here would count them twice. + case _: CommandResultExec => Nil + case a: AdaptiveSparkPlanExec => writes(a.executedPlan) + case s: QueryStageExec => writes(s.plan) + case w: CometIcebergWriteExec => Seq(IcebergWrite(Native, w.nodeName, Nil)) + case w: IcebergWriteExec => + val reasons = w.getTagValue(CometExplainInfo.FALLBACK_REASONS).getOrElse(Set.empty) + Seq(IcebergWrite(Jvm, w.nodeName, reasons.toSeq.sorted)) + case w: V2ExistingTableWriteExec if isIceberg(w.write) => + Seq(IcebergWrite(Spark, w.nodeName, Nil)) + // A streaming micro-batch, which Comet's split operator never plans. + case w: WriteToDataSourceV2Exec if isIcebergMicroBatch(w.batchWrite) => + Seq(IcebergWrite(Spark, w.nodeName, Nil)) + // Spark 3.4 writes a CTAS or RTAS from the create or replace exec itself. Later versions run + // that write as a nested append or overwrite, which this listener sees as a query of its own. + case w if IcebergTableAsSelectShim.writeCatalog(w).exists(isIceberg) => + Seq(IcebergWrite(Spark, w.nodeName, Nil)) + case p => p.children.flatMap(writes) + } + + private def isIcebergMicroBatch(write: BatchWrite): Boolean = write match { + case m: MicroBatchWrite => isIceberg(m.writeSupport) + case _ => false + } + + private def isIceberg(obj: AnyRef): Boolean = + obj.getClass.getName.startsWith("org.apache.iceberg.") +} diff --git a/spark/src/main/scala/org/apache/spark/Plugins.scala b/spark/src/main/scala/org/apache/spark/Plugins.scala index 6014259111e..b0680ae695d 100644 --- a/spark/src/main/scala/org/apache/spark/Plugins.scala +++ b/spark/src/main/scala/org/apache/spark/Plugins.scala @@ -31,9 +31,10 @@ import org.apache.spark.sql.internal.StaticSQLConf import org.apache.comet.{COMET_VERSION, CometSparkSessionExtensions, NativeBase} import org.apache.comet.{CometConf, ConfigEntry} -import org.apache.comet.CometConf.{COMET_METRICS_ENABLED, COMET_ONHEAP_ENABLED} +import org.apache.comet.CometConf.{COMET_ICEBERG_WRITE_REPORT_DIR, COMET_METRICS_ENABLED, COMET_ONHEAP_ENABLED} import org.apache.comet.CometKryoRegistrator import org.apache.comet.annotation.Public +import org.apache.comet.iceberg.IcebergWriteReportListener /** * Comet driver plugin. This class is loaded by Spark's plugin framework. It will be instantiated @@ -73,6 +74,7 @@ class CometDriverPlugin extends DriverPlugin with Logging { // Register Comet metrics CometDriverPlugin.registerCometMetrics(sc) + CometDriverPlugin.registerIcebergWriteReport(sc.conf) CometDriverPlugin.warnIfExecutorMemoryOverheadUnset(sc.getConf) CometDriverPlugin.warnIfMemoryPoolFractionSet(sc.getConf) @@ -229,27 +231,39 @@ object CometDriverPlugin extends Logging { COMET_METRICS_ENABLED.key, COMET_METRICS_ENABLED.defaultValue.get)) { sc.env.metricsSystem.registerSource(CometSource) - - val listenerKey = "spark.sql.queryExecutionListeners" - val listenerClass = "org.apache.comet.CometMetricsListener" - val listeners = sc.conf.get(listenerKey, "") - if (listeners.isEmpty) { - logInfo(s"Setting $listenerKey=$listenerClass") - sc.conf.set(listenerKey, listenerClass) - } else { - val currentListeners = listeners.split(",").map(_.trim) - if (!currentListeners.contains(listenerClass)) { - val newValue = s"$listeners,$listenerClass" - logInfo(s"Setting $listenerKey=$newValue") - sc.conf.set(listenerKey, newValue) - } - } + registerQueryExecutionListener(sc.conf, "org.apache.comet.CometMetricsListener") } else { logInfo( "Comet metrics reporting is disabled. Set spark.comet.metrics.enabled=true to enable.") } } + // Test-only: see COMET_ICEBERG_WRITE_REPORT_DIR. The value may come from the environment, which + // lets the Iceberg Spark test jobs turn the report on without changing the Iceberg diffs. + def registerIcebergWriteReport(conf: SparkConf): Unit = { + if (conf + .get(COMET_ICEBERG_WRITE_REPORT_DIR.key, COMET_ICEBERG_WRITE_REPORT_DIR.defaultValue.get) + .nonEmpty) { + registerQueryExecutionListener(conf, classOf[IcebergWriteReportListener].getName) + } + } + + private def registerQueryExecutionListener(conf: SparkConf, listenerClass: String): Unit = { + val listenerKey = "spark.sql.queryExecutionListeners" + val listeners = conf.get(listenerKey, "") + if (listeners.isEmpty) { + logInfo(s"Setting $listenerKey=$listenerClass") + conf.set(listenerKey, listenerClass) + } else { + val currentListeners = listeners.split(",").map(_.trim) + if (!currentListeners.contains(listenerClass)) { + val newValue = s"$listeners,$listenerClass" + logInfo(s"Setting $listenerKey=$newValue") + conf.set(listenerKey, newValue) + } + } + } + def registerCometSessionExtension(conf: SparkConf): Unit = { val extensionKey = StaticSQLConf.SPARK_SESSION_EXTENSIONS.key val extensionClass = classOf[CometSparkSessionExtensions].getName diff --git a/spark/src/main/spark-3.4/org/apache/comet/iceberg/IcebergTableAsSelectShim.scala b/spark/src/main/spark-3.4/org/apache/comet/iceberg/IcebergTableAsSelectShim.scala new file mode 100644 index 00000000000..4d66a374663 --- /dev/null +++ b/spark/src/main/spark-3.4/org/apache/comet/iceberg/IcebergTableAsSelectShim.scala @@ -0,0 +1,40 @@ +/* + * 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. + */ + +package org.apache.comet.iceberg + +import org.apache.spark.sql.connector.catalog.TableCatalog +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.datasources.v2.{AtomicCreateTableAsSelectExec, AtomicReplaceTableAsSelectExec, CreateTableAsSelectExec, ReplaceTableAsSelectExec} + +/** + * Spark 3.4: CTAS and RTAS write the table from the create or replace exec itself, through + * `TableWriteExecHelper.writeWithV2`, so the write never appears as a write node of its own. + */ +private[iceberg] object IcebergTableAsSelectShim { + + /** The catalog `plan` writes through, when `plan` is a CTAS or RTAS exec. */ + def writeCatalog(plan: SparkPlan): Option[TableCatalog] = plan match { + case p: CreateTableAsSelectExec => Some(p.catalog) + case p: AtomicCreateTableAsSelectExec => Some(p.catalog) + case p: ReplaceTableAsSelectExec => Some(p.catalog) + case p: AtomicReplaceTableAsSelectExec => Some(p.catalog) + case _ => None + } +} diff --git a/spark/src/main/spark-3.5/org/apache/comet/iceberg/IcebergTableAsSelectShim.scala b/spark/src/main/spark-3.5/org/apache/comet/iceberg/IcebergTableAsSelectShim.scala new file mode 100644 index 00000000000..ef02700a9c8 --- /dev/null +++ b/spark/src/main/spark-3.5/org/apache/comet/iceberg/IcebergTableAsSelectShim.scala @@ -0,0 +1,34 @@ +/* + * 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. + */ + +package org.apache.comet.iceberg + +import org.apache.spark.sql.connector.catalog.TableCatalog +import org.apache.spark.sql.execution.SparkPlan + +/** + * Spark 3.5+: CTAS and RTAS run their write as a nested `AppendData` or `OverwriteByExpression` + * query, which is planned and reported like any other write, so no create or replace exec writes + * a table itself. + */ +private[iceberg] object IcebergTableAsSelectShim { + + /** The catalog `plan` writes through, when `plan` is a CTAS or RTAS exec. */ + def writeCatalog(plan: SparkPlan): Option[TableCatalog] = None +} diff --git a/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergTableAsSelectShim.scala b/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergTableAsSelectShim.scala new file mode 100644 index 00000000000..ef02700a9c8 --- /dev/null +++ b/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergTableAsSelectShim.scala @@ -0,0 +1,34 @@ +/* + * 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. + */ + +package org.apache.comet.iceberg + +import org.apache.spark.sql.connector.catalog.TableCatalog +import org.apache.spark.sql.execution.SparkPlan + +/** + * Spark 3.5+: CTAS and RTAS run their write as a nested `AppendData` or `OverwriteByExpression` + * query, which is planned and reported like any other write, so no create or replace exec writes + * a table itself. + */ +private[iceberg] object IcebergTableAsSelectShim { + + /** The catalog `plan` writes through, when `plan` is a CTAS or RTAS exec. */ + def writeCatalog(plan: SparkPlan): Option[TableCatalog] = None +} diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index 3edb69f3e2b..20a383f31d8 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -20,6 +20,7 @@ package org.apache.comet import java.io.File +import java.nio.file.{FileAlreadyExistsException, Files, Path} import java.sql.Timestamp import java.util.concurrent.{CountDownLatch, TimeUnit} @@ -28,8 +29,12 @@ import scala.concurrent.{Await, Future} import scala.concurrent.ExecutionContext.Implicits.global import scala.concurrent.duration.DurationInt import scala.jdk.CollectionConverters._ +import scala.util.control.NonFatal -import org.apache.spark.{SparkConf, SparkException, Success} +import org.json4s.{DefaultFormats, Formats} +import org.json4s.jackson.JsonMethods.parse + +import org.apache.spark.{CometListenerBusUtils, SparkConf, SparkException, Success, TaskContext} import org.apache.spark.rdd.RDD import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} import org.apache.spark.sql.CometTestBase @@ -43,13 +48,38 @@ import org.apache.spark.sql.connector.write.{BatchWrite, DataWriterFactory, Phys import org.apache.spark.sql.execution.{ColumnarToRowTransition, LeafExecNode, SparkPlan} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.streaming.Trigger import org.apache.spark.sql.types.{DoubleType, IntegerType, StringType, StructField, StructType} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark41Plus} -import org.apache.comet.iceberg.IcebergReflection +import org.apache.comet.iceberg.{IcebergReflection, IcebergWriteReportListener} private case class WriteSnapshot(snapshotDelta: Long, plans: Seq[SparkPlan]) +private object IcebergTestFiles { + + /** Relative paths of every regular parquet file under `root`. */ + def parquetFiles(root: Path): Set[String] = { + if (!Files.exists(root)) return Set.empty + val stream = Files.walk(root) + try { + stream + .iterator() + .asScala + .filter(path => + Files.isRegularFile(path) && path.getFileName.toString.endsWith(".parquet")) + .map(path => root.relativize(path).toString) + .toSet + } finally stream.close() + } +} + +private case class ReportedWrite( + writer: String, + node: String, + reasons: Seq[String], + failed: Boolean) + class CometIcebergWriteActionSuite extends CometTestBase with AdaptiveSparkPlanHelper @@ -58,6 +88,9 @@ class CometIcebergWriteActionSuite override protected def sparkConf: SparkConf = { super.sparkConf .set(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key, "true") + // local[N,M] sets task max failures to M; the retry test needs one retry, and + // spark.task.maxFailures does not override this part of a local master URL. + .setMaster("local[5,2]") .set( "spark.sql.extensions", "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions") @@ -1825,6 +1858,76 @@ class CometIcebergWriteActionSuite } } + test("native acceleration: a mid-write failure retries without orphan files") { + assumeNativeAcceleration() + withIcebergCatalog { warehouseDir => + val session = spark + import session.implicits._ + (1 to 10000) + .map(i => (i, s"r$i", i.toDouble)) + .toDF("id", "region", "amount") + .coalesce(1) + .createOrReplaceTempView("retry_src") + createTable( + warehouseDir, + "retry_target", + partitionSpec = "", + properties = Some("'write.target-file-size-bytes'='1'")) + NativeWriteRetryProbe.reset() + val dataLocation = dataDir("retry_target").getAbsolutePath + spark.udf.register( + "reject_next_native_file_once", + (id: Int) => NativeWriteRetryProbe.check(id, dataLocation)) + + val snapshot = withNativeEnabled { + captureWrite("retry_target") { + withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "1000") { + spark.sql(s"INSERT INTO $catalog.$ns.retry_target " + + "SELECT reject_next_native_file_once(id), region, amount FROM retry_src") + } + } + } + assert(snapshot.snapshotDelta == 1L, s"expected one snapshot, got $snapshot") + assert( + NativeWriteRetryProbe.blockerCreated, + "the storage failure was not armed: " + + NativeWriteRetryProbe.failureReason.getOrElse("unknown reason")) + assert(NativeWriteRetryProbe.retrySeen, "Spark did not run a retry attempt") + assert(NativeWriteRetryProbe.blockersRemoved, "the retry did not remove every blocker") + val failedPaths = NativeWriteRetryProbe.failedPaths + assert(failedPaths.nonEmpty, "the failing attempt had not finalized a data file") + assert( + snapshot.plans.exists(p => + collectWithSubqueries(p) { case w: CometIcebergWriteExec => w }.nonEmpty), + s"retry did not use the native writer: ${snapshot.plans.mkString("\n--\n")}") + + val physical = parquetFiles(dataDir("retry_target")) + val root = new File(dataLocation).toPath.toAbsolutePath + val referenced = spark + .sql(s"SELECT file_path FROM $catalog.$ns.retry_target.files") + .collect() + .map { row => + val location = row.getString(0) + val uri = new java.net.URI(location) + val file = if (uri.getScheme == null) new File(location) else new File(uri) + root.relativize(file.toPath).toString + } + .toSet + assert(referenced.nonEmpty) + assert(physical == referenced, s"orphan files: ${physical -- referenced}") + assert( + (failedPaths intersect physical).isEmpty, + s"failed attempt files survived: $failedPaths") + assert( + (failedPaths intersect referenced).isEmpty, + s"a failed attempt file was referenced by the manifest: $failedPaths") + val counts = spark + .sql(s"SELECT count(*), count(DISTINCT id) FROM $catalog.$ns.retry_target") + .head() + assert(counts.getLong(0) == 10000L && counts.getLong(1) == 10000L) + } + } + // A three-task write where one task fails only after the other two have finished: their // commit messages reached the driver, so it is the committer's job abort, not task cleanup, // that has to remove their data files. @@ -2447,19 +2550,8 @@ class CometIcebergWriteActionSuite } /** Relative paths of every parquet file under `dir`, or empty when it does not exist yet. */ - private def parquetFiles(dir: File): Set[String] = { - if (!dir.exists()) return Set.empty - val root = dir.toPath - val stream = java.nio.file.Files.walk(root) - try { - stream - .iterator() - .asScala - .filter(p => p.toString.endsWith(".parquet")) - .map(p => root.relativize(p).toString) - .toSet - } finally stream.close() - } + private def parquetFiles(dir: File): Set[String] = + IcebergTestFiles.parquetFiles(dir.toPath) private def countSnapshots(tableName: String): Long = try { @@ -2629,6 +2721,128 @@ class CometIcebergWriteActionSuite } } + test("write report records which writer ran each Iceberg write") { + assumeNativeAcceleration() + withIcebergCatalog { warehouseDir => + createTable(warehouseDir, "report_parquet", partitionSpec = "PARTITIONED BY (region)") + createTable( + warehouseDir, + "report_orc", + partitionSpec = "", + properties = Some("'write.format.default'='orc'")) + + val writes = reportedWrites { + withNativeEnabled { + // Collecting the result runs a second query over the command's result; the write + // must still be reported once. + spark + .sql(s"INSERT INTO $catalog.$ns.report_parquet VALUES (1, 'us-east', 1.5)") + .collect() + spark.sql(s"INSERT INTO $catalog.$ns.report_orc VALUES (2, 'eu', 2.5)") + } + withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "false") { + spark.sql(s"INSERT INTO $catalog.$ns.report_parquet VALUES (3, 'eu', 3.5)") + } + } + assert(writes.map(_.writer) == Seq("native", "jvm", "spark"), writes.mkString("\n")) + assert(writes(1).reasons.exists(_.contains("only parquet"))) + assert(writes(2).node == "AppendData") + assert(writes.forall(!_.failed)) + } + } + + test("write report records a CTAS and an RTAS once each") { + assumeNativeAcceleration() + withIcebergCatalog { _ => + val writes = reportedWrites { + withNativeEnabled { + spark.sql(s""" + CREATE TABLE $catalog.$ns.report_ctas USING iceberg AS + SELECT * FROM VALUES (1, 'us', 1.0), (2, 'eu', 2.0) AS t(id, region, amount) + """) + spark.sql(s""" + REPLACE TABLE $catalog.$ns.report_ctas USING iceberg AS + SELECT * FROM VALUES (3, 'us', 3.0) AS t(id, region, amount) + """) + } + } + // Spark 3.5+ runs each write as a nested append or overwrite, which Comet's split operator + // plans. Spark 3.4 writes from the create and replace execs, which it never sees. + val expected = + if (isSpark35Plus) Seq("native" -> "CometIcebergWrite", "native" -> "CometIcebergWrite") + else Seq("spark" -> "AtomicCreateTableAsSelect", "spark" -> "AtomicReplaceTableAsSelect") + assert(writes.map(w => w.writer -> w.node) == expected, writes.mkString("\n")) + assert(writes.forall(!_.failed)) + assertRows("report_ctas", Seq(3)) + } + } + + test("write report records a streaming micro-batch write") { + assumeNativeAcceleration() + withIcebergCatalog { warehouseDir => + createTable(warehouseDir, "report_stream", partitionSpec = "") + withTempIcebergDir { dir => + val source = new File(dir, "source").getAbsolutePath + val session = spark + import session.implicits._ + Seq((1, "us", 1.0), (2, "eu", 2.0)).toDF("id", "region", "amount").write.parquet(source) + val schema = spark.table(s"$catalog.$ns.report_stream").schema + + val writes = reportedWrites { + withNativeEnabled { + spark.readStream + .schema(schema) + .parquet(source) + .writeStream + .format("iceberg") + .outputMode("append") + .trigger(Trigger.AvailableNow()) + .option("checkpointLocation", new File(dir, "checkpoint").getAbsolutePath) + .toTable(s"$catalog.$ns.report_stream") + .awaitTermination() + } + } + assert( + writes.map(w => w.writer -> w.node) == Seq("spark" -> "WriteToDataSourceV2"), + writes.mkString("\n")) + assert(writes.forall(!_.failed)) + assertRows("report_stream", Seq(1, 2)) + } + } + } + + /** The Iceberg writes `IcebergWriteReportListener` records while `action` runs, in order. */ + private def reportedWrites(action: => Unit): Seq[ReportedWrite] = { + var writes = Seq.empty[ReportedWrite] + withTempIcebergDir { reportDir => + val listener = new IcebergWriteReportListener( + new SparkConf() + .set(CometConf.COMET_ICEBERG_WRITE_REPORT_DIR.key, reportDir.getAbsolutePath)) + spark.listenerManager.register(listener) + try { + action + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + } finally { + spark.listenerManager.unregister(listener) + } + + implicit val formats: Formats = DefaultFormats + writes = reportDir + .listFiles() + .toSeq + .flatMap(f => Files.readAllLines(f.toPath).asScala) + .map { line => + val w = parse(line) + ReportedWrite( + (w \ "writer").extract[String], + (w \ "node").extract[String], + (w \ "reasons").extract[Seq[String]], + (w \ "failed").extract[Boolean]) + } + } + writes + } + private def assertNativeWriteEngages(tableName: String, expectedIds: Seq[Int])( action: => Unit): Unit = { val snapshot = withNativeEnabled { captureWrite(tableName)(action) } @@ -2735,6 +2949,103 @@ class CometIcebergWriteActionSuite } +/** Makes the local file store reject a later data-file write in the first native attempt. */ +private object NativeWriteRetryProbe { + @volatile private var firstAttemptFiles = Set.empty[String] + @volatile private var sawRetry = false + @volatile private var blockerPaths = Set.empty[String] + @volatile private var createdBlocker = false + @volatile private var firstAttemptFailure: Option[String] = None + + def reset(): Unit = synchronized { + firstAttemptFiles = Set.empty + sawRetry = false + blockerPaths = Set.empty + createdBlocker = false + firstAttemptFailure = None + } + + def failedPaths: Set[String] = firstAttemptFiles + + def retrySeen: Boolean = sawRetry + + def blockerCreated: Boolean = createdBlocker + + def blockersRemoved: Boolean = blockerPaths.isEmpty + + def failureReason: Option[String] = firstAttemptFailure + + private def failProbe(reason: String): Nothing = { + firstAttemptFailure = Some(reason) + throw new IllegalStateException(reason) + } + + private def armBlockers(dataLocation: String): Unit = { + val root = new File(dataLocation).toPath + val files = IcebergTestFiles.parquetFiles(root) + firstAttemptFiles = files + // The unpartitioned writer has at most one file open; two paths mean at least one + // earlier file has already been finalized before this input-side failure. + if (files.size < 2) { + failProbe("native writer did not finalize before injection") + } + val numbered = files.map { relative => + val name = new File(relative).getName + val pattern = "^(.*)-(\\d{5})\\.parquet$".r + name match { + case pattern(prefix, number) => (prefix, number.toInt) + case _ => failProbe(s"unexpected native file name: $name") + } + } + val prefixes = numbered.map(_._1) + if (prefixes.size != 1) { + failProbe(s"expected one native task prefix, got $prefixes") + } + val maxFileNumber = numbered.map(_._2).max + (1 to 4).foreach { offset => + val number = maxFileNumber + offset + val blocker = root.resolve(s"${prefixes.head}-${"%05d".format(number)}.parquet") + try { + Files.createDirectory(blocker) + blockerPaths += blocker.toString + createdBlocker = true + } catch { + // The writer can advance between the directory walk and blocker creation. Keep + // arming later file numbers rather than turning that harmless skew into the failure. + case _: FileAlreadyExistsException => + } + } + if (!createdBlocker) { + failProbe(s"native writer advanced past all blocker candidates after $maxFileNumber") + } + // The file:// store rejects writes at these paths with EISDIR. A small range tolerates + // the native pipeline being a few files ahead of the projection's filesystem view. + // Attempt-unique names keep the blockers clear of the second writer until it removes them. + } + + def check(id: Int, dataLocation: String): Int = { + if (id == 7000) { + val attempt = TaskContext.get().attemptNumber() + if (attempt == 0) { + try armBlockers(dataLocation) + catch { + case NonFatal(e) => + if (firstAttemptFailure.isEmpty) { + val detail = Option(e.getMessage).filter(_.nonEmpty).getOrElse(e.getClass.getName) + firstAttemptFailure = Some(detail) + } + throw e + } + } else { + blockerPaths.foreach(path => Files.deleteIfExists(new File(path).toPath)) + blockerPaths = Set.empty + sawRetry = true + } + } + id + } +} + /** * Blocks the DELETE's write job between its scan-snapshot pin and its commit so the test can * inject a conflicting commit. Top-level so the UDF closure doesn't capture the suite. diff --git a/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala b/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala index 4207bd6621e..82ea0a18894 100644 --- a/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometPluginsSuite.scala @@ -88,6 +88,22 @@ class CometPluginsSuite extends CometTestBase { } } + test("Iceberg write report listener is registered only when a report directory is set") { + val listenerKey = "spark.sql.queryExecutionListeners" + val listenerClass = "org.apache.comet.iceberg.IcebergWriteReportListener" + + val unset = new SparkConf() + CometDriverPlugin.registerIcebergWriteReport(unset) + assert(!unset.contains(listenerKey)) + + val set = new SparkConf() + .set(CometConf.COMET_ICEBERG_WRITE_REPORT_DIR.key, "/tmp/report") + .set(listenerKey, "foo") + CometDriverPlugin.registerIcebergWriteReport(set) + CometDriverPlugin.registerIcebergWriteReport(set) + assert(set.get(listenerKey) == s"foo,$listenerClass") + } + test("Comet version is exposed as a Spark config") { // The driver plugin sets spark.comet.version, which is then visible both on the SparkContext // conf and through the session runtime config (SET / spark.conf.get).