comphead commented on code in PR #5903:
URL: https://github.com/apache/datafusion-comet/pull/5903#discussion_r4265451952
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -366,6 +366,54 @@ class CometAggregateSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
classOf[LocalTableScanExec])
}
+ test("collect_list and collect_set over non-nullable nested fields") {
+ withTempPath { path =>
+ // Keep every group in one input and final partition so aggregation
outgrows its memory pool.
+ spark
+ .createDataFrame((0 until 32768).map { i =>
+ (i % 4096, i % 8192, s"${i % 8192}-" + "x" * 256)
+ })
+ .coalesce(1)
+ .write
+ .parquet(path.getAbsolutePath)
+ withParquetTable(path.getAbsolutePath, "tbl") {
+ // The literal and coalesce create required fields after the Parquet
scan, which
+ // otherwise widens field nullability. Sorting removes collect order
differences.
+ val struct = "named_struct('flag', true, 'id', coalesce(_2, 0),
'value', _3)"
+ // Ungrouped aggregation uses Accumulator rather than
GroupsAccumulator and fails
+ // on mismatched nested nullability even without spilling.
+ checkSparkAnswerAndOperator(
+ sql(s"""
+ SELECT sort_array(collect_list(s)), sort_array(collect_set(s))
+ FROM (SELECT $struct AS s FROM tbl WHERE _1 < 8)"""),
+ Seq(classOf[CometHashAggregateExec]))
+
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1",
+ CometConf.COMET_BATCH_SIZE.key -> "1024",
+ CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> "0.002",
+ CometConf.COMET_RESPECT_DATAFUSION_CONFIGS.key -> "true",
+ "spark.comet.datafusion.execution.sort_spill_reservation_bytes" ->
"65536") {
+ Seq(struct, s"array($struct)").foreach { value =>
+ val (_, cometPlan) = checkSparkAnswerAndOperator(
+ sql(s"""
+ SELECT _1, sort_array(collect_list(s)),
sort_array(collect_set(s))
+ FROM (
+ SELECT _1, $value AS s
+ FROM tbl
+ ) GROUP BY _1"""),
+ Seq(classOf[CometHashAggregateExec]))
+ val aggregates = collect(cometPlan) { case agg:
CometHashAggregateExec => agg }
+ assert(
+ aggregates.map(_.metrics("spill_count").value).sum > 0L,
+ s"Expected the native collection aggregate to
spill:\n$cometPlan")
Review Comment:
This file already has `withDistinctMergeMemoryPressure` (AQE off, one
shuffle partition, batch size 1024, pool fraction 0.002) and
`spilledAggregates(df, mode)`, which the #6254 and #6857 tests use. The first
four settings here are the helper's. The hand-written check sums `spill_count`
over every `CometHashAggregateExec`, so it does not show that the final
aggregate (the stage #5239 is about) spilled. Could this reuse the helper and
call `spilledAggregates(df, Final)` on the query's DataFrame? Also, are
`COMET_RESPECT_DATAFUSION_CONFIGS` and `sort_spill_reservation_bytes` needed
here? The #6857 tests do not set them. I did not run either variant, so I would
drop them unless the test fails without them.
##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -2488,11 +2488,12 @@ trait CometBaseAggregate[T <: BaseAggregateExec]
extends CometOperatorSerde[T] {
val bufferAttrs = aggFunc.aggBufferAttributes
aggFunc match {
case cs: CollectSet =>
- val elementType = cs.children.head.dataType
+ // Match the native planner's collect argument normalization,
including nested fields.
+ val elementType = cs.children.head.dataType.asNullable
val nativeStateType = ArrayType(elementType, containsNull = true)
output(bufferIdx) = output(bufferIdx).withDataType(nativeStateType)
case cl: CollectList =>
- val elementType = cl.children.head.dataType
+ val elementType = cl.children.head.dataType.asNullable
val nativeStateType = ArrayType(elementType, containsNull = true)
output(bufferIdx) = output(bufferIdx).withDataType(nativeStateType)
Review Comment:
The `CollectSet` and `CollectList` branches are now identical apart from the
pattern, and the new comment is only on the first. Since both lines are being
edited anyway, could this be one branch, for example `case _: CollectSet | _:
CollectList =>` using `aggFunc.children.head.dataType.asNullable`? That keeps a
single place that mirrors the native normalization.
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -366,6 +366,54 @@ class CometAggregateSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
classOf[LocalTableScanExec])
}
+ test("collect_list and collect_set over non-nullable nested fields") {
+ withTempPath { path =>
+ // Keep every group in one input and final partition so aggregation
outgrows its memory pool.
+ spark
+ .createDataFrame((0 until 32768).map { i =>
+ (i % 4096, i % 8192, s"${i % 8192}-" + "x" * 256)
+ })
+ .coalesce(1)
+ .write
+ .parquet(path.getAbsolutePath)
+ withParquetTable(path.getAbsolutePath, "tbl") {
+ // The literal and coalesce create required fields after the Parquet
scan, which
+ // otherwise widens field nullability. Sorting removes collect order
differences.
+ val struct = "named_struct('flag', true, 'id', coalesce(_2, 0),
'value', _3)"
+ // Ungrouped aggregation uses Accumulator rather than
GroupsAccumulator and fails
+ // on mismatched nested nullability even without spilling.
+ checkSparkAnswerAndOperator(
+ sql(s"""
+ SELECT sort_array(collect_list(s)), sort_array(collect_set(s))
+ FROM (SELECT $struct AS s FROM tbl WHERE _1 < 8)"""),
+ Seq(classOf[CometHashAggregateExec]))
Review Comment:
This ungrouped query needs no memory configuration or spill assertion, so it
fits a SQL file test better than this Scala test. It would also be cheap to
cover both shapes and both functions there: `collect_list.sql` and
`collect_set.sql` could each get a query over `named_struct('flag', true, 'id',
coalesce(...), ...)` and over `array(named_struct(...))` from a Parquet table,
with `sort_array`, which `CometSqlFileTestSuite` checks for answer and native
execution. Here only the struct shape is covered ungrouped. The grouped spill
case has to stay in Scala because the SQL harness has no spill-count assertion.
--
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]
