Support fully-native multi-stage (distinct-combined) collect_list / collect_set

Describe the bug

collect_list / collect_set declare their aggregate buffer as BinaryType in Spark (serialized TypedImperativeAggregate state) but produce a native ArrayType (Arrow List) state in Comet. Comet bridges this only for the simple two-stage shape: CometObjectHashAggregateExec.adjustOutputForNativeState rewrites the buffer column of a pure-Partial aggregate to ArrayType, so Partial -> Final runs natively and correctly.

It does not handle multi-stage aggregates that contain a PartialMerge stage, which Spark introduces for the distinct-aggregate rewrite, e.g.:

SELECT x, count(DISTINCT y), collect_list(z) FROM t GROUP BY x

This plans as Partial(x,y) -> PartialMerge(x,y) -> [PartialMerge, Partial](x) -> Final(x). The intermediate PartialMerge outputs are still declared BinaryType, so a fully-native pipeline crashes at runtime:

CometNativeException: Cast error: Cannot cast LIST to non-list data type Binary

(and a related nullability drift List(non-null T) vs List(nullable T) for nested element types). This is a facet of the broader Arrow-type-drift tracked in #4515.

This affects both collect_list (added in #4720) and the already-shipped collect_set — collect_set has the same latent crash on this shape.

Current behavior (workaround in #4720)

To avoid the crash, multi-stage collect_list/collect_set aggregates now fall back to Spark consistently:

CometExecRule.tagUnsafePartialAggregates tags the feeding pure-Partial when a PartialMerge stage of a CollectList/CollectSet aggregate is present.
CometBaseAggregate.doConvert falls back a PartialMerge/Final stage of these functions when no Comet partial produced the buffer (the cross-engine LocalTableScan case).

The simple two-stage collect_list/collect_set cases continue to run natively.

Expected behavior

Enable fully-native multi-stage execution by correcting the intermediate buffer schema for PartialMerge (and multi-mode) stages of CollectList/CollectSet (extend adjustOutputForNativeState beyond pure-Partial, and fix the element-nullability drift), so the ArrayType buffer round-trips across all stages. Then remove the fallback guards added in #4720.

Additional context

Related: DataFusion / DataFusion-Spark functions whose Arrow return type drifts from Spark catalyst's declared type #4515 (DataFusion / datafusion-spark functions whose Arrow return type drifts from Spark catalyst's declared type)
Introduced alongside: feat: native collect_list / array_agg aggregate #4720 (native collect_list / array_agg)

fix: fall back collect_list/collect_set in multi-stage distinct aggre…

take

feat: native collect_list / array_agg aggregate (#4720)

above Issue reproducible with v0.17.0 using below sample queries and data.

SELECT imp_campaign_id, imp_adgroup_id, ARRAY_JOIN(COLLECT_SET(STRING(imp_creative_id)), ',') AS creatives, COUNT(DISTINCT imp_impression_id) AS imp_count FROM impressions GROUP BY imp_campaign_id, imp_adgroup_id

SELECT imp_campaign_id, COUNT(DISTINCT imp_user_id) AS users, COLLECT_SET(imp_creative_id) AS creatives FROM impressions GROUP BY imp_campaign_id

impressions.csv

@chetan-habu

@chetan-habu thanks for reporting this.

above Issue reproducible with v0.17.0 using below sample queries and data....

#4727 should fix this.

you can

checkout to that pr and run

make release PROFILES="-Pspark-4.1 -Drat.skip=true"
export COMET_JAR=spark/target/comet-spark-spark4.1_2.13-1.0.0-SNAPSHOT.jar
sdk use spark 4.1.0

store following scala script as 4724.scala
download csv with curl -L https://github.com/user-attachments/files/30888137/impressions.csv -o impressions.csv
start spark shell

spark-shell \
    --jars $COMET_JAR \
    --conf spark.driver.extraClassPath=$COMET_JAR \
    --conf spark.executor.extraClassPath=$COMET_JAR \
    --conf spark.plugins=org.apache.spark.CometPlugin \
    --conf spark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager \
    --conf spark.comet.explainFallback.enabled=true \
    --conf spark.memory.offHeap.enabled=true \
    --conf spark.memory.offHeap.size=4g

in spark shell, run :load 4724.scala

4724.scala:

val input = "impressions.csv"
val parquet = "impressions.parquet"

spark.read
  .option("header", "true")
  .option("inferSchema", "true")
  .csv(input)
  .repartition(2)
  .write
  .mode("overwrite")
  .parquet(parquet)

spark.read.parquet(parquet).createOrReplaceTempView("impressions")

val q1 =
  """
    |SELECT
    |  imp_campaign_id,
    |  imp_adgroup_id,
    |  ARRAY_JOIN(COLLECT_SET(STRING(imp_creative_id)), ',') AS creatives,
    |  COUNT(DISTINCT imp_impression_id) AS imp_count
    |FROM impressions
    |GROUP BY imp_campaign_id, imp_adgroup_id
    |""".stripMargin

val q2 =
  """
    |SELECT
    |  imp_campaign_id,
    |  COUNT(DISTINCT imp_user_id) AS users,
    |  COLLECT_SET(imp_creative_id) AS creatives
    |FROM impressions
    |GROUP BY imp_campaign_id
    |""".stripMargin

def run(comet: Boolean) = {
  spark.conf.set("spark.comet.enabled", comet.toString)

  val df1 = spark.sql(q1)
  val result1 = df1.collect().map { r =>
    (
      r.getInt(0),
      r.getInt(1),
      r.getString(2).split(",").toSet,
      r.getLong(3))
  }.toSet

  val df2 = spark.sql(q2)
  val result2 = df2.collect().map { r =>
    (r.getInt(0), r.getLong(1), r.getSeq[Int](2).toSet)
  }.toSet

  (result1, result2, Seq(
    df1.queryExecution.executedPlan.toString,
    df2.queryExecution.executedPlan.toString))
}

val sparkResult = run(false)
val cometResult = run(true)

val answersMatch = (
  sparkResult._1 == cometResult._1 &&
  sparkResult._2 == cometResult._2
)

val requiredPlanParts = Seq(
  "CometNativeScan",
  "CometExchange",
  "CometHashAggregate",
  "[Partial]",
  "[PartialMerge]",
  "[PartialMerge, Partial]",
  "[Final]",
  "merge_collect_set")

val fullyNative =
  cometResult._3.forall(plan => requiredPlanParts.forall(plan.contains))

println(s"answersMatch=$answersMatch fullyNative=$fullyNative")
cometResult._3.foreach(println)

if (!answersMatch || !fullyNative) System.exit(1)
System.exit(0)

feat: native collect_list / array_agg aggregate (apache#4720)

fix: walk through sorts when pairing sort aggregate buffers

test: expect the distinct collect_list sort aggregate chain to run na…

feat: support SortAggregateExec (apache#4565)

Navigation Menu

Repository navigation

Support fully-native multi-stage (distinct-combined) collect_list / collect_set #4724

Describe the bug

Current behavior (workaround in #4720)

Expected behavior

Additional context

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Support fully-native multi-stage (distinct-combined) collect_list / collect_set #4724

Description

Describe the bug

Current behavior (workaround in #4720)

Expected behavior

Additional context

Activity

peterxcli commented on Jun 25, 2026

chetan-habu commented on Aug 10, 2026

peterxcli commented on Aug 10, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions