sunchao commented on code in PR #7155:
URL: https://github.com/apache/datafusion-comet/pull/7155#discussion_r4367710634
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1178,10 +1178,96 @@ case class CometExecRule(session: SparkSession,
queryStagePrep: Boolean = false)
serde.convert(op, builder)
}
}
- return nativeOp.map { nativeOp =>
- val exec = serde.createExec(nativeOp, op)
- rollUpInfoMessages(op, exec)
- exec
+ return nativeOp
+ .filter { proto =>
+ // A serde may emit several native nodes for one Spark operator.
Inspect those nodes,
+ // but stop at real Spark children, whose consumption context was
tagged independently.
+ def containsOwnAssertion(node: Operator): Boolean =
+
CometAssertNotNull.containsNativeAssertion(node.toBuilder.clearChildren().build())
||
+ node.getChildrenList.asScala.exists { child =>
+ child.getPlanId == proto.getPlanId &&
containsOwnAssertion(child)
+ }
+
+ def hasNativeAssertion: Boolean = op match {
+ // TopK's sink proto omits its runtime projection over the
selected rows.
+ case topK: TakeOrderedAndProjectExec =>
+ QueryPlanSerde.withSharedSubexpressions(topK.projectList) {
+ topK.projectList.exists { expression =>
+ QueryPlanSerde
+ .exprToProto(expression, topK.child.output)
+ .exists(CometAssertNotNull.containsNativeAssertion)
+ }
+ }
+ case _: HashAggregateExec | _: ObjectHashAggregateExec =>
+ // Input/grouping expressions consume all rows; only the result
projection is lazy.
+ proto.hasProjection &&
CometAssertNotNull.containsNativeAssertion(
+ proto.getProjection)
+ case _ => containsOwnAssertion(proto)
+ }
+ val fallback = nativeAssertionJoinFallback(proto).orElse {
+ op.getTagValue(CometExecRule.PARTIAL_READER)
+ .filter(_ => hasNativeAssertion)
+ .map(reader =>
+ s"Native AssertNotNull below ${reader.description} could
change which values " +
+ "Spark checks")
+ }
+ fallback.foreach(reason => withFallbackReason(op, reason))
+ fallback.isEmpty
+ }
+ .map { nativeOp =>
+ val exec = serde.createExec(nativeOp, op)
+ rollUpInfoMessages(op, exec)
+ exec
+ }
+ }
+ None
+ }
+
+ private def nativeAssertionJoinFallback(proto: Operator): Option[String] = {
+ import OperatorOuterClass.JoinType
+
+ // These consumers can skip candidates even without a LIMIT above the
join. Inspect only
+ // the serialized condition. Child projections use the partial-reader tags
independently;
+ // assertions captured inside opaque JVM dispatch keep their existing
execution path.
+ val condition = if (proto.hasHashJoin) {
+ val join = proto.getHashJoin
+ Some(join.getJoinType -> join.getCondition)
+ } else if (proto.hasSortMergeJoin) {
+ val join = proto.getSortMergeJoin
+ Some(join.getJoinType -> join.getCondition)
+ } else if (proto.hasBroadcastNestedLoopJoin) {
+ val join = proto.getBroadcastNestedLoopJoin
+ Some(join.getJoinType -> join.getCondition)
+ } else None
+
+ if (condition.exists {
+ case (JoinType.LeftSemi | JoinType.LeftAnti, expression) =>
+ CometAssertNotNull.containsNativeAssertion(expression)
+ case _ => false
+ }) {
+ return Some(
+ "Native AssertNotNull in an early-exit join condition could check
candidates Spark skips")
+ }
+
+ if (proto.hasHashJoin) {
+ val join = proto.getHashJoin
+ val leftKeys = join.getLeftJoinKeysList.asScala
+ val rightKeys = join.getRightJoinKeysList.asScala
+ // Empty and all-NULL build relations can skip different streamed keys
in Spark and
+ // DataFusion. Null-aware anti joins also stop building at the first
NULL key in Spark,
+ // while native execution builds on the opposite side and evaluates keys
in batches.
+ val keys = if (join.getNullAwareAntiJoin) {
+ leftKeys ++ rightKeys
+ } else {
+ join.getJoinType match {
+ case JoinType.Inner | JoinType.LeftSemi | JoinType.LeftAnti =>
+ if (join.getBuildSide == OperatorOuterClass.BuildSide.BuildRight)
leftKeys
Review Comment:
[P2] Validate asserted broadcast build keys when materializing the
broadcast. For `BuildRight`, this guard checks only the left keys. Spark
evaluates the right key to build its hashed relation. Comet broadcasts batches
and evaluates the key inside a join task. If the streamed input has zero
partitions, there is no task to raise the required error. Keep asserted build
keys on Spark or validate them at broadcast materialization with Spark’s
structured error.
Evidence: Read `z(k INT)` from an empty Parquet directory using an explicit
schema. Broadcast `b(k INT)` containing one null and run `SELECT /*+
BROADCAST(b) */ * FROM z JOIN b ON z.k=ann(b.k)`, where `ann` is Catalyst
`AssertNotNull`. Spark 4.1.3 and the base raised `NOT_NULL_ASSERT_VIOLATION`.
The head returned zero rows with codegen on and off. A nonempty streamed-side
control using a left semi join raised a generic `CometNativeException` on the
head and retained SQLSTATE `42000` on the references.
##########
native/spark-expr/src/predicate_funcs/assert_not_null.rs:
##########
@@ -0,0 +1,213 @@
+// 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.
+
+use std::fmt::{Display, Formatter};
+use std::hash::Hash;
+use std::sync::Arc;
+
+use arrow::array::new_empty_array;
+use arrow::datatypes::{DataType, Schema};
+use arrow::record_batch::RecordBatch;
+use datafusion::common::Result;
+use datafusion::logical_expr::ColumnarValue;
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion_comet_common::SparkError;
+
+/// Spark's runtime nullability check. Only the value itself is checked, not
its nested fields.
+#[derive(Debug, Eq)]
+pub struct AssertNotNull {
+ child: Arc<dyn PhysicalExpr>,
+ walked_type_path: String,
+}
+
+impl AssertNotNull {
+ pub fn new(child: Arc<dyn PhysicalExpr>, walked_type_path: String) -> Self
{
+ Self {
+ child,
+ walked_type_path,
+ }
+ }
+}
+
+impl PartialEq for AssertNotNull {
+ fn eq(&self, other: &Self) -> bool {
+ self.child.eq(&other.child) && self.walked_type_path ==
other.walked_type_path
+ }
+}
+
+impl Hash for AssertNotNull {
+ fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+ self.child.hash(state);
+ self.walked_type_path.hash(state);
+ }
+}
+
+impl Display for AssertNotNull {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ write!(
+ f,
+ "AssertNotNull({}, {:?})",
+ self.child, self.walked_type_path
+ )
+ }
+}
+
+impl PhysicalExpr for AssertNotNull {
+ fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ Display::fmt(self, f)
+ }
+
+ fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
+ self.child.data_type(input_schema)
+ }
+
+ fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
+ Ok(false)
+ }
+
+ fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
+ // There is no row to check, including when a conditional filters out
every row.
+ if batch.num_rows() == 0 {
+ return Ok(ColumnarValue::Array(new_empty_array(
+ &self.data_type(batch.schema().as_ref())?,
+ )));
+ }
+ let value = self.child.evaluate(batch)?;
Review Comment:
[P2] Preserve the first error when an asserted child can throw.
`self.child.evaluate(batch)` evaluates later rows before checking whether the
first result is null. In a required-field write, an arithmetic error on row two
can therefore bypass the application’s handler for the null on row one.
Evaluate the child and assertion in row order for this case, or keep assertions
over potentially throwing children in Spark.
Evidence: On Spark 4.1.3 with ANSI enabled, insert `SELECT 10L DIV den FROM
src` into an Iceberg table with one `BIGINT NOT NULL` column. Use one ordered
Parquet batch with denominators `[NULL, 0]`. Spark and the base raised
`NOT_NULL_ASSERT_VIOLATION` (`42000`), and a `42000` exit handler returned
`caught`. The head used `CometProject` feeding `CometIcebergWrite` and raised
`DIVIDE_BY_ZERO` (`22012`) past the handler. All three left zero snapshots.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1178,10 +1178,96 @@ case class CometExecRule(session: SparkSession,
queryStagePrep: Boolean = false)
serde.convert(op, builder)
}
}
- return nativeOp.map { nativeOp =>
- val exec = serde.createExec(nativeOp, op)
- rollUpInfoMessages(op, exec)
- exec
+ return nativeOp
+ .filter { proto =>
+ // A serde may emit several native nodes for one Spark operator.
Inspect those nodes,
+ // but stop at real Spark children, whose consumption context was
tagged independently.
+ def containsOwnAssertion(node: Operator): Boolean =
+
CometAssertNotNull.containsNativeAssertion(node.toBuilder.clearChildren().build())
||
+ node.getChildrenList.asScala.exists { child =>
+ child.getPlanId == proto.getPlanId &&
containsOwnAssertion(child)
+ }
+
+ def hasNativeAssertion: Boolean = op match {
+ // TopK's sink proto omits its runtime projection over the
selected rows.
+ case topK: TakeOrderedAndProjectExec =>
+ QueryPlanSerde.withSharedSubexpressions(topK.projectList) {
+ topK.projectList.exists { expression =>
+ QueryPlanSerde
+ .exprToProto(expression, topK.child.output)
+ .exists(CometAssertNotNull.containsNativeAssertion)
+ }
+ }
+ case _: HashAggregateExec | _: ObjectHashAggregateExec =>
+ // Input/grouping expressions consume all rows; only the result
projection is lazy.
+ proto.hasProjection &&
CometAssertNotNull.containsNativeAssertion(
+ proto.getProjection)
+ case _ => containsOwnAssertion(proto)
+ }
+ val fallback = nativeAssertionJoinFallback(proto).orElse {
+ op.getTagValue(CometExecRule.PARTIAL_READER)
+ .filter(_ => hasNativeAssertion)
+ .map(reader =>
+ s"Native AssertNotNull below ${reader.description} could
change which values " +
+ "Spark checks")
+ }
+ fallback.foreach(reason => withFallbackReason(op, reason))
+ fallback.isEmpty
+ }
+ .map { nativeOp =>
+ val exec = serde.createExec(nativeOp, op)
+ rollUpInfoMessages(op, exec)
+ exec
+ }
+ }
+ None
+ }
+
+ private def nativeAssertionJoinFallback(proto: Operator): Option[String] = {
+ import OperatorOuterClass.JoinType
+
+ // These consumers can skip candidates even without a LIMIT above the
join. Inspect only
+ // the serialized condition. Child projections use the partial-reader tags
independently;
+ // assertions captured inside opaque JVM dispatch keep their existing
execution path.
+ val condition = if (proto.hasHashJoin) {
+ val join = proto.getHashJoin
+ Some(join.getJoinType -> join.getCondition)
+ } else if (proto.hasSortMergeJoin) {
+ val join = proto.getSortMergeJoin
+ Some(join.getJoinType -> join.getCondition)
+ } else if (proto.hasBroadcastNestedLoopJoin) {
+ val join = proto.getBroadcastNestedLoopJoin
+ Some(join.getJoinType -> join.getCondition)
+ } else None
+
+ if (condition.exists {
+ case (JoinType.LeftSemi | JoinType.LeftAnti, expression) =>
+ CometAssertNotNull.containsNativeAssertion(expression)
+ case _ => false
+ }) {
+ return Some(
+ "Native AssertNotNull in an early-exit join condition could check
candidates Spark skips")
+ }
+
+ if (proto.hasHashJoin) {
+ val join = proto.getHashJoin
+ val leftKeys = join.getLeftJoinKeysList.asScala
+ val rightKeys = join.getRightJoinKeysList.asScala
+ // Empty and all-NULL build relations can skip different streamed keys
in Spark and
+ // DataFusion. Null-aware anti joins also stop building at the first
NULL key in Spark,
+ // while native execution builds on the opposite side and evaluates keys
in batches.
+ val keys = if (join.getNullAwareAntiJoin) {
+ leftKeys ++ rightKeys
+ } else {
+ join.getJoinType match {
+ case JoinType.Inner | JoinType.LeftSemi | JoinType.LeftAnti =>
+ if (join.getBuildSide == OperatorOuterClass.BuildSide.BuildRight)
leftKeys
+ else rightKeys
+ case _ => Seq.empty
Review Comment:
[P2] Preserve short-circuit evaluation of packed outer-join keys. `case _ =>
Seq.empty` admits every asserted outer hash-join key. Spark’s
`HashJoin.rewriteKeyExpr` packs two `INT` keys into one null-propagating `LONG`
expression. A null first key skips the second. Comet evaluates the component
keys independently and can fail a left outer join that should return its
unmatched row. Decline assertions in skippable packed-key components or
preserve Spark’s packed evaluation.
Evidence: Use Parquet inputs `l(a,b)=[(NULL,NULL)]` and `r(c,d)=[(1,2)]`,
all `INT`. Run `SELECT /*+ BROADCAST(r) */ * FROM l LEFT JOIN r ON l.a=r.c AND
ann(l.b)=r.d`, with `ann` bound to Catalyst `AssertNotNull`. Spark 4.1.3 and
the base returned `[NULL,NULL,NULL,NULL]`. The head selected
`CometBroadcastHashJoin` and raised `NOT_NULL_ASSERT_VIOLATION` with
whole-stage codegen both on and off.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1328,37 +1414,74 @@ case class CometExecRule(session: SparkSession,
queryStagePrep: Boolean = false)
allExecs.get(op.getClass).exists(_.enabledConfig.forall(_.get(op.conf)))
/**
- * Tags each `SerializeFromObjectExec` whose output an operator above it can
stop reading early
- * with that operator, so `transform` does not convert it. Spark computes
the rows of a typed
- * Dataset operation one at a time, as the operator above reads them, while
the conversion fills
- * a whole Arrow batch first. Below a limit, a `mapPartitions` function such
as `_.take(1)`, or
- * code reading `Dataset.rdd`, the conversion would run the user function on
rows Spark never
- * reaches, and a function that throws on one of them would fail a query
that succeeds in Spark.
- * An operator that reads all of its input before it returns a row ends the
search, since Spark
- * computes every row below it anyway: an exchange, a sort, a hash
aggregate, or a top-k over
- * input that is not already sorted.
+ * Tags operators whose output a consumer can skip evaluating. Spark
computes rows as the
+ * consumer reads them, while native expressions and typed Dataset
conversion fill a whole Arrow
+ * batch first. Below a limit, a `mapPartitions` function such as
`_.take(1)`, or code reading
+ * `Dataset.rdd`, an assertion or user function could throw on rows Spark
never reaches. Joins
+ * can also skip an input, and generated projections, filters, samples and
joins can defer child
+ * fields until an expression needs them. Conversely, native joins can skip
input Spark would
+ * evaluate. An operator that reads all of its input before it returns a row
ends the search,
+ * since Spark computes every row below it anyway: an exchange, a sort, a
hash aggregate, or a
+ * top-k over input that is not already sorted.
*
* Conversion is bottom-up, so this runs first. TreeNode tags survive the
child copies made
- * during transformUp, while an identity set would not.
+ * during transformUp, while an identity set would not. Retain established
tags when AQE visits
+ * a fragment without its limiting ancestor.
*/
- private def tagPartiallyReadTypedDatasetOutputs(plan: SparkPlan): Unit = {
- def visit(op: SparkPlan, partialReader: Option[String]): Unit = {
- val childReader = op match {
- case serialize: SerializeFromObjectExec =>
- partialReader.foreach(
- serialize.setTagValue(CometExecRule.TYPED_DATASET_PARTIAL_READER,
_))
- partialReader
- case _: CollectLimitExec | _: LocalLimitExec | _: GlobalLimitExec =>
Some("a limit")
+ private def tagPartiallyReadOutputs(plan: SparkPlan): Unit = {
+ import CometExecRule.PartialReader
+
+ def visit(op: SparkPlan, partialReader: Option[PartialReader]): Unit = {
+ // AQE can revisit a fragment without the ancestor that could skip whole
rows. A new
+ // field-only reader must not weaken that established context or its
child propagation.
+ val inheritedReader = op
+ .getTagValue(CometExecRule.PARTIAL_READER)
+ .filter(_.canSkipInput)
+ .orElse(partialReader)
+ // A repeated transform can already contain Comet limits and consumption
barriers.
+ val original = op match {
+ case comet: CometExec => comet.originalPlan
+ case other => other
+ }
+ val childReader = original match {
+ case _: CollectLimitExec | _: LocalLimitExec | _: GlobalLimitExec =>
Review Comment:
[P2] Carry the scalar-subquery two-row bound into its child plan. Spark
prepares that child separately and calls `executeTake(2)` to detect a
cardinality violation. The child has no explicit limit for this match to
recognize. A native projection can check row three and replace the cardinality
error, changing which SQL handler runs. Propagate the bound during subquery
preparation or keep the affected projection in Spark.
Evidence: With native range enabled, run `SELECT (SELECT
ann(IF(id<2,id,CAST(NULL AS BIGINT))) FROM range(0,3,1,1))`, using Catalyst
`AssertNotNull` as `ann`. Spark 4.1.3 and the base raised
`SCALAR_SUBQUERY_TOO_MANY_ROWS` (`21000`). The head’s `CometProject` raised
`NOT_NULL_ASSERT_VIOLATION` (`42000`). A `21000` exit handler returned `caught`
on both references and failed on the head.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1178,10 +1178,96 @@ case class CometExecRule(session: SparkSession,
queryStagePrep: Boolean = false)
serde.convert(op, builder)
}
}
- return nativeOp.map { nativeOp =>
- val exec = serde.createExec(nativeOp, op)
- rollUpInfoMessages(op, exec)
- exec
+ return nativeOp
+ .filter { proto =>
+ // A serde may emit several native nodes for one Spark operator.
Inspect those nodes,
+ // but stop at real Spark children, whose consumption context was
tagged independently.
+ def containsOwnAssertion(node: Operator): Boolean =
+
CometAssertNotNull.containsNativeAssertion(node.toBuilder.clearChildren().build())
||
+ node.getChildrenList.asScala.exists { child =>
+ child.getPlanId == proto.getPlanId &&
containsOwnAssertion(child)
+ }
+
+ def hasNativeAssertion: Boolean = op match {
+ // TopK's sink proto omits its runtime projection over the
selected rows.
+ case topK: TakeOrderedAndProjectExec =>
+ QueryPlanSerde.withSharedSubexpressions(topK.projectList) {
+ topK.projectList.exists { expression =>
+ QueryPlanSerde
+ .exprToProto(expression, topK.child.output)
+ .exists(CometAssertNotNull.containsNativeAssertion)
+ }
+ }
+ case _: HashAggregateExec | _: ObjectHashAggregateExec =>
+ // Input/grouping expressions consume all rows; only the result
projection is lazy.
+ proto.hasProjection &&
CometAssertNotNull.containsNativeAssertion(
+ proto.getProjection)
+ case _ => containsOwnAssertion(proto)
+ }
+ val fallback = nativeAssertionJoinFallback(proto).orElse {
+ op.getTagValue(CometExecRule.PARTIAL_READER)
+ .filter(_ => hasNativeAssertion)
Review Comment:
[P2] Apply the assertion guard when Parquet predicates are appended after
planning. `hasNativeAssertion` examines the current protobuf.
`CometNativeScanExec.serializedPartitionData` later resolves scalar subqueries
and appends their filters. A pushed assertion can then inspect a later row even
though the residual Spark filter under `LIMIT` stops at the first row. Leave
asserted late predicates to the residual filter.
Evidence: In one Parquet partition put `p(id,flag)=[(0,true),(1,NULL)]`. Run
`SELECT id FROM p WHERE IF((SELECT max(id) FROM p)>0,ann(flag),true) LIMIT 1`,
with `ann` bound to Catalyst `AssertNotNull`. Spark 4.1.3 and the base returned
`[0]`. With `spark.comet.parquet.rowFilterPushdown.enabled=true`, the head
raised `CometNativeException` in the scan. Turning pushdown off restored `[0]`.
Without the limit, a `42000` handler caught the expected assertion on Spark and
the base and failed on the head.
##########
spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala:
##########
@@ -109,6 +109,7 @@ object QueryPlanSerde extends Logging with CometExprShim
with CometTypeShim {
classOf[LessThanOrEqual] -> CometLessThanOrEqual,
classOf[In] -> CometIn,
classOf[IsNotNull] -> CometIsNotNull,
+ classOf[AssertNotNull] -> CometAssertNotNull,
Review Comment:
[P2] Preserve hash-key evaluation when column pruning leaves a shuffle with
zero data columns. Registering `AssertNotNull` makes an asserted partition
expression eligible for native shuffle. Even with multiple reducers,
`EmptySchemaShufflePartitioner` counts rows and sends them to partition zero
without evaluating the key. Evaluate throwing keys before counting rows, or
decline this native exchange.
Evidence: With AQE off, native range enabled, and two shuffle partitions,
run `SELECT count(*) FROM (SELECT 1 FROM range(0,3,1,1) DISTRIBUTE BY
ann(CAST(NULL AS INT))) q`, using Catalyst `AssertNotNull` as `ann`. Spark
4.1.3 and the base raised the assertion. The head planned `CometExchange
hashpartitioning(assertnotnull(null), 2)` above an empty projection and
returned `3`. Replacing the null with `1` returned `3` on all three.
##########
spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala:
##########
@@ -109,6 +109,7 @@ object QueryPlanSerde extends Logging with CometExprShim
with CometTypeShim {
classOf[LessThanOrEqual] -> CometLessThanOrEqual,
classOf[In] -> CometIn,
classOf[IsNotNull] -> CometIsNotNull,
+ classOf[AssertNotNull] -> CometAssertNotNull,
Review Comment:
[P2] Preserve hash-key evaluation when a nonempty-schema exchange has one
reducer. The new registration admits an asserted key, but
`CometNativeShuffleWriter` serializes this physical `HashPartitioning` as
`SinglePartition` and drops the key expression. Spark still evaluates
`partitionIdExpression`. Keep this exchange on Spark or evaluate the asserted
keys before routing the batch to partition zero.
Evidence: With AQE off, use `l(a,b)=[(NULL,NULL)]` and `r(c,d)=[(1,2)]`, all
`INT`. Run `SELECT /*+ SHUFFLE_HASH(r) */ l.a FROM l LEFT JOIN (SELECT c,d FROM
r WHERE c<0) r ON l.a=r.c AND ann(l.b)=r.d`, where `ann` is Catalyst
`AssertNotNull`. At one reducer, Spark 4.1.3 and the base threw. The head
returned `[NULL]` through a Spark join above `CometExchange
hashpartitioning(a,assertnotnull(b),1)`. At two reducers the head raised the
expected assertion.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1178,10 +1178,96 @@ case class CometExecRule(session: SparkSession,
queryStagePrep: Boolean = false)
serde.convert(op, builder)
}
}
- return nativeOp.map { nativeOp =>
- val exec = serde.createExec(nativeOp, op)
- rollUpInfoMessages(op, exec)
- exec
+ return nativeOp
+ .filter { proto =>
+ // A serde may emit several native nodes for one Spark operator.
Inspect those nodes,
+ // but stop at real Spark children, whose consumption context was
tagged independently.
+ def containsOwnAssertion(node: Operator): Boolean =
+
CometAssertNotNull.containsNativeAssertion(node.toBuilder.clearChildren().build())
||
Review Comment:
[P2] Skip the operator protobuf walk when the source expressions contain no
assertion. A `LIMIT` tags the filter beneath it, causing `containsOwnAssertion`
to inspect every literal in an ordinary `IN` filter. The scalar serializer’s
`NULL_CHECK` fast path does not cover this call. Check for a Catalyst assertion
candidate first or retain its presence during serialization.
Evidence: On Spark 4.1.3 with AQE off and native range enabled, plan
`spark.range(200001).filter(col("id").isin(ids: _*)).limit(1)` with 100,000
distinct `Long` values. After forcing `sparkPlan`, warmed `ThreadMXBean`
measurements for `executedPlan` were 266.6 MB on the base and 360.2 MB on the
head. Both serialized all 100,000 values. Calling the production assertion
search on the filter protobuf alone allocated 89.6 MB and returned `false`.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1372,7 +1495,9 @@ case class CometExecRule(session: SparkSession,
queryStagePrep: Boolean = false)
case _: FilterExec | _: ProjectExec => producesObjects(op.children.head)
case _ => false
}
- visit(plan, if (producesObjects(plan)) Some("code reading Dataset.rdd")
else None)
+ visit(
+ plan,
+ if (producesObjects(plan)) Some(PartialReader("code reading
Dataset.rdd")) else None)
Review Comment:
[P2] Protect partial reads through PySpark `DataFrame.rdd`.
`producesObjects(plan)` recognizes the object-producing root used by Scala
`Dataset.rdd`. PySpark’s `javaToPython` consumes `queryExecution.toRdd`
directly, so an ordinary native projection can check a whole batch before
`rdd.take(1)` returns its first row. Carry partial-consumption information into
that export or restrict native assertions at an externally consumed root.
Evidence: In PySpark 4.1.3, read IDs `0..2047` from one Parquet partition.
Select `ann(IF(id<2047,id,CAST(NULL AS BIGINT))) AS v` and `concat(CAST(id AS
STRING),repeat('p',65536)) AS pad`, with `ann` bound to Catalyst
`AssertNotNull`. The wide payload allows the Python reader to consume a prefix.
`df.rdd.take(1)` returned `v=0` and a 65,537-character payload on Spark and the
base. The head raised `NOT_NULL_ASSERT_VIOLATION` on the last input row.
##########
spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala:
##########
@@ -900,7 +898,22 @@ object QueryPlanSerde extends Logging with CometExprShim
with CometTypeShim {
extractQueryContext(fn).foreach { ctx =>
builder.setQueryContext(ctx)
}
- Some(builder.build())
+ val result = builder.build()
+ // Aggregate arguments are serialized as independent roots, so the
scalar parent guard
+ // cannot see their consumer. Spark can skip them in lazy updates (for
example First),
+ // filtered updates or null checks, while native aggregates can evaluate
whole arrays.
+ if (CometAssertNotNull.containsNativeAssertion(result)) {
Review Comment:
[P2] Gate the aggregate protobuf walk before traversing assertion-free
arguments and filters. `aggExprToProto` calls `containsNativeAssertion(result)`
for every successful aggregate conversion, including fully consumed plans with
no partial reader. A large conditional aggregate pays for reflective
enumeration of every literal during planning. Check the Catalyst aggregate and
filter for assertion candidates before inspecting their native representation.
Evidence: On Spark 4.1.3, plan
`spark.range(200001).agg(sum(when(col("id").isin(ids:
_*),lit(1L)).otherwise(lit(0L))))` with 100,000 distinct `Long` values, native
range enabled, and AQE off. Both revisions produced a native aggregate. The
head’s production assertion search on one 1.68 MB aggregate protobuf allocated
115.2 MB on the driver thread and returned `false`. In the same warmed run,
obtaining `executedPlan` allocated 533.4 MB on the base and 992.6 MB on the
head.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1328,37 +1414,74 @@ case class CometExecRule(session: SparkSession,
queryStagePrep: Boolean = false)
allExecs.get(op.getClass).exists(_.enabledConfig.forall(_.get(op.conf)))
/**
- * Tags each `SerializeFromObjectExec` whose output an operator above it can
stop reading early
- * with that operator, so `transform` does not convert it. Spark computes
the rows of a typed
- * Dataset operation one at a time, as the operator above reads them, while
the conversion fills
- * a whole Arrow batch first. Below a limit, a `mapPartitions` function such
as `_.take(1)`, or
- * code reading `Dataset.rdd`, the conversion would run the user function on
rows Spark never
- * reaches, and a function that throws on one of them would fail a query
that succeeds in Spark.
- * An operator that reads all of its input before it returns a row ends the
search, since Spark
- * computes every row below it anyway: an exchange, a sort, a hash
aggregate, or a top-k over
- * input that is not already sorted.
+ * Tags operators whose output a consumer can skip evaluating. Spark
computes rows as the
+ * consumer reads them, while native expressions and typed Dataset
conversion fill a whole Arrow
+ * batch first. Below a limit, a `mapPartitions` function such as
`_.take(1)`, or code reading
+ * `Dataset.rdd`, an assertion or user function could throw on rows Spark
never reaches. Joins
+ * can also skip an input, and generated projections, filters, samples and
joins can defer child
+ * fields until an expression needs them. Conversely, native joins can skip
input Spark would
+ * evaluate. An operator that reads all of its input before it returns a row
ends the search,
+ * since Spark computes every row below it anyway: an exchange, a sort, a
hash aggregate, or a
+ * top-k over input that is not already sorted.
*
* Conversion is bottom-up, so this runs first. TreeNode tags survive the
child copies made
- * during transformUp, while an identity set would not.
+ * during transformUp, while an identity set would not. Retain established
tags when AQE visits
+ * a fragment without its limiting ancestor.
*/
- private def tagPartiallyReadTypedDatasetOutputs(plan: SparkPlan): Unit = {
- def visit(op: SparkPlan, partialReader: Option[String]): Unit = {
- val childReader = op match {
- case serialize: SerializeFromObjectExec =>
- partialReader.foreach(
- serialize.setTagValue(CometExecRule.TYPED_DATASET_PARTIAL_READER,
_))
- partialReader
- case _: CollectLimitExec | _: LocalLimitExec | _: GlobalLimitExec =>
Some("a limit")
+ private def tagPartiallyReadOutputs(plan: SparkPlan): Unit = {
+ import CometExecRule.PartialReader
+
+ def visit(op: SparkPlan, partialReader: Option[PartialReader]): Unit = {
+ // AQE can revisit a fragment without the ancestor that could skip whole
rows. A new
+ // field-only reader must not weaken that established context or its
child propagation.
+ val inheritedReader = op
+ .getTagValue(CometExecRule.PARTIAL_READER)
+ .filter(_.canSkipInput)
+ .orElse(partialReader)
+ // A repeated transform can already contain Comet limits and consumption
barriers.
+ val original = op match {
+ case comet: CometExec => comet.originalPlan
+ case other => other
+ }
+ val childReader = original match {
+ case _: CollectLimitExec | _: LocalLimitExec | _: GlobalLimitExec =>
+ Some(PartialReader("a limit"))
// A top-k reads only its first rows when its input is already sorted.
case topK: TakeOrderedAndProjectExec
if SortOrder.orderingSatisfies(topK.child.outputOrdering,
topK.sortOrder) =>
- Some("a limit")
- case _: MapPartitionsExec => Some("a mapPartitions function")
+ Some(PartialReader("a limit"))
+ // Global window limits can stop below an exchange, without an upper
reader tag.
+ case w if
ShimCometWindowGroupLimit.extract(w).exists(_.partitionSpec.isEmpty) =>
+ Some(PartialReader("a window group limit"))
+ case _: MapPartitionsExec => Some(PartialReader("a mapPartitions
function"))
+ // These consumers defer projected fields, but still read every input
object. Keep
+ // typed Dataset conversion eligible unless an ancestor can skip whole
input rows.
+ case _: ProjectExec =>
+ inheritedReader.orElse(Some(PartialReader("a projection",
canSkipInput = false)))
+ case _: FilterExec =>
+ inheritedReader.orElse(Some(PartialReader("a filter", canSkipInput =
false)))
+ case _: SampleExec =>
+ inheritedReader.orElse(Some(PartialReader("a sample", canSkipInput =
false)))
+ case _: BroadcastHashJoinExec | _: ShuffledHashJoinExec | _:
SortMergeJoinExec |
+ _: BroadcastNestedLoopJoinExec =>
+ Some(PartialReader("a join"))
Review Comment:
[P2] Keep typed Dataset conversion eligible on the preserved streamed side
of a fully consumed broadcast left outer join. `PartialReader("a join")` marks
both children as able to skip whole input rows and disables
`SerializeFromObjectExec` conversion. Spark’s `HashJoin.outerJoin` consumes
every left row, including unmatched rows. Use a field-only reader for this side
while retaining any ancestor that can genuinely stop consumption.
Evidence: On Spark 4.1.3 `local[1]`, enable typed Dataset conversion and
native range, disable AQE, and use one shuffle partition. Map 1.2 million typed
rows to `(id+1,(id%500000).toString)`, left-outer-join broadcast `range(2048)`
on `k%2048=id`, then group by the string and sum the build-side `id+7`. Both
revisions returned the same grand total, `1236479040`. The base planned
`CometSparkRowToColumnar`, `CometBroadcastHashJoin`, and native aggregates. The
head used Spark for the join and aggregates. With whole-stage codegen off, five
warmed runs had medians of 743 ms on the base and 1,220 ms on the head.
--
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]
