diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index b38d33fabd9..b031c051bb0 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -1278,7 +1278,10 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( if let Some(producer) = &mut exec_context.batch_producer { match producer.next_batch()? { Some(batch) => { - update_metrics(env, exec_context)?; + // Publish on the configured interval, as the ScanExec path below does, + // since each publish walks the whole metric tree and calls into the JVM. + // `releasePlan` publishes the final values. + update_metrics_on_interval(env, exec_context)?; return prepare_output( env, array_addrs, diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala index 29271220f71..c447c7e775e 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala @@ -36,7 +36,7 @@ import org.apache.spark.sql.comet.execution.shuffle.CometBypassMergeSortShuffleH import org.apache.spark.sql.comet.execution.shuffle.CometColumnarShuffle import org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffle import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec -import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.{SparkPlan, SQLExecution} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike import org.apache.spark.sql.execution.metric.SQLMetric @@ -1154,6 +1154,55 @@ class CometTaskMetricsSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("native block without a JVM input publishes SQL metrics on the update interval") { + withTempPath { dir => + spark.range(0, 10000, 1, 1).write.parquet(dir.getAbsolutePath) + // Ten output batches from one task. Passes `check` the output_rows the task sees right after + // the first batch, and again after the last. The task creates the native plan, which starts + // the interval's clock, when it builds the iterator, so the sleep before the first batch runs + // after that. + def checkOutputRows(interval: String)(check: (Long, Long) => Unit): Unit = + withSQLConf( + CometConf.COMET_BATCH_SIZE.key -> "1000", + CometConf.COMET_METRICS_UPDATE_INTERVAL.key -> interval) { + val plan = spark.read.parquet(dir.getAbsolutePath).queryExecution.executedPlan + val scan = find(plan)(_.isInstanceOf[CometNativeScanExec]) + .getOrElse(fail(s"Expected CometNativeScanExec in plan:\n${plan.treeString}")) + .asInstanceOf[CometNativeScanExec] + // The task deserializes this metric together with the scan's metric node, so both + // refer to the task-side copy that native execution publishes into. + val outputRows = scan.metrics("output_rows") + val (midStream, atEnd) = SQLExecution.withSQLConfPropagated(spark) { + scan + .executeColumnar() + .mapPartitions { batches => + Thread.sleep(100) + batches.next() + val mid = outputRows.value + batches.foreach(_ => ()) + Iterator((mid, outputRows.value)) + } + .collect() + .head + } + check(midStream, atEnd) + } + + // With the interval disabled, the only publish is the one in releasePlan, which a + // per-batch publish would break after the first batch. + checkOutputRows("-1") { (midStream, atEnd) => + assert(midStream == 0, s"output_rows was published mid-stream: $midStream") + assert(atEnd == 10000, s"releasePlan should publish the final output_rows, got $atEnd") + } + + // A 1 ms interval has passed by the time the first batch arrives, so that batch publishes. + checkOutputRows("1") { (midStream, atEnd) => + assert(midStream > 0, "output_rows was not published mid-stream once the interval passed") + assert(atEnd == 10000, s"releasePlan should publish the final output_rows, got $atEnd") + } + } + } + test("native scan left unconsumed by a limit still reports task input metrics") { withTempPath { dir => spark