andygrove commented on code in PR #5821:
URL: https://github.com/apache/datafusion-comet/pull/5821#discussion_r3973834013
##########
spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala:
##########
@@ -4038,6 +4045,282 @@ class CometExecSuite extends CometTestBase {
}
}
+ private def emptyRelation(
+ output: Seq[org.apache.spark.sql.catalyst.expressions.Attribute]):
SparkPlan = {
+ assume(isSpark40Plus, "EmptyRelationExec requires Spark 4.0+")
+ ShimCometEmptyRelation.create(LocalRelation(output)).get
+ }
+
+ private def localTableScan(
+ output: Seq[org.apache.spark.sql.catalyst.expressions.Attribute]):
LocalTableScanExec =
+
spark.sessionState.planner.plan(LocalRelation(output)).next().asInstanceOf[LocalTableScanExec]
+
+ private def prepareEmptyInputPlan(plan: SparkPlan, native: Boolean):
SparkPlan = {
+ val required = EnsureRequirements().apply(plan)
+ val converted = if (native) CometExecRule(spark).apply(required) else
required
+ ApplyColumnarRulesAndInsertTransitions(Seq.empty, false).apply(converted)
+ }
+
+ private def collectEmptyInputPlan(plan: SparkPlan): Seq[Row] = {
+ val toRow = CatalystTypeConverters.createToScalaConverter(plan.schema)
+ plan.executeCollect().map(row => toRow(row).asInstanceOf[Row]).toSeq
+ }
+
+ private def withCometDisabled(plan: => SparkPlan): SparkPlan = {
+ var result: SparkPlan = null
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ result = plan
+ }
+ result
+ }
+
+ test("EmptyRelationExec preserves attributes and zero partitions") {
+ val attributes = Seq(
+ AttributeReference("id", IntegerType, nullable = false)(),
+ AttributeReference("amount", DecimalType(18, 4), nullable = true)(),
+ AttributeReference("nested", ArrayType(IntegerType, containsNull =
false))())
+ val original = emptyRelation(attributes)
+ val converted =
CometExecRule(spark).apply(original).asInstanceOf[CometEmptyRelationExec]
+ assert(converted.output == original.output)
+ assert(converted.schema == original.schema)
+ assert(converted.children.isEmpty)
+ assert(converted.executeColumnar().getNumPartitions == 0)
+ assert(converted.doExecuteAsArrowStream().getNumPartitions == 0)
+ assert(converted.executeCollect().isEmpty)
+ assert(
+ converted.canonicalized ==
+
CometExecRule(spark).apply(emptyRelation(attributes.map(_.newInstance()))).canonicalized)
+
+ withSQLConf(CometConf.COMET_EXEC_EMPTY_RELATION_ENABLED.key -> "false") {
+ assert(CometExecRule(spark).apply(emptyRelation(attributes)).getClass ==
original.getClass)
+ }
+ val unsupported = emptyRelation(Seq(AttributeReference("empty_struct",
StructType(Nil))()))
+ assert(CometExecRule(spark).apply(unsupported).getClass ==
unsupported.getClass)
+ }
+
+ test("EmptyRelationExec supports global and grouped COUNT and SUM") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2") {
+ withTempView("empty_aggregate_input") {
+ Seq((1, 2)).toDF("k",
"v").createOrReplaceTempView("empty_aggregate_input")
+ val queries = Seq(
+ "SELECT count(*), count(v), sum(v) FROM empty_aggregate_input" ->
Seq(
+ Row(0L, 0L, null)),
+ "SELECT sum(v) FROM empty_aggregate_input" -> Seq(Row(null)),
+ "SELECT count(*) FROM empty_aggregate_input" -> Seq(Row(0L)),
+ "SELECT k, count(*), sum(v) FROM empty_aggregate_input GROUP BY k"
-> Seq.empty)
+ for {
+ (query, expected) <- queries
+ singleEmptyPartition <- Seq(false, true)
+ } {
+ def original: SparkPlan = withCometDisabled {
+ sql(query).queryExecution.sparkPlan.transformUp { case leaf:
LocalTableScanExec =>
+ val empty = emptyRelation(leaf.output)
+ // COALESCE(1) turns zero partitions into one empty columnar
input stream.
+ if (singleEmptyPartition) CoalesceExec(1, empty) else empty
+ }
+ }
+ val native = prepareEmptyInputPlan(original, native = true)
+ withClue(s"$query\n$native") {
+ assert(
+ collectEmptyInputPlan(prepareEmptyInputPlan(original, native =
false)) == expected)
+ assert(collectEmptyInputPlan(native) == expected)
+ assert(native.collect { case e: CometEmptyRelationExec => e }.size
== 1)
+ assert(native.collect { case a: CometHashAggregateExec => a }.size
== 2)
+ assert(native.collect {
+ case a:
org.apache.spark.sql.execution.aggregate.HashAggregateExec => a
+ }.isEmpty)
+ }
+ }
+ }
+ }
+ }
+
+ test("EmptyRelationExec supports empty build and probe hash joins") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2",
+ CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") {
+ for {
+ broadcast <- Seq(false, true)
+ joinType <- Seq(Inner, LeftOuter, RightOuter, FullOuter, LeftSemi,
LeftAnti)
+ if !(broadcast && joinType == FullOuter)
+ (emptyLeft, emptyRight) <- Seq((false, true), (true, false), (true,
true))
+ } {
+ def original: SparkPlan = {
+ val leftKey = AttributeReference("l", IntegerType, nullable = true)()
+ val rightKey = AttributeReference("r", IntegerType, nullable =
true)()
+ val left =
+ if (emptyLeft) emptyRelation(Seq(leftKey))
+ else
+ localTableScan(Seq(leftKey))
+ .copy(rows = Seq(InternalRow(1), InternalRow(1),
InternalRow(null)))
+ val right =
+ if (emptyRight) emptyRelation(Seq(rightKey))
+ else localTableScan(Seq(rightKey)).copy(rows = Seq(InternalRow(1),
InternalRow(null)))
+ val buildSide = if (joinType == RightOuter) BuildLeft else BuildRight
+ if (broadcast) {
+ BroadcastHashJoinExec(
+ Seq(leftKey),
+ Seq(rightKey),
+ joinType,
+ buildSide,
+ None,
+ left,
+ right)
+ } else {
+ ShuffledHashJoinExec(
+ Seq(leftKey),
+ Seq(rightKey),
+ joinType,
+ buildSide,
+ None,
+ left,
+ right)
+ }
+ }
+ val native = prepareEmptyInputPlan(original, native = true)
+ withClue(s"$joinType broadcast=$broadcast
empty=($emptyLeft,$emptyRight)\n$native") {
+ val expected = collectEmptyInputPlan(prepareEmptyInputPlan(original,
native = false))
+ assert(
+ collectEmptyInputPlan(native).groupBy(identity).map { case (r, rs)
=>
+ r -> rs.size
+ } ==
+ expected.groupBy(identity).map { case (r, rs) => r -> rs.size })
+ assert(native.collect { case e: CometEmptyRelationExec => e
}.nonEmpty)
+ val joins = native.collect {
+ case j: CometBroadcastHashJoinExec => j
+ case j: CometHashJoinExec => j
+ }
+ assert(joins.size == 1)
+ }
+ }
+ }
+ }
+
+ test("EmptyRelationExec retains incompatible aggregate buffer fallback") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.key -> "false") {
+ def original: SparkPlan = withCometDisabled {
+ sql("SELECT avg(v) FROM VALUES (CAST(1 AS DECIMAL(38, 0))) AS
t(v)").queryExecution.sparkPlan
+ .transformUp { case leaf: LocalTableScanExec =>
+ emptyRelation(leaf.output)
+ }
+ }
+ val native = prepareEmptyInputPlan(original, native = true)
+ assert(collectEmptyInputPlan(native) == Seq(Row(null)))
+ assert(native.collect { case e: CometEmptyRelationExec => e }.size == 1)
+ assert(native.collect { case a: CometHashAggregateExec => a }.isEmpty)
+ assert(native.collect {
+ case a: org.apache.spark.sql.execution.aggregate.HashAggregateExec => a
+ }.size == 2)
+ }
+ }
+
+ test("EmptyRelationExec retains existence sort-merge join fallback") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2",
+ CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_SORT_MERGE_JOIN_ENABLED.key -> "true") {
+ val leftKey = AttributeReference("l", IntegerType, nullable = true)()
+ val rightKey = AttributeReference("r", IntegerType, nullable = true)()
+ val left = localTableScan(Seq(leftKey)).copy(rows = Seq(InternalRow(1),
InternalRow(null)))
+ val join = SortMergeJoinExec(
+ Seq(leftKey),
+ Seq(rightKey),
+ ExistenceJoin(AttributeReference("exists", BooleanType, nullable =
false)()),
+ None,
+ left,
+ emptyRelation(Seq(rightKey)))
+ val native = prepareEmptyInputPlan(join, native = true)
+ assert(
+ collectEmptyInputPlan(native).sortBy(_.toString) ==
+ Seq(Row(1, false), Row(null, false)).sortBy(_.toString))
+ assert(native.collect { case e: CometEmptyRelationExec => e }.size == 1)
+ assert(native.collect { case j: SortMergeJoinExec => j }.size == 1)
+ assert(native.collect { case j: CometSortMergeJoinExec => j }.isEmpty)
+ }
+ }
+
+ test("EmptyRelationExec supports reused broadcast and AQE broadcast stages")
{
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") {
+ val leftKey = AttributeReference("l", IntegerType, nullable = true)()
+ val rightKey = AttributeReference("r", IntegerType, nullable = true)()
+ val left = localTableScan(Seq(leftKey)).copy(rows = Seq(InternalRow(1),
InternalRow(null)))
+ def join(right: SparkPlan): SparkPlan = BroadcastHashJoinExec(
+ Seq(leftKey),
+ Seq(rightKey),
+ LeftOuter,
+ BuildRight,
+ None,
+ left,
+ right)
+ val first = prepareEmptyInputPlan(join(emptyRelation(Seq(rightKey))),
native = true)
+ val exchange = first.collectFirst { case e: CometBroadcastExchangeExec
=> e }.get
+ assert(collectEmptyInputPlan(first) == Seq(Row(1, null), Row(null,
null)))
+
+ for (adaptive <- Seq(false, true)) {
+ val reused = ReusedExchangeExec(exchange.output, exchange)
+ val native = if (adaptive) {
+ val input = BroadcastQueryStageExec(0, reused, reused.canonicalized)
+ prepareEmptyInputPlan(join(input), native = true)
+ } else {
+ // Non-AQE reuse runs after native conversion; its consumers already
have native plans.
+ first.transformUp { case _: CometBroadcastExchangeExec => reused }
+ }
+ withClue(s"AQE=$adaptive\n$native") {
+ assert(collectEmptyInputPlan(native) == Seq(Row(1, null), Row(null,
null)))
+ assert(native.collect { case j: CometBroadcastHashJoinExec => j
}.size == 1)
+ assert(collect(native) { case e: ReusedExchangeExec => e }.nonEmpty)
+ }
+ }
+ assert(exchange.metrics("numOutputRows").value == 0)
+ }
+ }
+
+ test("EmptyRelationExec discovered by AQE feeds native aggregates and Spark
existence joins") {
Review Comment:
I could not reproduce the scenario from #5819 with Comet at its default
configuration. Against a Parquet backed table, `SELECT count(*), sum(v) FROM
(SELECT _1 % 2 AS k, sum(_2) AS v FROM t WHERE _1 < 0 GROUP BY _1 % 2)`
produces no `EmptyRelationExec` at all under Comet. Vanilla Spark produces one,
Comet produces neither that nor `CometEmptyRelationExec`, and the query just
runs natively end to end.
The reason looks like `AQEPropagateEmptyRelation.getEstimatedRowCount`. It
only learns a stage's row count from a `QueryStageExec` or a
`BaseAggregateExec` inside the `LogicalQueryStage`, and once Comet has
converted the aggregate that node is a `CometHashAggregateExec`, so the rule
returns `None` and no empty relation is ever created.
This test reaches the operator because `range()` together with
`spark.comet.sparkToColumnar.enabled=false` and
`spark.comet.shuffle.convertFromSparkPlan.enabled=false` keeps the inner
aggregate a Spark `HashAggregateExec`. The shapes that do reach
`CometEmptyRelationExec` at default configs are the join elimination ones. I
measured eleven shapes and only inner join and left semi join with an empty
right side got there, while sort, limit, window, distinct and sort merge join
over an empty stage all produce an `EmptyRelationExec` under plain Spark and
nothing at all under Comet.
Would you consider changing this test to a join shape that reaches the
operator without the config overrides? That would cover the path users are
actually on, and it would fail if a future change made the operator unreachable
again. It also seems worth saying in #5819 that the aggregate case needs the
`BaseAggregateExec` gap addressed separately, since this PR does not close it.
##########
docs/source/user-guide/latest/compatibility/operators.md:
##########
@@ -19,6 +19,19 @@ under the License.
# Operator Compatibility
+## Empty Relations
Review Comment:
`docs/source/user-guide/latest/operators.md` describes itself as the
complete reference for how Comet handles each Spark physical operator, and
`EmptyRelationExec` is not in any of its tables. Could you add a row there as
well, linking to this section and noting the Spark 4.0 and later restriction?
##########
spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala:
##########
@@ -59,6 +62,14 @@ object CometDataWritingCommand extends
CometOperatorSerde[DataWritingCommandExec
case cmd: InsertIntoHadoopFsRelationCommand =>
cmd.fileFormat match {
case _: ParquetFileFormat =>
+ // AQE can replace the write input with a zero-partition empty
relation. Keep
+ // Spark's writer, which creates an empty task to preserve the
output file schema.
+ // The native writer only maps existing partitions and cannot do
that yet.
+ if (hasEmptyRelationInput(op.child)) {
Review Comment:
I checked that this guard is load bearing rather than defensive. With it
patched out, a fresh path write and a CTAS over an AQE eliminated join both
fail with `PATH_NOT_FOUND`, so it is doing real work on the shapes that
actually occur and not only on the aggregate shape in the test.
One thought on how it is written. It keys on finding the operator anywhere
in the write subtree, but the hazard is really the write input having zero
partitions. `hasEmptyRelationInput` recurses through `QueryStageExec.plan` and
through exchange children, so it also fires when the empty relation sits under
a shuffle and the write input has a perfectly normal partition count. Could
this comment name #5303 so that whoever fixes the native writer knows to come
back and delete the guard?
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometEmptyRelationExec.scala:
##########
@@ -0,0 +1,62 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.comet
+
+import scala.reflect.ClassTag
+
+import org.apache.arrow.memory.BufferAllocator
+import org.apache.arrow.vector.ipc.ArrowReader
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.comet.execution.arrow.CometNativeArrowSource
+import org.apache.spark.sql.execution.{LeafExecNode, SparkPlan}
+
+import org.apache.comet.{CometConf, ConfigEntry}
+import org.apache.comet.serde.OperatorOuterClass.Operator
+import org.apache.comet.serde.operator.CometSink
+
+/**
+ * An empty native input. Spark's eliminated logical subtree is explanation
data only; neither it
+ * nor an Arrow reader needs to run. Preserve the zero partitions of
EmptyRelationExec so Spark's
+ * exchanges continue to control aggregate and join partitioning.
+ */
+case class CometEmptyRelationExec(originalPlan: SparkPlan, override val
output: Seq[Attribute])
Review Comment:
Spark's `EmptyRelationExec` overrides `generateTreeString` so the eliminated
logical subtree prints as a pseudo child, which is the main way to see what AQE
removed. This node does not, so that subtree disappears from explain output
once it converts. Here is the same query with
`spark.comet.exec.emptyRelation.enabled` off and then on:
```
EmptyRelation [plan_id=1471]
+- Join LeftSemi, (_1#2 = _1#6)
:- LogicalQueryStage Project [_1#2], ShuffleQueryStage 0
...
```
```
CometEmptyRelation EmptyRelation [plan_id=1592], [_1#2]
```
The docstring just above already calls that subtree explanation data, so
would you delegate `generateTreeString` to `originalPlan` to keep it visible?
##########
spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimCometEmptyRelation.scala:
##########
@@ -0,0 +1,31 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.comet.shims
+
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.execution.{EmptyRelationExec, SparkPlan}
+
+/** EmptyRelationExec is available starting in Spark 4.0. */
+object ShimCometEmptyRelation {
+
+ def emptyRelationClass: Option[Class[_ <: SparkPlan]] =
Some(classOf[EmptyRelationExec])
+
+ def create(logical: LogicalPlan): Option[SparkPlan] =
Some(EmptyRelationExec(logical))
Review Comment:
`create` only has one caller, and it is `emptyRelation` in `CometExecSuite`.
There is a `spark/src/test/spark-4.x` tree that already holds Spark 4 only
suites such as `CometWidthBucketSuite` and `CometShuffle4_0Suite`. Would you
rather put the new tests there? That would let `create` come out of both
production shims, and it would remove the eight `assume(isSpark40Plus)`
cancellations that Spark 3.4 and 3.5 runs report today. A new suite needs
registering in `dev/ci/check-suites.py` and in both `pr_build_linux.yml` and
`pr_build_macos.yml`.
If the tests do move, it might also be worth splitting them. Five of the
seven build a physical plan by hand and call `CometExecRule(spark).apply`
directly, which is the style used in `CometExecRuleSuite` and
`RevertNativeForTransitionHeavyStagesSuite` rather than in `CometExecSuite`.
--
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]
