andygrove commented on code in PR #6155: URL: https://github.com/apache/datafusion-comet/pull/6155#discussion_r4109328524
########## spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteReportListener.scala: ########## @@ -0,0 +1,119 @@ +/* + * 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.execution.{CommandResultExec, QueryExecution, SparkPlan} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, QueryStageExec} +import org.apache.spark.sql.execution.datasources.v2.V2ExistingTableWriteExec +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 Review Comment: On 3.4 the CTAS and RTAS execs write through `TableWriteExecHelper.writeWithV2` themselves, while 3.5+ runs the write as a nested `AppendData` or `OverwriteByExpression` that the listener already records. b4dd3d0c2 adds `IcebergTableAsSelectShim`, which recognizes the four execs on 3.4 only. There's also a CTAS plus RTAS reporting test that expects two `spark` records on 3.4 and one `native` record each on 3.5+. With the shim disabled, the 3.4 run records nothing, as you found. ########## dev/ci/summarize-iceberg-writes.py: ########## @@ -0,0 +1,149 @@ +# 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 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 report_files(roots): + """Every report file under roots, keeping only the latest attempt of each shard.""" + latest = {} + files = [] + for root in roots: + for path in sorted(Path(root).rglob("*.jsonl")): Review Comment: Fixed in b4dd3d0c2. The latest attempt now comes from the shard artifact directories, and a shard whose latest attempt recorded nothing is named above the table. `dev/ci/test-summarize-iceberg-writes.py` covers your empty-retry case and runs in Preflight. ########## spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteReportListener.scala: ########## @@ -0,0 +1,119 @@ +/* + * 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.execution.{CommandResultExec, QueryExecution, SparkPlan} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, QueryStageExec} +import org.apache.spark.sql.execution.datasources.v2.V2ExistingTableWriteExec +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 w.write.getClass.getName.startsWith("org.apache.iceberg.") => + Seq(IcebergWrite(Spark, w.nodeName, Nil)) + case p => p.children.flatMap(writes) Review Comment: The micro-batch does reach the listener, through the batch's `collect`, but as `WriteToDataSourceV2Exec`. b4dd3d0c2 records it as a `spark` write when its `MicroBatchWrite` wraps an Iceberg streaming write. The streaming reporting test fails without that match. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
