From 904a43187734ba486cf32f9dcef6e25f09f8b6cb Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 27 Jul 2026 11:51:34 -0600 Subject: [PATCH 1/2] test: harden docs and coverage for native collect_list / collect_set Follow-up to PR #4720. Adds the missing `collect_set` audit entry, covers the split-execution cascade and the multi-stage fallback guard with assertions, and extends the SQL fixture datatype coverage. - docs: add a `## collect_set` section to the agg_funcs audit page and link the `collect_list` entry to it instead of referring to it in prose - test: add mixed-engine `collect_list` tests for Comet partial + Spark final and Spark partial + Comet final, exercising `hasNativeArrayBufferAgg` and `tagUnsafePartialAggregates` - test: assert the exact fallback reason in the distinct-combined test instead of only checking the answer - test: add `timestamp_ntz` (native path) and ANSI interval (Spark fallback) blocks to the `collect_list` SQL fixture --- .../expression-audits/agg_funcs.md | 13 +++- .../expressions/aggregate/collect_list.sql | 43 +++++++++++++ .../comet/exec/CometAggregateSuite.scala | 64 ++++++++++++++++--- 3 files changed, 109 insertions(+), 11 deletions(-) diff --git a/docs/source/contributor-guide/expression-audits/agg_funcs.md b/docs/source/contributor-guide/expression-audits/agg_funcs.md index 63f4ca84b7..9c40ecb3c4 100644 --- a/docs/source/contributor-guide/expression-audits/agg_funcs.md +++ b/docs/source/contributor-guide/expression-audits/agg_funcs.md @@ -57,9 +57,20 @@ - Spark 3.5.8 (audited 2026-06-24): identical to 3.4.3. - Spark 4.0.1 (audited 2026-06-24): only structural change is adding `with UnaryLike[Expression]` to the case class (no behavior change). - Spark 4.1.1 (audited 2026-06-24): identical to 4.0.1. -- Comet implementation: native side delegates to `datafusion_spark::function::aggregate::collect::SparkCollectList`, which wraps `ArrayAggAccumulator` with `ignore_nulls = true` and converts a final NULL accumulator state to an empty array (matching Spark's `defaultResult`). The native return type is `List(Field, containsNull = true)`, while Spark uses `containsNull = false`. Because nulls are filtered before insertion, no nulls actually appear in the array, so this is a schema-shape difference only and tests using `checkSparkAnswerAndOperator` accept it (same pattern already in use for `collect_set`). +- Comet implementation: native side delegates to `datafusion_spark::function::aggregate::collect::SparkCollectList`, which wraps `ArrayAggAccumulator` with `ignore_nulls = true` and converts a final NULL accumulator state to an empty array (matching Spark's `defaultResult`). The native return type is `List(Field, containsNull = true)`, while Spark uses `containsNull = false`. Because nulls are filtered before insertion, no nulls actually appear in the array, so this is a schema-shape difference only and tests using `checkSparkAnswerAndOperator` accept it (the same pattern applies to [collect_set](#collect-set)). - Spark 4.2 (preview): `CollectList` and `CollectSet` gain an `ignoreNulls` field (default `true`); `RESPECT NULLS` sets it to `false` and keeps null elements. The native path always drops nulls, so `CometCollectShim` reads the field per Spark version (always `true` on 3.4-4.1) and `CometCollectList` / `CometCollectSet` report `Unsupported` when it is `false`, falling back to Spark. +## collect_set + +- Spark 3.4.3 (audited 2026-07-27): `CollectSet` extends `Collect[mutable.HashSet[Any]]`, returns `ArrayType(child.dataType, containsNull = false)`, ignores NULL inputs in `update()` (the same Hive-compatible semantics as `collect_list`), and yields an empty array as `defaultResult`. `nullable = false`. Unlike `CollectList` it overrides `checkInputDataTypes` and rejects any input whose type recursively contains a `MapType` (`UNSUPPORTED_INPUT_TYPE`). `convertToBufferElement` copies the value with `InternalRow.copyValue`, except for `BinaryType`, which is wrapped in an `UnsafeArrayData` so that byte arrays dedup by content rather than by identity. Deduplication is Scala `mutable.HashSet` equality on the boxed value, which for floating-point types is numeric `==`: repeated `NaN`s are each kept as separate elements, while `0.0` and `-0.0` collapse to one. Registered only as `collect_set` in `FunctionRegistry` (there is no second alias, unlike `collect_list`/`array_agg`). +- Spark 3.5.8 (audited 2026-07-27): identical to 3.4.3. +- Spark 4.0.1 (audited 2026-07-27): adds `with UnaryLike[Expression]` to the case class, and `checkInputDataTypes` additionally requires `UnsafeRowUtils.isBinaryStable(child.dataType)`, so non-default-collation strings are rejected along with maps. Deduplication semantics unchanged. +- Spark 4.1.1 (audited 2026-07-27): identical to 4.0.1. +- Comet implementation: the native side delegates to `datafusion_spark::function::aggregate::collect::SparkCollectSet`, which wraps `DistinctArrayAggAccumulator` with `ignore_nulls = true` in a `NullToEmptyListAccumulator` so a final NULL accumulator state becomes an empty array (matching Spark's `defaultResult`). As with [collect_list](#collect-list), the native return type is `List(Field, containsNull = true)` while Spark declares `containsNull = false`; nulls are filtered before insertion, so this is a schema-shape difference only. +- `CometCollectSet` reports `Incompatible` for float and double input when `spark.comet.exec.strictFloatingPoint=true`, because the native distinct comparison treats `NaN == NaN` and collapses repeated `NaN`s into a single element while Spark keeps each one. The native path for floating-point input is then opt-in via `spark.comet.expression.CollectSet.allowIncompatible=true`. All other input types are `Compatible`. +- Buffer shape: `CollectSet` is a `TypedImperativeAggregate`, so Spark's `aggBufferAttributes` declares the intermediate buffer as `BinaryType` (a serialized `HashSet`), while the native accumulator's `state_fields` is a `List`. `CometHashAggregateExec.adjustOutputForNativeState` rewrites the Comet-side Partial output to the list shape, and `supportsMixedPartialFinal` is left at the default `false`, so a Comet Partial cannot feed a Spark Final or the reverse. `QueryPlanSerde.hasNativeArrayBufferAgg` matches `CollectList`/`CollectSet` so that `CometExecRule.tagUnsafePartialAggregates` forces the whole chain back to Spark when a `PartialMerge` stage is present (the distinct-aggregate rewrite), a shape Comet cannot round-trip yet ([#4724](https://github.com/apache/datafusion-comet/issues/4724)). +- Spark 4.2 (preview): see the `ignoreNulls` note under [collect_list](#collect-list); `CometCollectSet` reports `Unsupported` for `RESPECT NULLS`. + ## median - Spark 3.4.3 (audited 2026-06-24): `Median(child)` is a `RuntimeReplaceableAggregate` with `replacement = Percentile(child, Literal(0.5))`. Catalyst rewrites `median(x)` to `percentile(x, 0.5)` before Comet sees the plan, so it is served by `CometPercentile`. diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql index e00d2b1969..54c2f52245 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql @@ -262,6 +262,49 @@ INSERT INTO cl_src_ts VALUES query SELECT grp, sort_array(collect_list(v)) FROM cl_src_ts GROUP BY grp ORDER BY grp +-- ============================================================ +-- Timestamp NTZ (with NULLs). TimestampNTZType is in +-- QueryPlanSerde.supportedDataType, so this runs through the +-- native SparkCollectList accumulator rather than falling back. +-- ============================================================ + +statement +CREATE TABLE cl_src_ts_ntz(v timestamp_ntz, grp string) USING parquet + +statement +INSERT INTO cl_src_ts_ntz VALUES + (timestamp_ntz '2024-01-01 00:00:00', 'a'), (timestamp_ntz '2024-06-15 12:30:00', 'a'), + (timestamp_ntz '2024-01-01 00:00:00', 'a'), (NULL, 'a'), + (timestamp_ntz '1970-01-01 00:00:00', 'b'), (NULL, 'b') + +query +SELECT grp, sort_array(collect_list(v)) FROM cl_src_ts_ntz GROUP BY grp ORDER BY grp + +-- ============================================================ +-- ANSI intervals (with NULLs). Neither YearMonthIntervalType nor +-- DayTimeIntervalType is in QueryPlanSerde.supportedDataType, so +-- these fall back to Spark and only the answer is checked. +-- ============================================================ + +statement +CREATE TABLE cl_src_interval(ym interval year to month, dt interval day to second, grp string) +USING parquet + +statement +INSERT INTO cl_src_interval VALUES + (INTERVAL '1-2' YEAR TO MONTH, INTERVAL '3 04:05:06' DAY TO SECOND, 'a'), + (INTERVAL '0-0' YEAR TO MONTH, INTERVAL '0 00:00:00' DAY TO SECOND, 'a'), + (INTERVAL '1-2' YEAR TO MONTH, INTERVAL '3 04:05:06' DAY TO SECOND, 'a'), + (NULL, NULL, 'a'), + (INTERVAL '-3-4' YEAR TO MONTH, INTERVAL '-5 06:07:08' DAY TO SECOND, 'b'), + (NULL, NULL, 'b') + +query spark_answer_only +SELECT grp, sort_array(collect_list(ym)) FROM cl_src_interval GROUP BY grp ORDER BY grp + +query spark_answer_only +SELECT grp, sort_array(collect_list(dt)) FROM cl_src_interval GROUP BY grp ORDER BY grp + -- ============================================================ -- Mixed with other aggregates -- ============================================================ diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index 3248101a2b..65c17ee74a 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -87,22 +87,34 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { // dependent PartialMerge/Final stages must also fall back rather than crash. See issue #4724 // for enabling the fully-native distinct path. import org.apache.spark.sql.functions.{collect_list, collect_set, sort_array} - // Non-native source (LocalTableScan): the buffer-producing Partial runs in Spark. + // Non-native source (LocalTableScan): the buffer-producing Partial runs in Spark, so the + // dependent PartialMerge/Final stages fall back via the buffer-source check in doConvert. + // tagUnsafePartialAggregates does not fire here because the Partial was never convertible. + def bufferSourceFallback(fn: String): String = + "Comet aggregate that merges intermediate buffers requires a Comet child aggregate when " + + "the intermediate buffer formats are incompatible with Spark. Incompatible aggregate " + + s"function(s): $fn" val df = Seq((1, 3, "a"), (1, 2, "b"), (3, 4, "c"), (3, 4, "c"), (3, 5, "d")) .toDF("x", "y", "z") - checkSparkAnswer( - df.groupBy(col("x")).agg(count_distinct(col("y")), sort_array(collect_list(col("z"))))) - checkSparkAnswer( - df.groupBy(col("x")).agg(count_distinct(col("y")), sort_array(collect_set(col("z"))))) - - // Native source (Parquet): the whole multi-stage distinct chain must still fall back to - // Spark consistently (issue #4724), rather than running a fully-native pipeline that crashes. + checkSparkAnswerAndFallbackReason( + df.groupBy(col("x")).agg(count_distinct(col("y")), sort_array(collect_list(col("z")))), + bufferSourceFallback("collect_list")) + checkSparkAnswerAndFallbackReason( + df.groupBy(col("x")).agg(count_distinct(col("y")), sort_array(collect_set(col("z")))), + bufferSourceFallback("collect_set")) + + // Native source (Parquet): the Partial would otherwise convert, so tagUnsafePartialAggregates + // must disable it and force the whole multi-stage distinct chain back to Spark (issue #4724), + // rather than running a fully-native pipeline that crashes. + val multiStageFallback = "Partial aggregate disabled: part of a multi-stage CollectList/" + + "CollectSet aggregate whose intermediate buffer cannot round-trip in Comet (issue #4724)" withParquetTable( Seq((1, 3, "a"), (1, 2, "b"), (3, 4, "c"), (3, 4, "c"), (3, 5, "d")), "t17616") { for (fn <- Seq("collect_list", "collect_set")) { - checkSparkAnswer( - sql(s"SELECT _1, count(distinct _2), sort_array($fn(_3)) FROM t17616 GROUP BY _1")) + checkSparkAnswerAndFallbackReasons( + sql(s"SELECT _1, count(distinct _2), sort_array($fn(_3)) FROM t17616 GROUP BY _1"), + Set(multiStageFallback, bufferSourceFallback(fn))) } } } @@ -293,6 +305,38 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("mixed engine collect_list: Comet partial + Spark final matches Spark") { + // collect_list has no Spark-compatible intermediate buffer (Spark declares BinaryType, the + // native accumulator produces a list), so a Spark final must never sit above a Comet partial. + // tagUnsafePartialAggregates is expected to cascade the fallback down to the partial. + val data = (0 until 100).map(i => (if (i % 11 == 0) None else Some(i), i % 7)) + withParquetTable(data, "tbl") { + withSQLConf( + CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.key -> "false", + CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { + checkSparkAnswer( + "SELECT _2, sort_array(collect_list(_1)), count(*) FROM tbl GROUP BY _2 ORDER BY _2") + } + } + } + + test("mixed engine collect_list: Spark partial + Comet final matches Spark") { + // The reverse direction: the Spark partial emits a serialized BinaryType buffer that Comet + // cannot read, and adjustOutputForNativeState would misinterpret it as a list if the final + // ran natively, so the final must fall back too. + val data = (0 until 100).map(i => (if (i % 11 == 0) None else Some(i), i % 7)) + withParquetTable(data, "tbl") { + withSQLConf( + CometConf.COMET_ENABLE_PARTIAL_HASH_AGGREGATE.key -> "false", + CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { + checkSparkAnswer( + "SELECT _2, sort_array(collect_list(_1)), count(*) FROM tbl GROUP BY _2 ORDER BY _2") + } + } + } + test("Aggregation without aggregate expressions should use correct result expressions") { withSQLConf( CometConf.COMET_ENABLED.key -> "true", From a34400e71a0d4459401b4b485b7718b3eb5d826a Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 27 Jul 2026 12:19:48 -0600 Subject: [PATCH 2/2] refactor: address simplification review feedback - test: assert on short discriminating fallback-reason fragments instead of copying the full message text from CometExecRule / operators - test: collapse the two mixed-engine tests into one parameterized pair and pin the Comet aggregate count at zero, so they fail if the partial/final cascade regresses instead of passing on any plan shape - test: drop the dead `ORDER BY`, which the answer helper already handles - test: fold the timestamp_ntz fixture into the existing timestamp table rather than duplicating its rows, matching the file's convention - test: pin the interval fallback with `expect_fallback` instead of `spark_answer_only`, so the block comment's claim is enforced - docs: move the shared buffer-shape note under collect_list, correct the `adjustOutputForNativeState` owner to CometBaseAggregate, and drop the redundant containsNull and Spark 4.2 restatements --- .../expression-audits/agg_funcs.md | 5 +- .../expressions/aggregate/collect_list.sql | 38 +++++------- .../comet/exec/CometAggregateSuite.scala | 62 ++++++++----------- 3 files changed, 43 insertions(+), 62 deletions(-) diff --git a/docs/source/contributor-guide/expression-audits/agg_funcs.md b/docs/source/contributor-guide/expression-audits/agg_funcs.md index 9c40ecb3c4..117841ea81 100644 --- a/docs/source/contributor-guide/expression-audits/agg_funcs.md +++ b/docs/source/contributor-guide/expression-audits/agg_funcs.md @@ -58,6 +58,7 @@ - Spark 4.0.1 (audited 2026-06-24): only structural change is adding `with UnaryLike[Expression]` to the case class (no behavior change). - Spark 4.1.1 (audited 2026-06-24): identical to 4.0.1. - Comet implementation: native side delegates to `datafusion_spark::function::aggregate::collect::SparkCollectList`, which wraps `ArrayAggAccumulator` with `ignore_nulls = true` and converts a final NULL accumulator state to an empty array (matching Spark's `defaultResult`). The native return type is `List(Field, containsNull = true)`, while Spark uses `containsNull = false`. Because nulls are filtered before insertion, no nulls actually appear in the array, so this is a schema-shape difference only and tests using `checkSparkAnswerAndOperator` accept it (the same pattern applies to [collect_set](#collect-set)). +- Buffer shape (applies equally to `collect_set`): both are `TypedImperativeAggregate`s, so Spark's `aggBufferAttributes` declares the intermediate buffer as `BinaryType` (serialized state) while the native accumulator's `state_fields` is a `List`. `CometBaseAggregate.adjustOutputForNativeState` rewrites the Comet-side Partial output to the list shape. Neither collector can split a Partial and Final across Comet and Spark, and a multi-stage distinct rewrite (which inserts a `PartialMerge` stage) forces the whole chain back to Spark ([#4724](https://github.com/apache/datafusion-comet/issues/4724)). - Spark 4.2 (preview): `CollectList` and `CollectSet` gain an `ignoreNulls` field (default `true`); `RESPECT NULLS` sets it to `false` and keeps null elements. The native path always drops nulls, so `CometCollectShim` reads the field per Spark version (always `true` on 3.4-4.1) and `CometCollectList` / `CometCollectSet` report `Unsupported` when it is `false`, falling back to Spark. ## collect_set @@ -66,10 +67,8 @@ - Spark 3.5.8 (audited 2026-07-27): identical to 3.4.3. - Spark 4.0.1 (audited 2026-07-27): adds `with UnaryLike[Expression]` to the case class, and `checkInputDataTypes` additionally requires `UnsafeRowUtils.isBinaryStable(child.dataType)`, so non-default-collation strings are rejected along with maps. Deduplication semantics unchanged. - Spark 4.1.1 (audited 2026-07-27): identical to 4.0.1. -- Comet implementation: the native side delegates to `datafusion_spark::function::aggregate::collect::SparkCollectSet`, which wraps `DistinctArrayAggAccumulator` with `ignore_nulls = true` in a `NullToEmptyListAccumulator` so a final NULL accumulator state becomes an empty array (matching Spark's `defaultResult`). As with [collect_list](#collect-list), the native return type is `List(Field, containsNull = true)` while Spark declares `containsNull = false`; nulls are filtered before insertion, so this is a schema-shape difference only. +- Comet implementation: the native side delegates to `datafusion_spark::function::aggregate::collect::SparkCollectSet`, which wraps `DistinctArrayAggAccumulator` with `ignore_nulls = true` in a `NullToEmptyListAccumulator` so a final NULL accumulator state becomes an empty array. The `containsNull` mismatch against Spark's declared output type, and its rationale, are identical to [collect_list](#collect-list). - `CometCollectSet` reports `Incompatible` for float and double input when `spark.comet.exec.strictFloatingPoint=true`, because the native distinct comparison treats `NaN == NaN` and collapses repeated `NaN`s into a single element while Spark keeps each one. The native path for floating-point input is then opt-in via `spark.comet.expression.CollectSet.allowIncompatible=true`. All other input types are `Compatible`. -- Buffer shape: `CollectSet` is a `TypedImperativeAggregate`, so Spark's `aggBufferAttributes` declares the intermediate buffer as `BinaryType` (a serialized `HashSet`), while the native accumulator's `state_fields` is a `List`. `CometHashAggregateExec.adjustOutputForNativeState` rewrites the Comet-side Partial output to the list shape, and `supportsMixedPartialFinal` is left at the default `false`, so a Comet Partial cannot feed a Spark Final or the reverse. `QueryPlanSerde.hasNativeArrayBufferAgg` matches `CollectList`/`CollectSet` so that `CometExecRule.tagUnsafePartialAggregates` forces the whole chain back to Spark when a `PartialMerge` stage is present (the distinct-aggregate rewrite), a shape Comet cannot round-trip yet ([#4724](https://github.com/apache/datafusion-comet/issues/4724)). -- Spark 4.2 (preview): see the `ignoreNulls` note under [collect_list](#collect-list); `CometCollectSet` reports `Unsupported` for `RESPECT NULLS`. ## median diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql index 54c2f52245..a3cfbf9eed 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/collect_list.sql @@ -247,43 +247,33 @@ query SELECT grp, sort_array(collect_list(v)) FROM cl_src_date GROUP BY grp ORDER BY grp -- ============================================================ --- Timestamp (with NULLs) +-- Timestamp / Timestamp NTZ (with NULLs). Both are in +-- QueryPlanSerde.supportedDataType, so both run through the +-- native SparkCollectList accumulator. -- ============================================================ statement -CREATE TABLE cl_src_ts(v timestamp, grp string) USING parquet +CREATE TABLE cl_src_ts(v timestamp, v_ntz timestamp_ntz, grp string) USING parquet statement INSERT INTO cl_src_ts VALUES - (TIMESTAMP '2024-01-01 00:00:00', 'a'), (TIMESTAMP '2024-06-15 12:30:00', 'a'), - (TIMESTAMP '2024-01-01 00:00:00', 'a'), (NULL, 'a'), - (TIMESTAMP '1970-01-01 00:00:00', 'b'), (NULL, 'b') + (TIMESTAMP '2024-01-01 00:00:00', TIMESTAMP_NTZ '2024-01-01 00:00:00', 'a'), + (TIMESTAMP '2024-06-15 12:30:00', TIMESTAMP_NTZ '2024-06-15 12:30:00', 'a'), + (TIMESTAMP '2024-01-01 00:00:00', TIMESTAMP_NTZ '2024-01-01 00:00:00', 'a'), + (NULL, NULL, 'a'), + (TIMESTAMP '1970-01-01 00:00:00', TIMESTAMP_NTZ '1970-01-01 00:00:00', 'b'), + (NULL, NULL, 'b') query SELECT grp, sort_array(collect_list(v)) FROM cl_src_ts GROUP BY grp ORDER BY grp --- ============================================================ --- Timestamp NTZ (with NULLs). TimestampNTZType is in --- QueryPlanSerde.supportedDataType, so this runs through the --- native SparkCollectList accumulator rather than falling back. --- ============================================================ - -statement -CREATE TABLE cl_src_ts_ntz(v timestamp_ntz, grp string) USING parquet - -statement -INSERT INTO cl_src_ts_ntz VALUES - (timestamp_ntz '2024-01-01 00:00:00', 'a'), (timestamp_ntz '2024-06-15 12:30:00', 'a'), - (timestamp_ntz '2024-01-01 00:00:00', 'a'), (NULL, 'a'), - (timestamp_ntz '1970-01-01 00:00:00', 'b'), (NULL, 'b') - query -SELECT grp, sort_array(collect_list(v)) FROM cl_src_ts_ntz GROUP BY grp ORDER BY grp +SELECT grp, sort_array(collect_list(v_ntz)) FROM cl_src_ts GROUP BY grp ORDER BY grp -- ============================================================ -- ANSI intervals (with NULLs). Neither YearMonthIntervalType nor -- DayTimeIntervalType is in QueryPlanSerde.supportedDataType, so --- these fall back to Spark and only the answer is checked. +-- the scan rejects the column and the query falls back to Spark. -- ============================================================ statement @@ -299,10 +289,10 @@ INSERT INTO cl_src_interval VALUES (INTERVAL '-3-4' YEAR TO MONTH, INTERVAL '-5 06:07:08' DAY TO SECOND, 'b'), (NULL, NULL, 'b') -query spark_answer_only +query expect_fallback(Unsupported ym of type YearMonthIntervalType) SELECT grp, sort_array(collect_list(ym)) FROM cl_src_interval GROUP BY grp ORDER BY grp -query spark_answer_only +query expect_fallback(Unsupported dt of type DayTimeIntervalType) SELECT grp, sort_array(collect_list(dt)) FROM cl_src_interval GROUP BY grp ORDER BY grp -- ============================================================ diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index 65c17ee74a..c7a55b9e9b 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -90,10 +90,7 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { // Non-native source (LocalTableScan): the buffer-producing Partial runs in Spark, so the // dependent PartialMerge/Final stages fall back via the buffer-source check in doConvert. // tagUnsafePartialAggregates does not fire here because the Partial was never convertible. - def bufferSourceFallback(fn: String): String = - "Comet aggregate that merges intermediate buffers requires a Comet child aggregate when " + - "the intermediate buffer formats are incompatible with Spark. Incompatible aggregate " + - s"function(s): $fn" + def bufferSourceFallback(fn: String): String = s"Incompatible aggregate function(s): $fn" val df = Seq((1, 3, "a"), (1, 2, "b"), (3, 4, "c"), (3, 4, "c"), (3, 5, "d")) .toDF("x", "y", "z") checkSparkAnswerAndFallbackReason( @@ -106,8 +103,8 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { // Native source (Parquet): the Partial would otherwise convert, so tagUnsafePartialAggregates // must disable it and force the whole multi-stage distinct chain back to Spark (issue #4724), // rather than running a fully-native pipeline that crashes. - val multiStageFallback = "Partial aggregate disabled: part of a multi-stage CollectList/" + - "CollectSet aggregate whose intermediate buffer cannot round-trip in Comet (issue #4724)" + val multiStageFallback = "multi-stage CollectList/CollectSet aggregate whose intermediate " + + "buffer cannot round-trip in Comet" withParquetTable( Seq((1, 3, "a"), (1, 2, "b"), (3, 4, "c"), (3, 4, "c"), (3, 5, "d")), "t17616") { @@ -305,36 +302,31 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } - test("mixed engine collect_list: Comet partial + Spark final matches Spark") { - // collect_list has no Spark-compatible intermediate buffer (Spark declares BinaryType, the - // native accumulator produces a list), so a Spark final must never sit above a Comet partial. - // tagUnsafePartialAggregates is expected to cascade the fallback down to the partial. - val data = (0 until 100).map(i => (if (i % 11 == 0) None else Some(i), i % 7)) - withParquetTable(data, "tbl") { - withSQLConf( - CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.key -> "false", - CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true", - CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { - checkSparkAnswer( - "SELECT _2, sort_array(collect_list(_1)), count(*) FROM tbl GROUP BY _2 ORDER BY _2") - } - } - } - - test("mixed engine collect_list: Spark partial + Comet final matches Spark") { - // The reverse direction: the Spark partial emits a serialized BinaryType buffer that Comet - // cannot read, and adjustOutputForNativeState would misinterpret it as a list if the final - // ran natively, so the final must fall back too. - val data = (0 until 100).map(i => (if (i % 11 == 0) None else Some(i), i % 7)) - withParquetTable(data, "tbl") { - withSQLConf( - CometConf.COMET_ENABLE_PARTIAL_HASH_AGGREGATE.key -> "false", - CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true", - CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { - checkSparkAnswer( - "SELECT _2, sort_array(collect_list(_1)), count(*) FROM tbl GROUP BY _2 ORDER BY _2") + // collect_list has no Spark-compatible intermediate buffer: Spark declares BinaryType while the + // native accumulator produces a list. Disabling either half of the aggregate must therefore + // cascade so that neither half runs in Comet - a Spark final cannot read a Comet-produced list, + // and adjustOutputForNativeState would misinterpret Spark's Binary buffer as a list if a Comet + // final ran above a Spark partial. count(*) is included so the aggregate also carries a + // buffer-compatible function, which is the shape most likely to be split by mistake. + Seq( + ("Comet partial + Spark final", CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE), + ("Spark partial + Comet final", CometConf.COMET_ENABLE_PARTIAL_HASH_AGGREGATE)).foreach { + case (name, disabledHalf) => + test(s"mixed engine collect_list: $name matches Spark") { + val data = (0 until 100).map(i => (if (i % 11 == 0) None else Some(i), i % 7)) + withParquetTable(data, "tbl") { + withSQLConf( + disabledHalf.key -> "false", + CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { + val df = sql("SELECT _2, sort_array(collect_list(_1)), count(*) FROM tbl GROUP BY _2") + checkSparkAnswer(df) + // Without the cascade the surviving half would still convert, so pinning this at zero + // is what keeps the test from passing on a regression. + assert(getNumCometHashAggregate(df) == 0) + } + } } - } } test("Aggregation without aggregate expressions should use correct result expressions") {