CometHashAggregate prefixed with ! in explain plan

[rishvin]

Describe the bug

I invoked the following in the spark-shell, version 3.5.6.

$SPARK_HOME/bin/spark-shell $COMET/spark/target/comet-spark-spark3.5_2.12-0.10.0-SNAPSHOT --conf spark.plugins=org.apache.spark.CometPlugin --conf spark.comet.enabled=true --conf spark.comet.exec.enabled=true

import org.apache.spark.sql.types._
import org.apache.spark.sql.functions._
import org.apache.spark.sql.Row

val schema = StructType(Seq(StructField("id", IntegerType, nullable = false), StructField("value", IntegerType, nullable = false)))
val data = Seq(Row(1, 10), Row(2, 20), Row(3, 10), Row(4, 30), Row(5, 20), Row(6, 10))
val df = spark.createDataFrame(spark.sparkContext.parallelize(data), schema)
val out = "groupby"
df.write.mode("overwrite").parquet(out)
val parquetDF = spark.read.parquet(out)
val grouped = parquetDF.groupBy("id").count()
grouped.explain()

And saw that the explain output has CometHashAggregate prefixed with !.

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- HashAggregate(keys=[id#137], functions=[count(1)])
   +- Exchange hashpartitioning(id#137, 4), ENSURE_REQUIREMENTS, [plan_id=420]
      +- !CometHashAggregate [id#137], Partial, [id#137], [partial_count(1)]
         +- CometNativeScan parquet [id#137] Batched: true, DataFilters: [], Format: CometParquet, Location: InMemoryFileIndex(1 paths)[file:/home/testing/groupby], PartitionFilters: [], PushedFilters: [], ReadSchema: struct<id:int>

There was no correctness/failures issues.

The spark codebase here says - they use "!" to indicate an invalid plan, and "'" to indicate an unresolved plan. However, I haven't verified, if this is the only place in the code from where ! gets added to the plan.

I filed this issue after seeing that comment in the spark code to bring this to notice.

Steps to reproduce

No response

Expected behavior

No response

Additional context

No response

[andygrove]

I started debugging this, and CometHashAggregateExec.missingInputs contains values, which leads to the warnings.

In CometHashAggregateExec, I updated the code:

  override def stringArgs: Iterator[Any] =
    Iterator(input, mode, groupingExpressions, aggregateExpressions, child, "missing", missingInput)

Here is an example output:

   !CometHashAggregate [_2#7, count#25L], Final, [_2#7], [count(_1#6)], missing, {count(_1)#21L, count(_1#6)#20L}
   +- CometExchange hashpartitioning(_2#7, 10), ENSURE_REQUIREMENTS, CometNativeShuffle, [plan_id=211]
      +- !CometHashAggregate [_1#6, _2#7], Partial, [_2#7], [partial_count(_1#6)], missing, {count#25L}
         +- CometScan [native_iceberg_compat] parquet [_1#6,_2#7] ... ReadSchema: struct<_1:int,_2:int>

[comphead]

I remember hitting this ! on other operators, like CometTakeOrderedAndProject, checking it

[comphead]

@rishvin does it happen with AQE only?

[comphead]

test("invalid CometHashAggregate") {
    sql("CREATE TABLE t(id int, value int) USING PARQUET")
    sql("INSERT INTO t VALUES (1, 10)")
    sql("INSERT INTO t VALUES (2, 20)")
    sql("INSERT INTO t VALUES (3, 10)")
    sql("INSERT INTO t VALUES (4, 30)")
    sql("INSERT INTO t VALUES (5, 20)")
    sql("INSERT INTO t VALUES (6, 10)")

    sql("select * from t").write.mode("overwrite").parquet("/tmp/test_11")

    withSQLConf(
      CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
      "spark.comet.cast.allowIncompatible" -> "true",
      "spark.comet.enabled" -> "true",
      "spark.comet.exec.replaceSortMergeJoin" -> "true",
      "spark.comet.exec.shuffle.enableFastEncoding" -> "true",
      "spark.comet.exec.shuffle.enabled" -> "true",
      "spark.comet.exec.shuffle.fallbackToColumnar" -> "true",
      "spark.comet.explain.native.enabled" -> "true",
      "spark.comet.explainFallback.enabled" -> "true",
      "spark.shuffle.manager" -> "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager",
      "spark.comet.logFallbackReasons.enabled" -> "true") {
      spark.read.parquet("/tmp/test_11").createOrReplaceTempView("t2")
      spark.sql("select count(id) from t2 group by id").explain("extended")
      sql("drop table t")
    }


== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- !CometHashAggregate [id#43, count#51L], Final, [id#43], [count(id#43)]
   +- CometExchange hashpartitioning(id#43, 10), ENSURE_REQUIREMENTS, CometNativeShuffle, [plan_id=141]
      +- !CometHashAggregate [id#43], Partial, [id#43], [partial_count(id#43)]
         +- CometScan [native_iceberg_compat] parquet [id#43] Batched: true, DataFilters: [], Format: CometParquet, Location: InMemoryFileIndex(1 paths)[file:/tmp/test_11], PartitionFilters: [], PushedFilters: [], ReadSchema: struct<id:int>

This is interesting as ! usually says the reference from preceding plan node is missing, which looks okay here

and it still warns without AQE

[comphead]

Spark considers count#51L to be as input, and thinks it is missing. This is the reason for exclamation mark

UPD: Spark HashAggregate.producedAttributes contains count values, where CometHashAggregate does not, refer to the code below, that's causing the exclamation mark.

  final def missingInput: AttributeSet = references -- inputSet

  lazy val references: AttributeSet = {
    AttributeSet.fromAttributeSets(expressions.map(_.references)) -- producedAttributes
  }

It does seem it is not an issue with wrong plan
@andygrove should we address this in 0.11.0?

[andygrove]

Spark considers count#51L to be as input, and thinks it is missing. This is the reason for exclamation mark

UPD: Spark HashAggregate.producedAttributes contains count values, where CometHashAggregate does not, refer to the code below, that's causing the exclamation mark.

  final def missingInput: AttributeSet = references -- inputSet

  lazy val references: AttributeSet = {
    AttributeSet.fromAttributeSets(expressions.map(_.references)) -- producedAttributes
  }

It does seem it is not an issue with wrong plan @andygrove should we address this in 0.11.0?

Yes, I think it would be good to fix this for 0.11.0.