sunchao commented on code in PR #7325: URL: https://github.com/apache/datafusion-comet/pull/7325#discussion_r4418748203
########## spark/src/main/scala/org/apache/spark/sql/comet/CometCollectMetricsExec.scala: ########## @@ -0,0 +1,170 @@ +/* + * 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.types.{CalendarIntervalType, DataType} +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) + + // Forwarded batches can carry widened nested nullability, as local table scans do. Preserve + // it at the next native scan instead of narrowing back to the original Catalyst field types. + override protected def scanFieldType(dt: DataType): DataType = dt.asNullable + + 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 = op match { + case collector: CometCollectMetricsExec => + // Replanning must retain the existing accumulator while using the same scan schema. + CometScanWrapper(nativeOp, collector) + case _ => + 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() + } + val inputAttributes = child.output + val requiresUnsafeRows = metricExpressions.exists(_.references.exists(attribute => + SupportLevel.containsType(attribute.dataType, classOf[CalendarIntervalType]))) + child.executeColumnar().mapPartitions { batches => + // Retaining aggregates copy their inputs, but ColumnarRow.copy cannot copy calendar + // intervals, including nested ones. Preserve the accumulator's input ordinals. + val toObservationRow: InternalRow => InternalRow = + if (requiresUnsafeRows) CometUnsafeProjection.create(inputAttributes) else identity Review Comment: **`mode` can silently report the less frequent struct.** I used a Parquet view `source` with IDs 0–9, where `v` is `struct(1 AS x)` for IDs 0–5 and `struct(2 AS x)` for IDs 6–9. I observed the following metric, then ran `.select("id").repartition(2).collect()`: ```sql mode(if(id between 3 and 5, (select v from source where id = 0), v), true) ``` The subquery returns the same `struct(1)` those three rows already contain, so the counts are six and four. Spark and pre-PR Comet report `struct(1)`. With the new collector, all ten rows are returned but the observation reports `struct(2)`. Disabling just this collector restores `struct(1)`. Could we give Spark's accumulator a consistent row representation for complex keys, or keep these metrics on Spark's collector? This case would make a good regression test because checking only the output rows misses it. ########## spark/src/main/scala/org/apache/spark/sql/comet/CometCollectMetricsExec.scala: ########## @@ -0,0 +1,170 @@ +/* + * 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.types.{CalendarIntervalType, DataType} +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) + + // Forwarded batches can carry widened nested nullability, as local table scans do. Preserve + // it at the next native scan instead of narrowing back to the original Catalyst field types. + override protected def scanFieldType(dt: DataType): DataType = dt.asNullable + + 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 = op match { + case collector: CometCollectMetricsExec => + // Replanning must retain the existing accumulator while using the same scan schema. + CometScanWrapper(nativeOp, collector) + case _ => + 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() + } + val inputAttributes = child.output + val requiresUnsafeRows = metricExpressions.exists(_.references.exists(attribute => + SupportLevel.containsType(attribute.dataType, classOf[CalendarIntervalType]))) + child.executeColumnar().mapPartitions { batches => + // Retaining aggregates copy their inputs, but ColumnarRow.copy cannot copy calendar + // intervals, including nested ones. Preserve the accumulator's input ordinals. + val toObservationRow: InternalRow => InternalRow = + if (requiresUnsafeRows) CometUnsafeProjection.create(inputAttributes) else identity + // 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(toObservationRow(rows.next())) Review Comment: **Observing `collect_set(to_json(v))` can abort the whole query.** With a native Parquet scan over `id` and `v: struct<x:int>`, I reproduced it on Spark 4.1.3 with: ```scala df.observe(obs, expr("collect_set(to_json(v))").as("m")) .select("id").repartition(2).collect() ``` For `v.x` values 1 and 2, the pre-PR base and this head with just the collector disabled both return the two distinct JSON strings. With the collector enabled, the task fails at result serialization with `had a not serializable result: org.apache.comet.vector.CometStructVector`. Plain `collect_set(v)` passes on the same input, and `mode(to_json(v), true)` also triggers the failure. Could we prevent these expressions from retaining a batch-backed value in the task accumulator, or fall back to Spark for them? The JSON metric above would cover the failure at task completion. -- 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]
