sunchao commented on code in PR #7325: URL: https://github.com/apache/datafusion-comet/pull/7325#discussion_r4417919497
########## spark/src/main/scala/org/apache/spark/sql/comet/CometCollectMetricsExec.scala: ########## @@ -0,0 +1,154 @@ +/* + * 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.spark.sql.comet + +import org.apache.spark.TaskContext +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.NamedExpression +import org.apache.spark.sql.catalyst.trees.TreeNodeTag +import org.apache.spark.sql.execution.{AggregatingAccumulator, CollectMetricsExec, ColumnarToRowExec, SparkPlan} +import org.apache.spark.sql.vectorized.ColumnarBatch + +import org.apache.comet.{CometConf, ConfigEntry} +import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus +import org.apache.comet.serde.{Compatible, OperatorOuterClass, SupportLevel, Unsupported} +import org.apache.comet.serde.operator.CometSink + +object CometCollectMetricsExec extends CometSink[CollectMetricsExec] { + + val REQUIRES_ROW_CONSUMPTION: TreeNodeTag[Boolean] = + TreeNodeTag[Boolean]("comet.collectMetrics.requiresRowConsumption") + + // Keep Spark's own accumulator for observation discovery and reset lifecycle. The accessor + // is private on supported Spark 3/4 versions, so resolve it once rather than duplicating + // those version-specific contracts. + private lazy val accumulatorAccessor = { + val method = classOf[CollectMetricsExec].getDeclaredMethod("accumulator") + method.setAccessible(true) + method + } + + private def accumulator(op: CollectMetricsExec): AggregatingAccumulator = + accumulatorAccessor.invoke(op).asInstanceOf[AggregatingAccumulator] + + // Spark's setState is public on the JVM but scoped to the execution package in Scala. + private lazy val setStateMethod = + classOf[AggregatingAccumulator].getMethod("setState", classOf[AggregatingAccumulator]) + + private def publishTaskState( + collector: AggregatingAccumulator, + updater: AggregatingAccumulator): Unit = { + if (collector.isZero) { + // Adopt the first buffer without copying large aggregate states such as collect_list. + val _ = setStateMethod.invoke(collector, updater) + } else { + collector.merge(updater) + } + } + + override def enabledConfig: Option[ConfigEntry[Boolean]] = + Some(CometConf.COMET_EXEC_COLLECT_METRICS_ENABLED) + + override def getSupportLevel(op: CollectMetricsExec): SupportLevel = { + if (op.getTagValue(REQUIRES_ROW_CONSUMPTION).contains(true)) { + Unsupported(Some("Observation below a short-circuiting consumer requires Spark rows")) + } else { + Compatible() + } + } + + override def createExec( + nativeOp: OperatorOuterClass.Operator, + op: CollectMetricsExec): CometNativeExec = { + val collector = new CometCollectMetricsExec(op.name, op.metricExpressions, op.child) + collector.copyTagsFrom(op) + op.logicalLink.foreach(collector.setLogicalLink) + CometSinkPlaceHolder(nativeOp, op, collector) + } +} + +/** + * Evaluate observation expressions with Spark while forwarding the original Arrow batches. + * Extending CollectMetricsExec keeps the node visible to Spark's observation discovery. + */ +class CometCollectMetricsExec( + name: String, + metricExpressions: Seq[NamedExpression], + child: SparkPlan) + extends CollectMetricsExec(name, metricExpressions, child) + with CometPlan { + + override def supportsColumnar: Boolean = true + + override protected def doExecuteColumnar(): RDD[ColumnarBatch] = { + val collector = CometCollectMetricsExec.accumulator(this) + // Spark 3 resets observations before execution; Spark 4 does so in resetMetrics(). + if (!isSpark40Plus) { + collector.reset() + } + child.executeColumnar().mapPartitions { batches => + // Publish only at task completion, as Spark does, so heartbeats cannot expose partial + // updates. Spark's accumulator registration also excludes failed task attempts. + val updater = collector.copyAndReset() + val _ = TaskContext.get().addTaskCompletionListener[Unit] { _ => + CometCollectMetricsExec.publishTaskState(collector, updater) + } + batches.map { batch => + val rows = batch.rowIterator() + while (rows.hasNext) { + updater.add(rows.next()) + } + batch Review Comment: A `CometLocalTableScan` returns batches that share the Arrow reader’s reusable vectors. Forwarding one here sends it to the downstream native projection’s `ColumnarBatchArrowReader`, which closes it after loading it. That clears the struct’s child vectors before the source can fill the next batch. I reproduced this on Spark 4.1.3 with ten `VALUES` rows containing `named_struct('v', i)`, `spark.comet.batchSize=2`, native local table scans enabled, whole-stage codegen and AQE off, and `ConvertToLocalRelation` excluded. `observe(count(lit(1)))` followed by `selectExpr("xxhash64(a)").repartition(2)` returns all ten hashes and a count of 10 with only `collectMetrics.enabled=false`. Enabling it fails on the second batch with `no more field nodes for field v`. Could we preserve reader ownership or keep the Spark collector for these sources, and cover a multi-batch struct? `CometShuffleExchangeExec` already handles reader-owned batches separately. ########## spark/src/main/scala/org/apache/spark/sql/comet/CometCollectMetricsExec.scala: ########## @@ -0,0 +1,154 @@ +/* + * 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.spark.sql.comet + +import org.apache.spark.TaskContext +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.NamedExpression +import org.apache.spark.sql.catalyst.trees.TreeNodeTag +import org.apache.spark.sql.execution.{AggregatingAccumulator, CollectMetricsExec, ColumnarToRowExec, SparkPlan} +import org.apache.spark.sql.vectorized.ColumnarBatch + +import org.apache.comet.{CometConf, ConfigEntry} +import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus +import org.apache.comet.serde.{Compatible, OperatorOuterClass, SupportLevel, Unsupported} +import org.apache.comet.serde.operator.CometSink + +object CometCollectMetricsExec extends CometSink[CollectMetricsExec] { + + val REQUIRES_ROW_CONSUMPTION: TreeNodeTag[Boolean] = + TreeNodeTag[Boolean]("comet.collectMetrics.requiresRowConsumption") + + // Keep Spark's own accumulator for observation discovery and reset lifecycle. The accessor + // is private on supported Spark 3/4 versions, so resolve it once rather than duplicating + // those version-specific contracts. + private lazy val accumulatorAccessor = { + val method = classOf[CollectMetricsExec].getDeclaredMethod("accumulator") + method.setAccessible(true) + method + } + + private def accumulator(op: CollectMetricsExec): AggregatingAccumulator = + accumulatorAccessor.invoke(op).asInstanceOf[AggregatingAccumulator] + + // Spark's setState is public on the JVM but scoped to the execution package in Scala. + private lazy val setStateMethod = + classOf[AggregatingAccumulator].getMethod("setState", classOf[AggregatingAccumulator]) + + private def publishTaskState( + collector: AggregatingAccumulator, + updater: AggregatingAccumulator): Unit = { + if (collector.isZero) { + // Adopt the first buffer without copying large aggregate states such as collect_list. + val _ = setStateMethod.invoke(collector, updater) + } else { + collector.merge(updater) + } + } + + override def enabledConfig: Option[ConfigEntry[Boolean]] = + Some(CometConf.COMET_EXEC_COLLECT_METRICS_ENABLED) + + override def getSupportLevel(op: CollectMetricsExec): SupportLevel = { + if (op.getTagValue(REQUIRES_ROW_CONSUMPTION).contains(true)) { + Unsupported(Some("Observation below a short-circuiting consumer requires Spark rows")) + } else { + Compatible() + } + } + + override def createExec( + nativeOp: OperatorOuterClass.Operator, + op: CollectMetricsExec): CometNativeExec = { + val collector = new CometCollectMetricsExec(op.name, op.metricExpressions, op.child) + collector.copyTagsFrom(op) + op.logicalLink.foreach(collector.setLogicalLink) + CometSinkPlaceHolder(nativeOp, op, collector) + } +} + +/** + * Evaluate observation expressions with Spark while forwarding the original Arrow batches. + * Extending CollectMetricsExec keeps the node visible to Spark's observation discovery. + */ +class CometCollectMetricsExec( + name: String, + metricExpressions: Seq[NamedExpression], + child: SparkPlan) + extends CollectMetricsExec(name, metricExpressions, child) + with CometPlan { + + override def supportsColumnar: Boolean = true + + override protected def doExecuteColumnar(): RDD[ColumnarBatch] = { + val collector = CometCollectMetricsExec.accumulator(this) + // Spark 3 resets observations before execution; Spark 4 does so in resetMetrics(). + if (!isSpark40Plus) { + collector.reset() + } + child.executeColumnar().mapPartitions { batches => + // Publish only at task completion, as Spark does, so heartbeats cannot expose partial + // updates. Spark's accumulator registration also excludes failed task attempts. + val updater = collector.copyAndReset() + val _ = TaskContext.get().addTaskCompletionListener[Unit] { _ => + CometCollectMetricsExec.publishTaskState(collector, updater) + } + batches.map { batch => + val rows = batch.rowIterator() + while (rows.hasNext) { + updater.add(rows.next()) Review Comment: Some valid Spark aggregates copy their input. With `spark.comet.exec.localTableScan.enabled=true`, I observed `collect_list(v)` over two `VALUES` rows where `v` was `named_struct('i', make_interval(0, 1), 'id', id)`, then selected `id` and repartitioned. Spark 4.1.3 and Comet with only `collectMetrics.enabled=false` both returned the two structs. With this collector enabled, `CollectList` calls `InternalRow.copyValue` on the row view and `ColumnarRow.copy` throws `Not implemented. CalendarIntervalType`. `collect_set` also fails on this input. Could we give these aggregates a row they can copy, or fall back for the affected inputs? A nested-interval regression test would exercise the operation that currently fails. ########## spark/src/test/scala/org/apache/spark/sql/comet/CometCollectMetricsSuite.scala: ########## @@ -0,0 +1,509 @@ +/* + * 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.spark.sql.comet + +import java.util.concurrent.{FutureTask, TimeUnit} + +import org.apache.spark.sql.{CometTestBase, Observation, Row} +import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute +import org.apache.spark.sql.execution.{CollectMetricsExec, MapPartitionsExec, QueryExecution, SortExec, SparkPlan} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec +import org.apache.spark.sql.execution.joins.SortMergeJoinExec +import org.apache.spark.sql.expressions.Window +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.util.QueryExecutionListener + +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus + +class CometCollectMetricsSuite extends CometTestBase { + + import testImplicits._ + + private def observedMetrics(observation: Observation): Map[String, Any] = { + // A missing observation in the executed plan must fail instead of blocking in get forever. + // Observation.get is available on every supported Spark version, including Spark 3.4. + val result = new FutureTask[Map[String, Any]](() => observation.get) + val reader = new Thread(result, "comet-observation-reader") + reader.setDaemon(true) + reader.start() + try { + result.get(10, TimeUnit.SECONDS) + } finally { + result.cancel(true) + } + } + + private def withCometEnabled[T](enabled: Boolean)(body: => T): T = { + // Spark 3.x withSQLConf returns Unit, so capture the result explicitly. + var result: Option[T] = None + withSQLConf(CometConf.COMET_ENABLED.key -> enabled.toString) { + result = Some(body) + } + result.get + } + + for (adaptive <- Seq(false, true)) { + test(s"observation preserves counters and native consumers, AQE=$adaptive") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + CometConf.COMET_SHUFFLE_MODE.key -> "native") { + val data = (0 until 12).map(i => (i, i % 3, if (i % 4 == 0) None else Some(i))) + withParquetTable(data, "observed_input") { + def run(comet: Boolean): (Seq[Row], Map[String, Any], SparkPlan) = { + withCometEnabled(comet) { + val observation = Observation("source_counts") + val result = spark + .table("observed_input") + .observe( + observation, + count(lit(1)).as("total_rows"), + count(when($"_2" === 1, 1)).as("selected_rows"), + sum($"_3").as("total_value")) + .filter($"_1" >= 4) + .select($"_2".as("k"), ($"_1" + 1).as("v")) + .groupBy("k") + .agg(sum("v").as("total")) + val rows = result.collect().toSeq + (rows, observedMetrics(observation), result.queryExecution.executedPlan) + } + } + + val (sparkRows, sparkMetrics, _) = run(comet = false) + val (cometRows, cometMetrics, plan) = run(comet = true) + assert(cometRows.toSet == sparkRows.toSet) + assert(cometMetrics == sparkMetrics) + assert( + cometMetrics == Map("total_rows" -> 12L, "selected_rows" -> 4L, "total_value" -> 54L)) + assert(collect(plan) { case _: CometCollectMetricsExec => true }.size == 1) + assert(collect(plan) { + case project: CometProjectExec if collect(project) { + case _: CometCollectMetricsExec => true + }.nonEmpty => + project + }.nonEmpty) + assert(collect(plan) { case _: CometHashAggregateExec => true }.nonEmpty) + if (adaptive) { + assert(plan.asInstanceOf[AdaptiveSparkPlanExec].isFinalPlan) + } + } + } + } + } + + test("observation publishes collect_list state without an executor merge") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withTempPath { path => + spark.range(0, 100, 1, 1).write.parquet(path.getCanonicalPath) + + def run(comet: Boolean): (Seq[Long], Long, SparkPlan) = { + withCometEnabled(comet) { + val executorMergedValues = spark.sparkContext.longAccumulator + val values = getColumnFromExpression( + new MergeTrackingCollectList(UnresolvedAttribute("id"), 0, 0, executorMergedValues) + .toAggregateExpression()) + val observation = Observation("collected_values") + val result = spark.read + .parquet(path.getCanonicalPath) + .observe(observation, values.as("values")) + .repartition(1) + assert(result.collect().map(_.getLong(0)).sorted.toSeq == (0L until 100L)) + val observed = observedMetrics(observation)("values") + .asInstanceOf[scala.collection.Seq[Long]] + .sorted + .toList + (observed, executorMergedValues.value, result.queryExecution.executedPlan) + } + } + + val (sparkValues, sparkMergedValues, _) = run(comet = false) + val (cometValues, cometMergedValues, plan) = run(comet = true) + assert(collect(plan) { case _: CometCollectMetricsExec => true }.size == 1) + assert(sparkValues == (0L until 100L)) + assert(cometValues == sparkValues) + assert(sparkMergedValues == 0L) + assert(cometMergedValues == sparkMergedValues) + } + } + } + + test("empty observation publishes aggregate defaults") { + withParquetTable(Seq((1, 2), (3, 4)), "observed_input") { + val observation = Observation("empty") + val result = spark + .table("observed_input") + .filter($"_1" < 0) + .observe(observation, count(lit(1)).as("rows"), sum($"_2").as("total")) + .repartition(1) + .select($"_1" + 1) + assert(result.collect().isEmpty) + assert(observedMetrics(observation) == Map("rows" -> 0L, "total" -> null)) + assert(collect(result.queryExecution.executedPlan) { case _: CometCollectMetricsExec => + true + }.nonEmpty) + } + } + + test("named observation preserves Spark producer reexecution and reset behavior") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withParquetTable(Seq((1, 2), (3, 4)), "observed_input") { + def run(comet: Boolean): Seq[Row] = { + withCometEnabled(comet) { + val result = spark + .table("observed_input") + .observe("repeated", count(lit(1)).as("rows"), sum($"_2").as("total")) + .repartition(1) + val plan = result.queryExecution.executedPlan + val collector = collect(plan) { case op: CollectMetricsExec => op }.head + if (comet) { + assert(collector.isInstanceOf[CometCollectMetricsExec]) + } + // Reexecute the actual producer: collecting the whole query would reuse its shuffle. + val metrics = (0 until 2).map { _ => + assert(collector.executeCollect().map(_.getInt(0)).toSet == Set(1, 3)) + CollectMetricsExec.collect(plan)("repeated") + } + plan.resetMetrics() + metrics :+ CollectMetricsExec.collect(plan)("repeated") + } + } + val expected = run(comet = false) + assert(run(comet = true) == expected) + } + } + } + + for (fraction <- Seq(0.0, 1.0)) { + test(s"observation below replacement sampling preserves row demand, fraction=$fraction") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "false", + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false") { + withTempPath { path => + spark.range(0, 100, 1, 1).write.parquet(path.getCanonicalPath) + + def run(comet: Boolean): (Seq[Long], Map[String, Any], SparkPlan) = { + withCometEnabled(comet) { + val observation = Observation("sampled_input") + val result = spark.read + .parquet(path.getCanonicalPath) + .observe(observation, count(lit(1)).as("rows")) + .select(($"id" + 1).as("k")) + .sample(withReplacement = true, fraction = fraction, seed = 42L) + .repartition(2) + val rows = result.collect().map(_.getLong(0)).sorted.toSeq + (rows, observedMetrics(observation), result.queryExecution.executedPlan) + } + } + + val (sparkRows, sparkMetrics, _) = run(comet = false) + val (cometRows, cometMetrics, plan) = run(comet = true) + assert(collect(plan) { case _: CometNativeScanExec => true }.nonEmpty) + assert(cometRows == sparkRows) + assert(sparkMetrics == Map("rows" -> (if (fraction == 0.0) 0L else 100L))) + assert(cometMetrics == sparkMetrics) + assert( + collect(plan) { case _: CometCollectMetricsExec => true }.nonEmpty == + (fraction > 0.0)) + } + } + } + } + + test("observation below a limit preserves Spark row demand") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withParquetTable((0 until 100).map(i => (i, i)), "observed_input") { + def run(comet: Boolean): (Map[String, Any], SparkPlan) = { + withCometEnabled(comet) { + val observation = Observation("limited") + val result = spark + .table("observed_input") + .observe(observation, count(lit(1)).as("rows")) + .limit(1) + assert(result.collect().length == 1) + (observedMetrics(observation), result.queryExecution.executedPlan) + } + } + val (expected, _) = run(comet = false) + val (actual, plan) = run(comet = true) + assert(actual == expected) + assert(actual == Map("rows" -> 1L)) + assert(collect(plan) { case _: CometCollectMetricsExec => true }.isEmpty) + assert(collect(plan) { case _: CollectMetricsExec => true }.nonEmpty) + } + } + } + + test("observation below a window group limit preserves Spark row demand") { + assume(isSpark35Plus, "WindowGroupLimit was added in Spark 3.5") + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + "spark.sql.optimizer.windowGroupLimitThreshold" -> "1000", + "spark.sql.execution.topKSortFallbackThreshold" -> "0", + CometConf.COMET_EXEC_RANGE_ENABLED.key -> "true") { + def run(comet: Boolean): (Map[String, Any], SparkPlan) = { + withCometEnabled(comet) { + val observation = Observation("window_limited") + // Ordered single-partition input avoids a sort consuming the complete observation. + val result = spark + .range(0, 100, 1, 1) + .filter($"id" % 2 === 0) + .observe(observation, count(lit(1)).as("rows")) + .withColumn("rn", row_number().over(Window.orderBy("id"))) + .filter($"rn" <= 1) + .select("id") + assert(result.collect().toSeq == Seq(Row(0L))) + (observedMetrics(observation), result.queryExecution.executedPlan) + } + } + val (expected, _) = run(comet = false) + val (actual, plan) = run(comet = true) + assert(expected("rows").asInstanceOf[Long] < 50L) + assert(actual == expected) + assert(collect(plan) { case _: CometCollectMetricsExec => true }.isEmpty) + assert(collect(plan) { case op if op.nodeName == "WindowGroupLimit" => op }.nonEmpty) + } + } + + test("observation below a sort merge join preserves Spark row demand") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + CometConf.COMET_FORCE_SHJ.key -> "false", + CometConf.COMET_EXEC_RANGE_ENABLED.key -> "true") { + def run(comet: Boolean): (Map[String, Any], SparkPlan) = { + withCometEnabled(comet) { + val observation = Observation("join_limited") + val left = spark + .range(0, 100, 1, 1) + .filter($"id" % 2 === 0) + .observe(observation, count(lit(1)).as("rows")) + val right = spark.range(0, 1, 1, 1) + val result = left.join(right, Seq("id")) + assert(result.collect().toSeq == Seq(Row(0L))) + (observedMetrics(observation), result.queryExecution.executedPlan) + } + } + val (expected, sparkPlan) = run(comet = false) + val (actual, cometPlan) = run(comet = true) + assert(expected("rows").asInstanceOf[Long] < 50L) + assert(actual == expected) + for (plan <- Seq(sparkPlan, cometPlan)) { + assert(collect(plan) { case _: SortMergeJoinExec => true }.nonEmpty) + assert(collect(plan) { case _: CometCollectMetricsExec => true }.isEmpty) + assert(collect(plan) { case _: CollectMetricsExec => true }.nonEmpty) + } + } + } + + for (sortBeforeTake <- Seq(false, true)) { + test(s"observation below mapPartitions preserves row demand, sort=$sortBeforeTake") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withTempPath { path => + spark.range(0, 100, 1, 1).write.parquet(path.getCanonicalPath) + + def run(comet: Boolean): (Seq[Long], Map[String, Any], SparkPlan) = { + withCometEnabled(comet) { + val observation = Observation("iterator_rows") + val observed = spark.read + .parquet(path.getCanonicalPath) + .observe(observation, count(lit(1)).as("rows")) + // A local sort drains the observation before the iterator takes its first row. + val input = + if (sortBeforeTake) observed.sortWithinPartitions($"id".desc) else observed + val result = input.as[Long].mapPartitions(_.take(1)) Review Comment: Could we append `.repartition(2)` after `mapPartitions(_.take(1))` here? As written, the plan root already requires row consumption, so this test passes even if the `MapPartitionsExec` guard is removed. I removed just that guard in a temporary build and all 17 suite tests still passed. With the trailing shuffle, both variants pass on this PR and fail without the guard. The unsorted case then reports 100 observed rows where Spark reports 1. -- 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]
