ErikBPF commented on code in PR #5903:
URL: https://github.com/apache/datafusion-comet/pull/5903#discussion_r4365068125
##########
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:
You are right that summing `spill_count` over every `CometHashAggregateExec`
does not show which stage spilled, and I measured it before changing anything.
On current `main` this fixture spills in neither stage: printing `spill_count`
for every aggregate with the helper's settings alone, and again with
`COMET_RESPECT_DATAFUSION_CONFIGS` plus `sort_spill_reservation_bytes=65536`,
gives `Final = 0` and `Partial = 0` in both runs. Shrinking the pool to 400 KiB
(fraction `0.0002`) does not spill either; the Final aggregate fails the
allocation instead: `Failed to acquire 1122504 bytes ... fair limit is 429496
bytes`.
So neither the old assertion nor `spilledAggregates(df, Final)` can pass on
this fixture, and the two extra configs buy nothing. In `8c5700c5a` I removed
the spill assertion and the hand-written config block. The grouped test keeps
both nested shapes plus the native-execution and Spark-answer assertions, which
is what actually reproduces #5239: the mismatch fires when the Final aggregate
materializes the collected state, and it needs no spilling. The RED evidence is
in the test comment: removing `asNullable` fails with `column types must match
schema types`.
##########
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:
Done in `8c5700c5a`. The ungrouped case moved out of the Scala test into
`collect_list.sql` and `collect_set.sql`: `cl_src_struct_required` /
`cs_src_struct_required` build the required-fields struct and
`array(named_struct(...))` after the Parquet scan, and both functions run
through `sort_array`. The harness's default `query` directive asserts the Spark
answer and native execution, so it covers the same `Accumulator` path without
the memory settings that the Scala copy needed and did not use.
Results: `collect_list` 2/2, `collect_set` 8/8. With `asNullable` removed,
both `collect_list` cases fail with `column types must match schema types`; the
ungrouped `collect_set` cases still pass, so that half is coverage rather than
a regression gate. The grouped Scala case remains the gate for both functions.
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -188,6 +188,46 @@ class CometAggregateSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
classOf[LocalTableScanExec])
}
+ test("collect_list and collect_set over non-nullable nested fields survive
spilling") {
+ 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") {
+ 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") {
+ // 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)"
+ Seq(struct, s"array($struct)").foreach { value =>
Review Comment:
Updated in `8c5700c5a`: your ungrouped query now lives in `collect_list.sql`
/ `collect_set.sql` rather than in the Scala test, following comphead's
placement suggestion. It covers the required-fields struct and `array<struct>`
for both functions, and the SQL harness still asserts the Spark answer and
native execution.
RED check on that head: removing `asNullable` makes both `collect_list`
cases fail with the nullability mismatch. The ungrouped `collect_set` cases
pass either way, so they are coverage rather than a gate; the grouped Scala
case exercises both functions and does fail without the fix.
--
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]
