sunchao commented on code in PR #7117: URL: https://github.com/apache/datafusion-comet/pull/7117#discussion_r4367129463
########## spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala: ########## @@ -0,0 +1,236 @@ +/* + * 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.rules + +import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, AttributeSet, BinaryComparison, Cast, ConditionalExpression, CreateArray, CreateNamedStruct, DateAdd, DateDiff, DateFormatClass, DateSub, EqualNullSafe, Expression, IsNotNull, IsNull, MakeDate, Not, Or, SortOrder} +import org.apache.spark.sql.catalyst.trees.TreeNodeTag +import org.apache.spark.sql.comet._ +import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec +import org.apache.spark.sql.execution._ +import org.apache.spark.sql.execution.adaptive.{BroadcastQueryStageExec, QueryStageExec} +import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec} +import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec} +import org.apache.spark.sql.execution.joins.BaseJoinExec + +import org.apache.comet.CometConf +import org.apache.comet.shims.ShimCometWindowGroupLimit + +/** Keeps make_date on Spark's row path where native batching can change whether it throws. */ +private[rules] object MakeDateEvaluation { + private val sparkEvaluation = TreeNodeTag[String]("comet.makeDateSparkEvaluation") + + private def enabled(plan: SparkPlan): Boolean = + CometConf.COMET_EXEC_MAKE_DATE_OPERATOR_FALLBACK_ENABLED.get(plan.conf) + + def reason(plan: SparkPlan): Option[String] = + if (enabled(plan)) plan.getTagValue(sparkEvaluation) else None + + private def containsMakeDate(expressions: Seq[Expression]): Boolean = + expressions.exists(_.exists(_.isInstanceOf[MakeDate])) + + private case class ProjectionDemand(boundary: String, deferred: AttributeSet) + + /** References that can be evaluated before the expression without changing skipped inputs. */ + private def eagerReferences(expression: Expression): AttributeSet = expression match { + case attribute: Attribute => AttributeSet(Seq(attribute)) + case _ => + val children = expression match { + case conditional: ConditionalExpression => conditional.alwaysEvaluatedInputs + case _: Alias | _: Cast | _: IsNull | _: IsNotNull | _: Not => expression.children + case _: EqualNullSafe | _: CreateArray | _: CreateNamedStruct => expression.children + case _: And | _: Or => expression.children.take(1) + case _: BinaryComparison | _: DateAdd | _: DateSub | _: DateDiff | _: DateFormatClass => + expression.children.take(if (expression.children.head.nullable) 1 else 2) + // Unknown expressions may skip inputs or catch errors raised by their children. + case _ => Nil + } + AttributeSet(children.flatMap(eagerReferences)) + } + + private def copyLogicalLinks(from: SparkPlan, to: SparkPlan): Unit = { + if (from ne to) { + // setLogicalLink promotes inherited links to direct links and rewrites descendants. + // Keep each tag's role and retain the fallback plan's links when the source has none. + Seq(SparkPlan.LOGICAL_PLAN_TAG, SparkPlan.LOGICAL_PLAN_INHERITED_TAG).foreach { tag => + from.getTagValue(tag).foreach(to.setTagValue(tag, _)) + } + } + } + + private def nativeBroadcastInput(plan: SparkPlan): Option[SparkPlan] = plan match { + case _: CometBroadcastExchangeExec => Some(plan) + case stage: BroadcastQueryStageExec => nativeBroadcastInput(stage.plan).map(_ => stage) + case reuse: ReusedExchangeExec => nativeBroadcastInput(reuse.child).map(_ => reuse) + case sink: CometSinkPlaceHolder => nativeBroadcastInput(sink.child) + case wrapper: CometScanWrapper => nativeBroadcastInput(wrapper.originalPlan) + case _ => None + } + + private def restoreBroadcast(plan: SparkPlan, reason: String): SparkPlan = + nativeBroadcastInput(plan) + .map { input => + // Spark joins need a HashedRelation (or rows for a nested-loop join), not Arrow buffers. + // The existing transition decodes the cached broadcast and applies its Spark broadcast + // mode, including through AQE/reuse wrappers, without re-executing the build input. + val restored = CometColumnarToRowExec(input) + restored.setTagValue(sparkEvaluation, reason) + copyLogicalLinks(plan, restored) + restored + } + .getOrElse(plan) + + /** + * Run before bottom-up conversion and join rewriting. Reasons propagate back to the consuming + * operator, keeping the row pipeline together instead of inserting an eager Arrow batch between + * a throwing projection and its lazy consumer. Exchanges materialize their input independently; + * local sorts and aggregates do so only after their consumer requests input. + */ + def preserve(plan: SparkPlan): SparkPlan = { + if (!enabled(plan)) return plan + + def visit( + node: SparkPlan, + rowDemand: Option[String], + projectionDemand: Option[ProjectionDemand]): (SparkPlan, Option[String]) = { + node match { + // Materialized stages are independent producers. Their plans have already been guarded. + case _: QueryStageExec | _: ReusedExchangeExec => return (node, None) + // These wrappers do not have the child shape of their original Spark operator. + case wrapper: CometScanWrapper => + val (child, childReason) = visit(wrapper.originalPlan, rowDemand, projectionDemand) + return if (childReason.isDefined) (child, childReason) + else (wrapper.copy(originalPlan = child), None) + case sink: CometSinkPlaceHolder => + val (child, childReason) = visit(sink.child, rowDemand, projectionDemand) + return if (childReason.isDefined) (child, childReason) + else (sink.withNewChildren(Seq(child)), None) + case _ => + } + + val original = node match { + case scan: CometScanExec => + scan.wrapped.copy( + partitionFilters = scan.partitionFilters, + dataFilters = scan.dataFilters) + case scan: CometNativeScanExec => scan.sparkFallback(Nil) + case native: CometExec => native.originalPlan + case shuffle: CometShuffleExchangeExec => shuffle.originalPlan + case broadcast: CometBroadcastExchangeExec => broadcast.originalPlan + case _ => node + } + val exchange = original.isInstanceOf[Exchange] + val join = original.isInstanceOf[BaseJoinExec] + val localMaterialization = original match { + case _: SortExec | _: HashAggregateExec | _: ObjectHashAggregateExec => true + case _: TakeOrderedAndProjectExec => true + case _ => false + } + val startsLimit = original match { + case _: CollectLimitExec | _: LocalLimitExec | _: GlobalLimitExec => true + case topK: TakeOrderedAndProjectExec => + SortOrder.orderingSatisfies(topK.child.outputOrdering, topK.sortOrder) + case _ + if ShimCometWindowGroupLimit.windowGroupLimitClass.exists(_.isInstance(original)) => + true + case _ => false + } + val childDemand = original match { + case _ if exchange => None + // A join can leave a local sort or aggregate completely unpolled. + case _ if join || rowDemand.contains("JOIN") => Some("JOIN") + case _ if startsLimit => Some("LIMIT") + case _: SampleExec => Some("SAMPLE") + case _: MapPartitionsExec => Some("a mapPartitions function") + case _ if localMaterialization => None + case _ => rowDemand + } + val childProjectionDemand = original match { + // Generated filters can defer deterministic projection outputs until a row is accepted. + case filter: FilterExec => Some(ProjectionDemand("FILTER", filter.child.outputSet)) + case project: ProjectExec => + val deferred = projectionDemand.map(_.deferred).getOrElse(AttributeSet.empty) + val eager = project.projectList + .filter(output => !output.deterministic || !deferred.contains(output.toAttribute)) + .flatMap(eagerReferences) + // ProjectExec forces repeated inputs and nondeterministic outputs. Other inputs can + // remain inlined in a consumer's conditional, even across intervening aliases. + val skipped = project.child.outputSet -- project.usedInputs -- AttributeSet(eager) + if (skipped.isEmpty) { + None + } else { + Some(ProjectionDemand(projectionDemand.map(_.boundary).getOrElse("PROJECT"), skipped)) + } + case _ if exchange || localMaterialization || join => None + case _ => projectionDemand + } + val visited = node.children.map(visit(_, childDemand, childProjectionDemand)) + val childReason = visited.flatMap(_._2).headOption + val ownReason = reason(node).orElse { + if (containsMakeDate(original.expressions)) { + rowDemand + .map(boundary => s"make_date requires Spark evaluation below $boundary") + .orElse(original match { + case project: ProjectExec => + projectionDemand + .filter(demand => + project.projectList.exists(output => + output.deterministic && demand.deferred.contains(output.toAttribute) && + containsMakeDate(Seq(output)))) + .map(demand => s"make_date requires Spark evaluation below ${demand.boundary}") + // Join keys/conditions and aggregate updates can be skipped inside the operator. + case _: BaseJoinExec | _: BaseAggregateExec | _: TakeOrderedAndProjectExec => + Some(s"make_date requires Spark evaluation in ${original.nodeName}") + case _ => None + }) + } else None + } + val fallbackReason = ownReason.orElse(if (exchange) None else childReason) + val children = visited.map { case (child, _) => + if (join && fallbackReason.isDefined) restoreBroadcast(child, fallbackReason.get) + else child + } + val rebuilt = node match { + case _: RowToColumnarExec | _: CometSparkToColumnarExec if childReason.isDefined => + children.head + case _: ColumnarToRowExec | _: CometColumnarToRowExec | _: CometNativeColumnarToRowExec + if childReason.isDefined && !children.head.supportsColumnar => + children.head + case native: CometExec if fallbackReason.isDefined || children != node.children => + native.sparkFallback(children) + case _ if (original ne node) && (fallbackReason.isDefined || children != node.children) => + original.withNewChildren(children) + case _ => node.withNewChildren(children) + } + copyLogicalLinks(node, rebuilt) + fallbackReason.foreach(rebuilt.setTagValue(sparkEvaluation, _)) + (rebuilt, fallbackReason) + } + + // Dataset.rdd adds an object-producing root, and its caller may read only a few objects. + // EliminateSerialization can leave a typed operator, filter, or project at that root. + def producesObjects(op: SparkPlan): Boolean = op match { + case _: ObjectProducerExec => true + case _: FilterExec | _: ProjectExec => producesObjects(op.children.head) + case _ => false + } + val rootDemand = if (producesObjects(plan)) Some("code reading Dataset.rdd") else None Review Comment: [P2] Preserve the implicit one-row read for Spark 3.x `Dataset.isEmpty`. It calls `executeTake(1)` on `select().queryExecution`, so the zero-column root has no physical limit and receives no row demand here. A native filter can then fail on a row after Spark has already found its first match. Carry the action's bound into planning or conservatively guard this zero-column shape on the affected versions. Evidence: On Spark 3.5.9 with ANSI off, put years `2024, 2025, 6000000` in that order in one Parquet batch. `df.where("make_date(y,1,1) <=> DATE '2024-01-01'").isEmpty` returned `false` in Spark and on the base. The head's `CometFilter` raised `ARITHMETIC_OVERFLOW`. Disabling native filtering restored `false`. Spark 4.1.3 passed the same control with its explicit `limit(1)`. [Spark 3.5.9 implementation](https://github.com/apache/spark/blob/v3.5.9/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala#L653-L655) ########## spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala: ########## @@ -0,0 +1,236 @@ +/* + * 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.rules + +import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, AttributeSet, BinaryComparison, Cast, ConditionalExpression, CreateArray, CreateNamedStruct, DateAdd, DateDiff, DateFormatClass, DateSub, EqualNullSafe, Expression, IsNotNull, IsNull, MakeDate, Not, Or, SortOrder} +import org.apache.spark.sql.catalyst.trees.TreeNodeTag +import org.apache.spark.sql.comet._ +import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec +import org.apache.spark.sql.execution._ +import org.apache.spark.sql.execution.adaptive.{BroadcastQueryStageExec, QueryStageExec} +import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec} +import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec} +import org.apache.spark.sql.execution.joins.BaseJoinExec + +import org.apache.comet.CometConf +import org.apache.comet.shims.ShimCometWindowGroupLimit + +/** Keeps make_date on Spark's row path where native batching can change whether it throws. */ +private[rules] object MakeDateEvaluation { + private val sparkEvaluation = TreeNodeTag[String]("comet.makeDateSparkEvaluation") + + private def enabled(plan: SparkPlan): Boolean = + CometConf.COMET_EXEC_MAKE_DATE_OPERATOR_FALLBACK_ENABLED.get(plan.conf) + + def reason(plan: SparkPlan): Option[String] = + if (enabled(plan)) plan.getTagValue(sparkEvaluation) else None + + private def containsMakeDate(expressions: Seq[Expression]): Boolean = + expressions.exists(_.exists(_.isInstanceOf[MakeDate])) + + private case class ProjectionDemand(boundary: String, deferred: AttributeSet) + + /** References that can be evaluated before the expression without changing skipped inputs. */ + private def eagerReferences(expression: Expression): AttributeSet = expression match { + case attribute: Attribute => AttributeSet(Seq(attribute)) + case _ => + val children = expression match { + case conditional: ConditionalExpression => conditional.alwaysEvaluatedInputs + case _: Alias | _: Cast | _: IsNull | _: IsNotNull | _: Not => expression.children + case _: EqualNullSafe | _: CreateArray | _: CreateNamedStruct => expression.children + case _: And | _: Or => expression.children.take(1) + case _: BinaryComparison | _: DateAdd | _: DateSub | _: DateDiff | _: DateFormatClass => + expression.children.take(if (expression.children.head.nullable) 1 else 2) + // Unknown expressions may skip inputs or catch errors raised by their children. + case _ => Nil + } + AttributeSet(children.flatMap(eagerReferences)) + } + + private def copyLogicalLinks(from: SparkPlan, to: SparkPlan): Unit = { + if (from ne to) { + // setLogicalLink promotes inherited links to direct links and rewrites descendants. + // Keep each tag's role and retain the fallback plan's links when the source has none. + Seq(SparkPlan.LOGICAL_PLAN_TAG, SparkPlan.LOGICAL_PLAN_INHERITED_TAG).foreach { tag => + from.getTagValue(tag).foreach(to.setTagValue(tag, _)) + } + } + } + + private def nativeBroadcastInput(plan: SparkPlan): Option[SparkPlan] = plan match { + case _: CometBroadcastExchangeExec => Some(plan) + case stage: BroadcastQueryStageExec => nativeBroadcastInput(stage.plan).map(_ => stage) + case reuse: ReusedExchangeExec => nativeBroadcastInput(reuse.child).map(_ => reuse) + case sink: CometSinkPlaceHolder => nativeBroadcastInput(sink.child) + case wrapper: CometScanWrapper => nativeBroadcastInput(wrapper.originalPlan) + case _ => None + } + + private def restoreBroadcast(plan: SparkPlan, reason: String): SparkPlan = + nativeBroadcastInput(plan) + .map { input => + // Spark joins need a HashedRelation (or rows for a nested-loop join), not Arrow buffers. + // The existing transition decodes the cached broadcast and applies its Spark broadcast + // mode, including through AQE/reuse wrappers, without re-executing the build input. + val restored = CometColumnarToRowExec(input) + restored.setTagValue(sparkEvaluation, reason) + copyLogicalLinks(plan, restored) + restored + } + .getOrElse(plan) + + /** + * Run before bottom-up conversion and join rewriting. Reasons propagate back to the consuming + * operator, keeping the row pipeline together instead of inserting an eager Arrow batch between + * a throwing projection and its lazy consumer. Exchanges materialize their input independently; + * local sorts and aggregates do so only after their consumer requests input. + */ + def preserve(plan: SparkPlan): SparkPlan = { + if (!enabled(plan)) return plan + + def visit( + node: SparkPlan, + rowDemand: Option[String], + projectionDemand: Option[ProjectionDemand]): (SparkPlan, Option[String]) = { + node match { + // Materialized stages are independent producers. Their plans have already been guarded. + case _: QueryStageExec | _: ReusedExchangeExec => return (node, None) + // These wrappers do not have the child shape of their original Spark operator. + case wrapper: CometScanWrapper => + val (child, childReason) = visit(wrapper.originalPlan, rowDemand, projectionDemand) + return if (childReason.isDefined) (child, childReason) + else (wrapper.copy(originalPlan = child), None) + case sink: CometSinkPlaceHolder => + val (child, childReason) = visit(sink.child, rowDemand, projectionDemand) + return if (childReason.isDefined) (child, childReason) + else (sink.withNewChildren(Seq(child)), None) + case _ => + } + + val original = node match { + case scan: CometScanExec => + scan.wrapped.copy( + partitionFilters = scan.partitionFilters, + dataFilters = scan.dataFilters) + case scan: CometNativeScanExec => scan.sparkFallback(Nil) + case native: CometExec => native.originalPlan + case shuffle: CometShuffleExchangeExec => shuffle.originalPlan + case broadcast: CometBroadcastExchangeExec => broadcast.originalPlan + case _ => node + } + val exchange = original.isInstanceOf[Exchange] + val join = original.isInstanceOf[BaseJoinExec] + val localMaterialization = original match { + case _: SortExec | _: HashAggregateExec | _: ObjectHashAggregateExec => true + case _: TakeOrderedAndProjectExec => true + case _ => false + } + val startsLimit = original match { + case _: CollectLimitExec | _: LocalLimitExec | _: GlobalLimitExec => true + case topK: TakeOrderedAndProjectExec => + SortOrder.orderingSatisfies(topK.child.outputOrdering, topK.sortOrder) + case _ + if ShimCometWindowGroupLimit.windowGroupLimitClass.exists(_.isInstance(original)) => + true + case _ => false + } + val childDemand = original match { + case _ if exchange => None + // A join can leave a local sort or aggregate completely unpolled. + case _ if join || rowDemand.contains("JOIN") => Some("JOIN") + case _ if startsLimit => Some("LIMIT") + case _: SampleExec => Some("SAMPLE") + case _: MapPartitionsExec => Some("a mapPartitions function") + case _ if localMaterialization => None + case _ => rowDemand + } + val childProjectionDemand = original match { + // Generated filters can defer deterministic projection outputs until a row is accepted. + case filter: FilterExec => Some(ProjectionDemand("FILTER", filter.child.outputSet)) + case project: ProjectExec => + val deferred = projectionDemand.map(_.deferred).getOrElse(AttributeSet.empty) + val eager = project.projectList + .filter(output => !output.deterministic || !deferred.contains(output.toAttribute)) + .flatMap(eagerReferences) + // ProjectExec forces repeated inputs and nondeterministic outputs. Other inputs can + // remain inlined in a consumer's conditional, even across intervening aliases. + val skipped = project.child.outputSet -- project.usedInputs -- AttributeSet(eager) + if (skipped.isEmpty) { + None + } else { + Some(ProjectionDemand(projectionDemand.map(_.boundary).getOrElse("PROJECT"), skipped)) + } + case _ if exchange || localMaterialization || join => None + case _ => projectionDemand + } + val visited = node.children.map(visit(_, childDemand, childProjectionDemand)) + val childReason = visited.flatMap(_._2).headOption + val ownReason = reason(node).orElse { + if (containsMakeDate(original.expressions)) { + rowDemand + .map(boundary => s"make_date requires Spark evaluation below $boundary") + .orElse(original match { + case project: ProjectExec => + projectionDemand + .filter(demand => + project.projectList.exists(output => + output.deterministic && demand.deferred.contains(output.toAttribute) && + containsMakeDate(Seq(output)))) + .map(demand => s"make_date requires Spark evaluation below ${demand.boundary}") + // Join keys/conditions and aggregate updates can be skipped inside the operator. + case _: BaseJoinExec | _: BaseAggregateExec | _: TakeOrderedAndProjectExec => + Some(s"make_date requires Spark evaluation in ${original.nodeName}") + case _ => None + }) + } else None + } + val fallbackReason = ownReason.orElse(if (exchange) None else childReason) + val children = visited.map { case (child, _) => + if (join && fallbackReason.isDefined) restoreBroadcast(child, fallbackReason.get) + else child + } + val rebuilt = node match { + case _: RowToColumnarExec | _: CometSparkToColumnarExec if childReason.isDefined => + children.head + case _: ColumnarToRowExec | _: CometColumnarToRowExec | _: CometNativeColumnarToRowExec + if childReason.isDefined && !children.head.supportsColumnar => + children.head + case native: CometExec if fallbackReason.isDefined || children != node.children => + native.sparkFallback(children) + case _ if (original ne node) && (fallbackReason.isDefined || children != node.children) => + original.withNewChildren(children) + case _ => node.withNewChildren(children) + } + copyLogicalLinks(node, rebuilt) + fallbackReason.foreach(rebuilt.setTagValue(sparkEvaluation, _)) + (rebuilt, fallbackReason) + } + + // Dataset.rdd adds an object-producing root, and its caller may read only a few objects. + // EliminateSerialization can leave a typed operator, filter, or project at that root. + def producesObjects(op: SparkPlan): Boolean = op match { + case _: ObjectProducerExec => true Review Comment: [P2] Carry partial-read protection through classic PySpark's `DataFrame.rdd` export. `Dataset.javaToPython` consumes `queryExecution.toRdd` from an ordinary SQL projection, so there is no `ObjectProducerExec` for this check to recognize. A public `rdd.take(1)` can now raise on a date thousands of rows beyond the requested result. Preserve row-wise evaluation for a date projection exposed through this bridge. Evidence: On PySpark 4.2.0 with ANSI off, write one Parquet partition with `k = 0..4095` and `y = 2024` except for `y = 6000000` on the last row. Run `df.selectExpr("unix_date(make_date(y,1,1)) AS d", "concat(cast(k as string), repeat('a',32768)) AS payload").rdd.take(1)`. The wide payload lets the Python reader stop before the last row. Both head attempts raised `ARITHMETIC_OVERFLOW`. Spark, base Comet, and the head with native projection disabled each returned `d = 19723` twice. [Spark's Python export path](https://github.com/apache/spark/blob/v4.2.0/sql/core/src/main/scala/org/apache/spark/sql/classic/Dataset.scala#L2130-L2136) ########## spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala: ########## @@ -1041,15 +1044,16 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) plan } } else { + val planWithDateEvaluation = MakeDateEvaluation.preserve(plan) Review Comment: [P2] Preserve the scalar subquery's two-row bound when converting its child. Spark prepares the inner plan before wrapping it in `SubqueryExec(maxNumRows = Some(2))`. `preserve(plan)` sees an unrestricted projection, which can evaluate a third row and replace the cardinality violation with an unrelated overflow. Pass the bound into child planning or revisit the prepared child from the physical subquery expression. Evidence: On Spark 4.1.3 with ANSI off and years `2024, 2025, 6000000` in one Parquet batch, run `SELECT (SELECT unix_date(make_date(y,1,1)) FROM t)`. Spark and the base reported `SCALAR_SUBQUERY_TOO_MANY_ROWS` (`21000`). The head reported `ARITHMETIC_OVERFLOW` (`22003`) with AQE both on and off. Disabling native projection restored the cardinality error in both runs. [Spark's scalar-subquery bound](https://github.com/apache/spark/blob/v4.1.3/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala#L964-L968) ########## spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala: ########## @@ -0,0 +1,236 @@ +/* + * 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.rules + +import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, AttributeSet, BinaryComparison, Cast, ConditionalExpression, CreateArray, CreateNamedStruct, DateAdd, DateDiff, DateFormatClass, DateSub, EqualNullSafe, Expression, IsNotNull, IsNull, MakeDate, Not, Or, SortOrder} +import org.apache.spark.sql.catalyst.trees.TreeNodeTag +import org.apache.spark.sql.comet._ +import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec +import org.apache.spark.sql.execution._ +import org.apache.spark.sql.execution.adaptive.{BroadcastQueryStageExec, QueryStageExec} +import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec} +import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec} +import org.apache.spark.sql.execution.joins.BaseJoinExec + +import org.apache.comet.CometConf +import org.apache.comet.shims.ShimCometWindowGroupLimit + +/** Keeps make_date on Spark's row path where native batching can change whether it throws. */ +private[rules] object MakeDateEvaluation { + private val sparkEvaluation = TreeNodeTag[String]("comet.makeDateSparkEvaluation") + + private def enabled(plan: SparkPlan): Boolean = + CometConf.COMET_EXEC_MAKE_DATE_OPERATOR_FALLBACK_ENABLED.get(plan.conf) + + def reason(plan: SparkPlan): Option[String] = + if (enabled(plan)) plan.getTagValue(sparkEvaluation) else None + + private def containsMakeDate(expressions: Seq[Expression]): Boolean = + expressions.exists(_.exists(_.isInstanceOf[MakeDate])) + + private case class ProjectionDemand(boundary: String, deferred: AttributeSet) + + /** References that can be evaluated before the expression without changing skipped inputs. */ + private def eagerReferences(expression: Expression): AttributeSet = expression match { + case attribute: Attribute => AttributeSet(Seq(attribute)) + case _ => + val children = expression match { + case conditional: ConditionalExpression => conditional.alwaysEvaluatedInputs + case _: Alias | _: Cast | _: IsNull | _: IsNotNull | _: Not => expression.children + case _: EqualNullSafe | _: CreateArray | _: CreateNamedStruct => expression.children + case _: And | _: Or => expression.children.take(1) + case _: BinaryComparison | _: DateAdd | _: DateSub | _: DateDiff | _: DateFormatClass => + expression.children.take(if (expression.children.head.nullable) 1 else 2) + // Unknown expressions may skip inputs or catch errors raised by their children. + case _ => Nil Review Comment: [P2] `YEAR(dt)` reaches this default case, so the rule treats `dt` as deferred and moves two projections to Spark. In `SELECT r, YEAR(dt), r + 1 FROM (SELECT rand(7) AS r, make_date(y,m,d) AS dt FROM t)`, reusing `r` keeps the projections separate, and Spark's `Year` always evaluates its sole child. Follow the child of eager date extractors while preserving the mask from any enclosing conditional. Evidence: A grouped version on two million rows returned identical counts and sums. The head converted all input rows and allocated 65.0 MiB on the JVM per action. The base converted one result row and allocated 37.2 MiB. In the overflow control, Spark raised for direct `YEAR(dt)` and returned all 4,097 rows for `IF(r < 0, YEAR(dt), 1)`. Disabling operator fallback restored both native projections for the eager case. ########## spark/src/main/scala/org/apache/comet/serde/datetime.scala: ########## @@ -436,10 +436,70 @@ object CometNextDay extends CometExpressionSerde[NextDay] with CodegenDispatchFa object CometMakeDate extends CometExpressionSerde[MakeDate] { + private val sharedArgumentReason = + "make_date requires Spark evaluation when a computed argument after a nullable " + + "argument cannot preserve Spark's shared-subexpression evaluation." + + override def getUnsupportedReasons(): Seq[String] = Seq(sharedArgumentReason) + + override def getSupportLevel(expr: MakeDate): SupportLevel = { + val firstNullable = expr.children.indexWhere(_.nullable) + val hasComputedMaskedArgument = firstNullable >= 0 && + expr.children.drop(firstNullable + 1).exists(!_.isInstanceOf[LeafExpression]) + if (hasComputedMaskedArgument && !canShortCircuitNulls(expr)) { + Unsupported(Some(sharedArgumentReason)) + } else { + Compatible() + } + } + + /** A native parent must not evaluate a date on rows where Spark skips that child. */ + private[serde] def unsupportedParentReason(expr: Expression): Option[String] = { + var pending: List[(Expression, Option[String])] = (expr, None) :: Nil + while (pending.nonEmpty) { + val (current, parent) = pending.head + pending = pending.tail + if (current.isInstanceOf[MakeDate] && parent.isDefined) { + return Some(s"make_date requires Spark evaluation under ${parent.get}") + } + def pushChildren(mask: Option[String]): Unit = { + current.children.reverseIterator.foreach(child => pending = (child, mask) :: pending) + } + if (parent.isDefined) { + pushChildren(parent) + } else { + current match { + // These native expressions already select the rows passed to their children. + case _: If | _: CaseWhen | _: MakeDate => pushChildren(None) + // COALESCE lowers to CASE and evaluates an earlier child again in its result branch. + case coalesce: Coalesce if coalesce.deterministic => pushChildren(None) + // These parents evaluate every child in Spark, including NULL values. + case _: EqualNullSafe | _: CreateArray | _: CreateNamedStruct => + pushChildren(None) + // These parents evaluate left first and skip right only when left is NULL. + case _: BinaryComparison | _: DateAdd | _: DateSub | _: DateDiff | _: DateFormatClass => + val left = current.children.head + val right = current.children(1) + val rightMask = if (left.nullable) Some(current.nodeName) else None + pending = (left, None) :: (right, rightMask) :: pending + case _: SortOrder => pushChildren(Some(current.nodeName)) Review Comment: [P2] Allow a direct `make_date` as the leading `DateType` key of `SortExec`. Spark computes that key's sort prefix for every inserted row. Masking every `SortOrder` sends the whole sort to Spark, even with operator fallback disabled. Pass key position into the sort serialization guard and retain protection for secondary keys, which Spark can skip. Evidence: Sorting two million valid rows with `sortWithinPartitions(expr("make_date(2024,m,d)"), col("value"))` produced the same order and checksum. Median process CPU was 2.83 s on the head and 1.16 s on the base. Projecting the same date under an alias restored `CometSort` on the head and used 1.06 s. On 4,097 rows with one overflow, Spark raised with the date first and completed when the date followed a unique key. [Spark's leading-key prefix](https://github.com/apache/spark/blob/v4.1.3/sql/core/src/main/scala/org/apache/spark/sql/execution/SortExec.scala#L81-L98) ########## spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala: ########## @@ -0,0 +1,236 @@ +/* + * 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.rules + +import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, AttributeSet, BinaryComparison, Cast, ConditionalExpression, CreateArray, CreateNamedStruct, DateAdd, DateDiff, DateFormatClass, DateSub, EqualNullSafe, Expression, IsNotNull, IsNull, MakeDate, Not, Or, SortOrder} +import org.apache.spark.sql.catalyst.trees.TreeNodeTag +import org.apache.spark.sql.comet._ +import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec +import org.apache.spark.sql.execution._ +import org.apache.spark.sql.execution.adaptive.{BroadcastQueryStageExec, QueryStageExec} +import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec} +import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec} +import org.apache.spark.sql.execution.joins.BaseJoinExec + +import org.apache.comet.CometConf +import org.apache.comet.shims.ShimCometWindowGroupLimit + +/** Keeps make_date on Spark's row path where native batching can change whether it throws. */ +private[rules] object MakeDateEvaluation { + private val sparkEvaluation = TreeNodeTag[String]("comet.makeDateSparkEvaluation") + + private def enabled(plan: SparkPlan): Boolean = + CometConf.COMET_EXEC_MAKE_DATE_OPERATOR_FALLBACK_ENABLED.get(plan.conf) + + def reason(plan: SparkPlan): Option[String] = + if (enabled(plan)) plan.getTagValue(sparkEvaluation) else None + + private def containsMakeDate(expressions: Seq[Expression]): Boolean = + expressions.exists(_.exists(_.isInstanceOf[MakeDate])) + + private case class ProjectionDemand(boundary: String, deferred: AttributeSet) + + /** References that can be evaluated before the expression without changing skipped inputs. */ + private def eagerReferences(expression: Expression): AttributeSet = expression match { + case attribute: Attribute => AttributeSet(Seq(attribute)) + case _ => + val children = expression match { + case conditional: ConditionalExpression => conditional.alwaysEvaluatedInputs + case _: Alias | _: Cast | _: IsNull | _: IsNotNull | _: Not => expression.children + case _: EqualNullSafe | _: CreateArray | _: CreateNamedStruct => expression.children + case _: And | _: Or => expression.children.take(1) + case _: BinaryComparison | _: DateAdd | _: DateSub | _: DateDiff | _: DateFormatClass => + expression.children.take(if (expression.children.head.nullable) 1 else 2) + // Unknown expressions may skip inputs or catch errors raised by their children. + case _ => Nil + } + AttributeSet(children.flatMap(eagerReferences)) + } + + private def copyLogicalLinks(from: SparkPlan, to: SparkPlan): Unit = { + if (from ne to) { + // setLogicalLink promotes inherited links to direct links and rewrites descendants. + // Keep each tag's role and retain the fallback plan's links when the source has none. + Seq(SparkPlan.LOGICAL_PLAN_TAG, SparkPlan.LOGICAL_PLAN_INHERITED_TAG).foreach { tag => + from.getTagValue(tag).foreach(to.setTagValue(tag, _)) + } + } + } + + private def nativeBroadcastInput(plan: SparkPlan): Option[SparkPlan] = plan match { + case _: CometBroadcastExchangeExec => Some(plan) + case stage: BroadcastQueryStageExec => nativeBroadcastInput(stage.plan).map(_ => stage) + case reuse: ReusedExchangeExec => nativeBroadcastInput(reuse.child).map(_ => reuse) + case sink: CometSinkPlaceHolder => nativeBroadcastInput(sink.child) + case wrapper: CometScanWrapper => nativeBroadcastInput(wrapper.originalPlan) + case _ => None + } + + private def restoreBroadcast(plan: SparkPlan, reason: String): SparkPlan = + nativeBroadcastInput(plan) + .map { input => + // Spark joins need a HashedRelation (or rows for a nested-loop join), not Arrow buffers. + // The existing transition decodes the cached broadcast and applies its Spark broadcast + // mode, including through AQE/reuse wrappers, without re-executing the build input. + val restored = CometColumnarToRowExec(input) + restored.setTagValue(sparkEvaluation, reason) + copyLogicalLinks(plan, restored) + restored + } + .getOrElse(plan) + + /** + * Run before bottom-up conversion and join rewriting. Reasons propagate back to the consuming + * operator, keeping the row pipeline together instead of inserting an eager Arrow batch between + * a throwing projection and its lazy consumer. Exchanges materialize their input independently; + * local sorts and aggregates do so only after their consumer requests input. + */ + def preserve(plan: SparkPlan): SparkPlan = { + if (!enabled(plan)) return plan + + def visit( + node: SparkPlan, + rowDemand: Option[String], + projectionDemand: Option[ProjectionDemand]): (SparkPlan, Option[String]) = { + node match { + // Materialized stages are independent producers. Their plans have already been guarded. + case _: QueryStageExec | _: ReusedExchangeExec => return (node, None) + // These wrappers do not have the child shape of their original Spark operator. + case wrapper: CometScanWrapper => + val (child, childReason) = visit(wrapper.originalPlan, rowDemand, projectionDemand) + return if (childReason.isDefined) (child, childReason) + else (wrapper.copy(originalPlan = child), None) + case sink: CometSinkPlaceHolder => + val (child, childReason) = visit(sink.child, rowDemand, projectionDemand) + return if (childReason.isDefined) (child, childReason) + else (sink.withNewChildren(Seq(child)), None) + case _ => + } + + val original = node match { + case scan: CometScanExec => + scan.wrapped.copy( + partitionFilters = scan.partitionFilters, + dataFilters = scan.dataFilters) + case scan: CometNativeScanExec => scan.sparkFallback(Nil) + case native: CometExec => native.originalPlan + case shuffle: CometShuffleExchangeExec => shuffle.originalPlan + case broadcast: CometBroadcastExchangeExec => broadcast.originalPlan + case _ => node + } + val exchange = original.isInstanceOf[Exchange] + val join = original.isInstanceOf[BaseJoinExec] + val localMaterialization = original match { + case _: SortExec | _: HashAggregateExec | _: ObjectHashAggregateExec => true + case _: TakeOrderedAndProjectExec => true + case _ => false + } + val startsLimit = original match { + case _: CollectLimitExec | _: LocalLimitExec | _: GlobalLimitExec => true + case topK: TakeOrderedAndProjectExec => + SortOrder.orderingSatisfies(topK.child.outputOrdering, topK.sortOrder) + case _ + if ShimCometWindowGroupLimit.windowGroupLimitClass.exists(_.isInstance(original)) => + true + case _ => false + } + val childDemand = original match { + case _ if exchange => None + // A join can leave a local sort or aggregate completely unpolled. + case _ if join || rowDemand.contains("JOIN") => Some("JOIN") + case _ if startsLimit => Some("LIMIT") + case _: SampleExec => Some("SAMPLE") + case _: MapPartitionsExec => Some("a mapPartitions function") + case _ if localMaterialization => None + case _ => rowDemand + } + val childProjectionDemand = original match { + // Generated filters can defer deterministic projection outputs until a row is accepted. + case filter: FilterExec => Some(ProjectionDemand("FILTER", filter.child.outputSet)) + case project: ProjectExec => + val deferred = projectionDemand.map(_.deferred).getOrElse(AttributeSet.empty) + val eager = project.projectList + .filter(output => !output.deterministic || !deferred.contains(output.toAttribute)) + .flatMap(eagerReferences) + // ProjectExec forces repeated inputs and nondeterministic outputs. Other inputs can + // remain inlined in a consumer's conditional, even across intervening aliases. + val skipped = project.child.outputSet -- project.usedInputs -- AttributeSet(eager) + if (skipped.isEmpty) { + None + } else { + Some(ProjectionDemand(projectionDemand.map(_.boundary).getOrElse("PROJECT"), skipped)) + } + case _ if exchange || localMaterialization || join => None + case _ => projectionDemand + } + val visited = node.children.map(visit(_, childDemand, childProjectionDemand)) Review Comment: [P2] Translate deferred union outputs to each child's attributes by ordinal before calling `visit`. Spark takes the union output IDs from the first child. Passing the same `AttributeSet` to every branch leaves a date projection in the second branch native, even when the parent skips that date. Reversing the branches changes whether the query succeeds. Evidence: On Spark 4.2.0 with ANSI and AQE off, enable `spark.sql.codegen.wholeStage.union.enabled` and disable `spark.sql.unionOutputPartitioning`. For Parquet `t(y)` containing `2024, 2025, 6000000`, run `SELECT r, IF(r < 0, dt, DATE '2000-01-01'), rand(2L) FROM (SELECT y AS r, DATE '2000-01-01' AS dt FROM t UNION ALL SELECT y AS r, make_date(y,1,1) AS dt FROM t)`. Spark and the base returned six rows with the selected constant date. The head raised `ARITHMETIC_OVERFLOW`. Moving the date-producing branch first restored all six rows. [Spark's per-child union mapping](https://github.com/apache/spark/blob/v4.2.0/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala#L959-L975) ########## spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala: ########## @@ -0,0 +1,236 @@ +/* + * 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.rules + +import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, AttributeSet, BinaryComparison, Cast, ConditionalExpression, CreateArray, CreateNamedStruct, DateAdd, DateDiff, DateFormatClass, DateSub, EqualNullSafe, Expression, IsNotNull, IsNull, MakeDate, Not, Or, SortOrder} +import org.apache.spark.sql.catalyst.trees.TreeNodeTag +import org.apache.spark.sql.comet._ +import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec +import org.apache.spark.sql.execution._ +import org.apache.spark.sql.execution.adaptive.{BroadcastQueryStageExec, QueryStageExec} +import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec} +import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec} +import org.apache.spark.sql.execution.joins.BaseJoinExec + +import org.apache.comet.CometConf +import org.apache.comet.shims.ShimCometWindowGroupLimit + +/** Keeps make_date on Spark's row path where native batching can change whether it throws. */ +private[rules] object MakeDateEvaluation { + private val sparkEvaluation = TreeNodeTag[String]("comet.makeDateSparkEvaluation") + + private def enabled(plan: SparkPlan): Boolean = + CometConf.COMET_EXEC_MAKE_DATE_OPERATOR_FALLBACK_ENABLED.get(plan.conf) + + def reason(plan: SparkPlan): Option[String] = + if (enabled(plan)) plan.getTagValue(sparkEvaluation) else None + + private def containsMakeDate(expressions: Seq[Expression]): Boolean = + expressions.exists(_.exists(_.isInstanceOf[MakeDate])) + + private case class ProjectionDemand(boundary: String, deferred: AttributeSet) + + /** References that can be evaluated before the expression without changing skipped inputs. */ + private def eagerReferences(expression: Expression): AttributeSet = expression match { + case attribute: Attribute => AttributeSet(Seq(attribute)) + case _ => + val children = expression match { + case conditional: ConditionalExpression => conditional.alwaysEvaluatedInputs + case _: Alias | _: Cast | _: IsNull | _: IsNotNull | _: Not => expression.children + case _: EqualNullSafe | _: CreateArray | _: CreateNamedStruct => expression.children + case _: And | _: Or => expression.children.take(1) + case _: BinaryComparison | _: DateAdd | _: DateSub | _: DateDiff | _: DateFormatClass => + expression.children.take(if (expression.children.head.nullable) 1 else 2) + // Unknown expressions may skip inputs or catch errors raised by their children. + case _ => Nil + } + AttributeSet(children.flatMap(eagerReferences)) Review Comment: [P2] Collect eager references once instead of rebuilding an `AttributeSet` at every enclosing expression. A projection over a leaf pays this cost for every field of a nested sibling struct, even though the leaf cannot supply a deferred date projection. A plan-wide check for the absence of `make_date` would still enter this path. Evidence: Apply `CometExecRule` to a `ProjectExec` over an `ExistingRDD` of integer columns, projecting `make_date(2024,c0,c1)` and a struct nested by wrapping the previous struct with the next group of fields. At 1,024 fields and 32 levels, median driver-thread allocation was 3.82 MiB on the head and 0.27 MiB on the base. At 4,096 fields and 64 levels it was 21.93 MiB and 21.8 ms on the head versus 1.03 MiB and 1.6 ms on the base. The resulting physical plan was the same. ########## spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala: ########## @@ -0,0 +1,236 @@ +/* + * 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.rules + +import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, AttributeSet, BinaryComparison, Cast, ConditionalExpression, CreateArray, CreateNamedStruct, DateAdd, DateDiff, DateFormatClass, DateSub, EqualNullSafe, Expression, IsNotNull, IsNull, MakeDate, Not, Or, SortOrder} +import org.apache.spark.sql.catalyst.trees.TreeNodeTag +import org.apache.spark.sql.comet._ +import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec +import org.apache.spark.sql.execution._ +import org.apache.spark.sql.execution.adaptive.{BroadcastQueryStageExec, QueryStageExec} +import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec} +import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec} +import org.apache.spark.sql.execution.joins.BaseJoinExec + +import org.apache.comet.CometConf +import org.apache.comet.shims.ShimCometWindowGroupLimit + +/** Keeps make_date on Spark's row path where native batching can change whether it throws. */ +private[rules] object MakeDateEvaluation { + private val sparkEvaluation = TreeNodeTag[String]("comet.makeDateSparkEvaluation") + + private def enabled(plan: SparkPlan): Boolean = + CometConf.COMET_EXEC_MAKE_DATE_OPERATOR_FALLBACK_ENABLED.get(plan.conf) + + def reason(plan: SparkPlan): Option[String] = + if (enabled(plan)) plan.getTagValue(sparkEvaluation) else None + + private def containsMakeDate(expressions: Seq[Expression]): Boolean = + expressions.exists(_.exists(_.isInstanceOf[MakeDate])) + + private case class ProjectionDemand(boundary: String, deferred: AttributeSet) + + /** References that can be evaluated before the expression without changing skipped inputs. */ + private def eagerReferences(expression: Expression): AttributeSet = expression match { + case attribute: Attribute => AttributeSet(Seq(attribute)) + case _ => + val children = expression match { + case conditional: ConditionalExpression => conditional.alwaysEvaluatedInputs + case _: Alias | _: Cast | _: IsNull | _: IsNotNull | _: Not => expression.children + case _: EqualNullSafe | _: CreateArray | _: CreateNamedStruct => expression.children + case _: And | _: Or => expression.children.take(1) + case _: BinaryComparison | _: DateAdd | _: DateSub | _: DateDiff | _: DateFormatClass => + expression.children.take(if (expression.children.head.nullable) 1 else 2) + // Unknown expressions may skip inputs or catch errors raised by their children. + case _ => Nil + } + AttributeSet(children.flatMap(eagerReferences)) + } + + private def copyLogicalLinks(from: SparkPlan, to: SparkPlan): Unit = { + if (from ne to) { + // setLogicalLink promotes inherited links to direct links and rewrites descendants. + // Keep each tag's role and retain the fallback plan's links when the source has none. + Seq(SparkPlan.LOGICAL_PLAN_TAG, SparkPlan.LOGICAL_PLAN_INHERITED_TAG).foreach { tag => + from.getTagValue(tag).foreach(to.setTagValue(tag, _)) + } + } + } + + private def nativeBroadcastInput(plan: SparkPlan): Option[SparkPlan] = plan match { + case _: CometBroadcastExchangeExec => Some(plan) + case stage: BroadcastQueryStageExec => nativeBroadcastInput(stage.plan).map(_ => stage) + case reuse: ReusedExchangeExec => nativeBroadcastInput(reuse.child).map(_ => reuse) + case sink: CometSinkPlaceHolder => nativeBroadcastInput(sink.child) + case wrapper: CometScanWrapper => nativeBroadcastInput(wrapper.originalPlan) + case _ => None + } + + private def restoreBroadcast(plan: SparkPlan, reason: String): SparkPlan = + nativeBroadcastInput(plan) + .map { input => + // Spark joins need a HashedRelation (or rows for a nested-loop join), not Arrow buffers. + // The existing transition decodes the cached broadcast and applies its Spark broadcast + // mode, including through AQE/reuse wrappers, without re-executing the build input. + val restored = CometColumnarToRowExec(input) + restored.setTagValue(sparkEvaluation, reason) + copyLogicalLinks(plan, restored) + restored + } + .getOrElse(plan) + + /** + * Run before bottom-up conversion and join rewriting. Reasons propagate back to the consuming + * operator, keeping the row pipeline together instead of inserting an eager Arrow batch between + * a throwing projection and its lazy consumer. Exchanges materialize their input independently; + * local sorts and aggregates do so only after their consumer requests input. + */ + def preserve(plan: SparkPlan): SparkPlan = { + if (!enabled(plan)) return plan + + def visit( + node: SparkPlan, + rowDemand: Option[String], + projectionDemand: Option[ProjectionDemand]): (SparkPlan, Option[String]) = { + node match { + // Materialized stages are independent producers. Their plans have already been guarded. + case _: QueryStageExec | _: ReusedExchangeExec => return (node, None) + // These wrappers do not have the child shape of their original Spark operator. + case wrapper: CometScanWrapper => + val (child, childReason) = visit(wrapper.originalPlan, rowDemand, projectionDemand) + return if (childReason.isDefined) (child, childReason) + else (wrapper.copy(originalPlan = child), None) + case sink: CometSinkPlaceHolder => + val (child, childReason) = visit(sink.child, rowDemand, projectionDemand) + return if (childReason.isDefined) (child, childReason) + else (sink.withNewChildren(Seq(child)), None) + case _ => + } + + val original = node match { + case scan: CometScanExec => + scan.wrapped.copy( + partitionFilters = scan.partitionFilters, + dataFilters = scan.dataFilters) + case scan: CometNativeScanExec => scan.sparkFallback(Nil) + case native: CometExec => native.originalPlan + case shuffle: CometShuffleExchangeExec => shuffle.originalPlan + case broadcast: CometBroadcastExchangeExec => broadcast.originalPlan + case _ => node + } + val exchange = original.isInstanceOf[Exchange] + val join = original.isInstanceOf[BaseJoinExec] + val localMaterialization = original match { + case _: SortExec | _: HashAggregateExec | _: ObjectHashAggregateExec => true + case _: TakeOrderedAndProjectExec => true + case _ => false + } + val startsLimit = original match { + case _: CollectLimitExec | _: LocalLimitExec | _: GlobalLimitExec => true + case topK: TakeOrderedAndProjectExec => + SortOrder.orderingSatisfies(topK.child.outputOrdering, topK.sortOrder) + case _ + if ShimCometWindowGroupLimit.windowGroupLimitClass.exists(_.isInstance(original)) => + true + case _ => false + } + val childDemand = original match { + case _ if exchange => None + // A join can leave a local sort or aggregate completely unpolled. + case _ if join || rowDemand.contains("JOIN") => Some("JOIN") + case _ if startsLimit => Some("LIMIT") + case _: SampleExec => Some("SAMPLE") + case _: MapPartitionsExec => Some("a mapPartitions function") + case _ if localMaterialization => None + case _ => rowDemand + } + val childProjectionDemand = original match { + // Generated filters can defer deterministic projection outputs until a row is accepted. + case filter: FilterExec => Some(ProjectionDemand("FILTER", filter.child.outputSet)) + case project: ProjectExec => + val deferred = projectionDemand.map(_.deferred).getOrElse(AttributeSet.empty) + val eager = project.projectList + .filter(output => !output.deterministic || !deferred.contains(output.toAttribute)) + .flatMap(eagerReferences) + // ProjectExec forces repeated inputs and nondeterministic outputs. Other inputs can + // remain inlined in a consumer's conditional, even across intervening aliases. + val skipped = project.child.outputSet -- project.usedInputs -- AttributeSet(eager) + if (skipped.isEmpty) { + None + } else { + Some(ProjectionDemand(projectionDemand.map(_.boundary).getOrElse("PROJECT"), skipped)) + } + case _ if exchange || localMaterialization || join => None + case _ => projectionDemand + } + val visited = node.children.map(visit(_, childDemand, childProjectionDemand)) + val childReason = visited.flatMap(_._2).headOption + val ownReason = reason(node).orElse { + if (containsMakeDate(original.expressions)) { + rowDemand + .map(boundary => s"make_date requires Spark evaluation below $boundary") + .orElse(original match { + case project: ProjectExec => + projectionDemand + .filter(demand => + project.projectList.exists(output => + output.deterministic && demand.deferred.contains(output.toAttribute) && + containsMakeDate(Seq(output)))) + .map(demand => s"make_date requires Spark evaluation below ${demand.boundary}") + // Join keys/conditions and aggregate updates can be skipped inside the operator. + case _: BaseJoinExec | _: BaseAggregateExec | _: TakeOrderedAndProjectExec => Review Comment: [P2] Keep an unfiltered, single-argument `COUNT(make_date(y,m,d))` eligible for native aggregation when the date is nullable. Spark's `Count.updateExpressions` checks that argument for NULL on every input row. The blanket `BaseAggregateExec` fallback moves this eager update to Spark. Retain the guard for filters or other arguments that can skip the date. Evidence: Grouping two million Parquet rows by `y,m,d` and summing `COUNT(make_date(cast(y AS int),m,d))` returned two million across 262,144 groups in every mode. The head converted two million input rows and allocated 80.2 MiB per action. The base converted one result row and allocated 32.9 MiB. With an overflow row, Spark raised for the unfiltered count and returned 4,096 with `FILTER (WHERE y < 3000)`. [Spark's count update](https://github.com/apache/spark/blob/v4.1.3/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Count.scala#L89-L98) ########## spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala: ########## @@ -0,0 +1,236 @@ +/* + * 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.rules + +import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, AttributeSet, BinaryComparison, Cast, ConditionalExpression, CreateArray, CreateNamedStruct, DateAdd, DateDiff, DateFormatClass, DateSub, EqualNullSafe, Expression, IsNotNull, IsNull, MakeDate, Not, Or, SortOrder} +import org.apache.spark.sql.catalyst.trees.TreeNodeTag +import org.apache.spark.sql.comet._ +import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec +import org.apache.spark.sql.execution._ +import org.apache.spark.sql.execution.adaptive.{BroadcastQueryStageExec, QueryStageExec} +import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec} +import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec} +import org.apache.spark.sql.execution.joins.BaseJoinExec + +import org.apache.comet.CometConf +import org.apache.comet.shims.ShimCometWindowGroupLimit + +/** Keeps make_date on Spark's row path where native batching can change whether it throws. */ +private[rules] object MakeDateEvaluation { + private val sparkEvaluation = TreeNodeTag[String]("comet.makeDateSparkEvaluation") + + private def enabled(plan: SparkPlan): Boolean = + CometConf.COMET_EXEC_MAKE_DATE_OPERATOR_FALLBACK_ENABLED.get(plan.conf) + + def reason(plan: SparkPlan): Option[String] = + if (enabled(plan)) plan.getTagValue(sparkEvaluation) else None + + private def containsMakeDate(expressions: Seq[Expression]): Boolean = + expressions.exists(_.exists(_.isInstanceOf[MakeDate])) + + private case class ProjectionDemand(boundary: String, deferred: AttributeSet) + + /** References that can be evaluated before the expression without changing skipped inputs. */ + private def eagerReferences(expression: Expression): AttributeSet = expression match { + case attribute: Attribute => AttributeSet(Seq(attribute)) + case _ => + val children = expression match { + case conditional: ConditionalExpression => conditional.alwaysEvaluatedInputs + case _: Alias | _: Cast | _: IsNull | _: IsNotNull | _: Not => expression.children + case _: EqualNullSafe | _: CreateArray | _: CreateNamedStruct => expression.children + case _: And | _: Or => expression.children.take(1) + case _: BinaryComparison | _: DateAdd | _: DateSub | _: DateDiff | _: DateFormatClass => + expression.children.take(if (expression.children.head.nullable) 1 else 2) + // Unknown expressions may skip inputs or catch errors raised by their children. + case _ => Nil + } + AttributeSet(children.flatMap(eagerReferences)) + } + + private def copyLogicalLinks(from: SparkPlan, to: SparkPlan): Unit = { + if (from ne to) { + // setLogicalLink promotes inherited links to direct links and rewrites descendants. + // Keep each tag's role and retain the fallback plan's links when the source has none. + Seq(SparkPlan.LOGICAL_PLAN_TAG, SparkPlan.LOGICAL_PLAN_INHERITED_TAG).foreach { tag => + from.getTagValue(tag).foreach(to.setTagValue(tag, _)) + } + } + } + + private def nativeBroadcastInput(plan: SparkPlan): Option[SparkPlan] = plan match { + case _: CometBroadcastExchangeExec => Some(plan) + case stage: BroadcastQueryStageExec => nativeBroadcastInput(stage.plan).map(_ => stage) + case reuse: ReusedExchangeExec => nativeBroadcastInput(reuse.child).map(_ => reuse) + case sink: CometSinkPlaceHolder => nativeBroadcastInput(sink.child) + case wrapper: CometScanWrapper => nativeBroadcastInput(wrapper.originalPlan) + case _ => None + } + + private def restoreBroadcast(plan: SparkPlan, reason: String): SparkPlan = + nativeBroadcastInput(plan) + .map { input => + // Spark joins need a HashedRelation (or rows for a nested-loop join), not Arrow buffers. + // The existing transition decodes the cached broadcast and applies its Spark broadcast + // mode, including through AQE/reuse wrappers, without re-executing the build input. + val restored = CometColumnarToRowExec(input) + restored.setTagValue(sparkEvaluation, reason) + copyLogicalLinks(plan, restored) + restored + } + .getOrElse(plan) + + /** + * Run before bottom-up conversion and join rewriting. Reasons propagate back to the consuming + * operator, keeping the row pipeline together instead of inserting an eager Arrow batch between + * a throwing projection and its lazy consumer. Exchanges materialize their input independently; + * local sorts and aggregates do so only after their consumer requests input. + */ + def preserve(plan: SparkPlan): SparkPlan = { + if (!enabled(plan)) return plan + + def visit( + node: SparkPlan, + rowDemand: Option[String], + projectionDemand: Option[ProjectionDemand]): (SparkPlan, Option[String]) = { + node match { + // Materialized stages are independent producers. Their plans have already been guarded. + case _: QueryStageExec | _: ReusedExchangeExec => return (node, None) + // These wrappers do not have the child shape of their original Spark operator. + case wrapper: CometScanWrapper => + val (child, childReason) = visit(wrapper.originalPlan, rowDemand, projectionDemand) + return if (childReason.isDefined) (child, childReason) + else (wrapper.copy(originalPlan = child), None) + case sink: CometSinkPlaceHolder => + val (child, childReason) = visit(sink.child, rowDemand, projectionDemand) + return if (childReason.isDefined) (child, childReason) + else (sink.withNewChildren(Seq(child)), None) + case _ => + } + + val original = node match { + case scan: CometScanExec => + scan.wrapped.copy( + partitionFilters = scan.partitionFilters, + dataFilters = scan.dataFilters) + case scan: CometNativeScanExec => scan.sparkFallback(Nil) + case native: CometExec => native.originalPlan + case shuffle: CometShuffleExchangeExec => shuffle.originalPlan + case broadcast: CometBroadcastExchangeExec => broadcast.originalPlan + case _ => node + } + val exchange = original.isInstanceOf[Exchange] + val join = original.isInstanceOf[BaseJoinExec] + val localMaterialization = original match { + case _: SortExec | _: HashAggregateExec | _: ObjectHashAggregateExec => true + case _: TakeOrderedAndProjectExec => true + case _ => false + } + val startsLimit = original match { + case _: CollectLimitExec | _: LocalLimitExec | _: GlobalLimitExec => true + case topK: TakeOrderedAndProjectExec => + SortOrder.orderingSatisfies(topK.child.outputOrdering, topK.sortOrder) + case _ + if ShimCometWindowGroupLimit.windowGroupLimitClass.exists(_.isInstance(original)) => + true + case _ => false + } + val childDemand = original match { + case _ if exchange => None + // A join can leave a local sort or aggregate completely unpolled. + case _ if join || rowDemand.contains("JOIN") => Some("JOIN") + case _ if startsLimit => Some("LIMIT") + case _: SampleExec => Some("SAMPLE") + case _: MapPartitionsExec => Some("a mapPartitions function") + case _ if localMaterialization => None + case _ => rowDemand + } + val childProjectionDemand = original match { + // Generated filters can defer deterministic projection outputs until a row is accepted. + case filter: FilterExec => Some(ProjectionDemand("FILTER", filter.child.outputSet)) + case project: ProjectExec => + val deferred = projectionDemand.map(_.deferred).getOrElse(AttributeSet.empty) + val eager = project.projectList + .filter(output => !output.deterministic || !deferred.contains(output.toAttribute)) + .flatMap(eagerReferences) + // ProjectExec forces repeated inputs and nondeterministic outputs. Other inputs can + // remain inlined in a consumer's conditional, even across intervening aliases. + val skipped = project.child.outputSet -- project.usedInputs -- AttributeSet(eager) + if (skipped.isEmpty) { + None + } else { + Some(ProjectionDemand(projectionDemand.map(_.boundary).getOrElse("PROJECT"), skipped)) + } + case _ if exchange || localMaterialization || join => None + case _ => projectionDemand + } + val visited = node.children.map(visit(_, childDemand, childProjectionDemand)) + val childReason = visited.flatMap(_._2).headOption + val ownReason = reason(node).orElse { + if (containsMakeDate(original.expressions)) { + rowDemand Review Comment: [P2] Clear an outer limit's row demand for a hash exchange's own partition key as well as its child. Spark must evaluate that key for every shuffle input before the limit reads the result. Applying `rowDemand` here sends a previously native hash shuffle through JVM rows without protecting any skipped date. Evidence: On two million valid Parquet rows, `df.repartition(8, expr("make_date(cast(y AS int),m,d)")).limit(10).collect()` returned ten rows in every mode. The head converted all two million rows for Spark `Exchange` and allocated 296.5 MiB on the JVM per action. The base used native shuffle and allocated 17.4 MiB. The head with operator fallback disabled allocated 17.1 MiB. With overflow on the last input row, Spark and both head configurations all raised despite `LIMIT 10`. -- 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]
