[SPARK-58399][SQL][PYTHON] Add collect_union aggregate function

[ChuckLin2025]

What changes were proposed in this pull request?

This PR adds a new aggregate function collect_union that takes an
array-typed column and returns the distinct union of the elements of the
arrays across rows.

collect_union(col: array<T>) : array<T>

It is equivalent to array_distinct(flatten(collect_list(col))), but the
aggregation buffer holds only the distinct elements (a HashSet), so its
size is bounded by the element universe rather than by the number of input
rows. This avoids buffering every row's whole array, which for a hot
grouping key can grow without bound.

The function is implemented as a Collect[mutable.HashSet[Any]]
(sibling of collect_set); the only material difference is that update
iterates the input array and adds each non-null element, and the result
element type is the input array's element type. NULL input arrays and NULL
elements are skipped, following collect_set semantics.

Added across the usual surfaces: Catalyst expression + registry, the Scala
DataFrame API, and PySpark (classic + Spark Connect). Spark Connect needs
no protocol change: the function travels as a generic UnresolvedFunction
resolved against the registry.

Why are the changes needed?

There is no built-in aggregate that unions the elements of an array column
across rows into a single distinct array. The workaround
array_distinct(flatten(collect_list(arr))) buffers every row's whole
array before de-duplicating, which can OOM on skewed keys. collect_union
de-duplicates during aggregation, keeping the buffer bounded by the
distinct-element universe.

Note that collect_set cannot replace this. collect_set(element) over
explode(col) does bound the buffer, but it stops being a plain aggregate:
each array column must be exploded and grouped on its own and then joined
back on the grouping keys. A query needing the distinct union of N array
columns therefore pays N explodes + N joins purely to work around the
missing array-input aggregate. collect_union keeps the bounded buffer
while staying an ordinary aggregate, so multiple array columns aggregate
together in one GROUP BY with no join.

Industry precedent: BigQuery (GoogleSQL) already supports this style of
array-input aggregate (ARRAY_CONCAT_AGG, which concatenates arrays across
rows; distinct is then applied), whereas PostgreSQL has no dedicated
function and users fall back to array_agg(DISTINCT ...) over unnest(...)
(the analogue of the explode + collect_set workaround above). collect_union
gives Spark a first-class, bounded-buffer form of this operation.

Does this PR introduce any user-facing change?

Yes. It adds a new SQL function collect_union and the corresponding
functions.collect_union in the Scala and Python DataFrame APIs.

How was this patch tested?

• New collect_union function case in DataFrameAggregateSuite (distinct
  union, NULL array, NULL element, per-group, empty result). Full suite:
  170 tests, all pass.
• New test_collect_union in python/pyspark/sql/tests/test_functions.py
  covering array<int>, array<string>, array<double>, NULL elements,
  and array<struct> (passes end-to-end through the PySpark runtime).
• Spark Connect parity check in test_connect_function.py.
• ExpressionsSchemaSuite regenerated sql-expression-schema.md.

[ChuckLin2025]

Hi @cloud-fan @srielau, could you help me review this PR.

Today for this kind of query:
select key1, key2, array_distinct(flatten(collect_list(array_col1))), array_distinct(flatten(collect_list(array_col2))) FROM ...

Spark can't run it in an effiecnt way, we have to buffer a huge array_of_array and perform array distinct on it later. Also it can't be rewrote by collect_set in a graceful way. An additional join cost need to be paid, like

with col1 as (
   select key1, key2, collect_set(col) as distinct_col1
   FROM ... LATERAL VIEW explode(array_col1) e AS col
),
col2 as (
   select key1, key2, collect_set(col) as distinct_col2
   FROM ... LATERAL VIEW explode(array_col2) e AS col
)
select key1, key2, distinct_col1, distinct_col2
from col1 join col2 
using (key1, key2) 

So I propose to introdue this collect_union function to the collect familiy. Also this is already supported in googleSQL by ARRAY_CONCAT_AGG(distinct ...).

[cloud-fan]

1 blocking, 0 non-blocking, 1 nit.
The bounded aggregation design is sound, but nullable-element semantics must be reconciled with the documented equivalent expression before merge.

Correctness (1)

• sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/collect.scala:396: Preserve one null element in the aggregation buffer, and declare the result element as nullable. The documented equivalent expression keeps one null from an input such as array(1, NULL), while this branch drops it, so the new function currently changes results for nullable arrays despite claiming equivalence. The existing null-element test locks in the mismatched result and should be updated with the implementation. -- see inline

Nits: 1 minor item (see inline comments).

Verification

I traced nullable elements through both implementations. CollectUnion.update skips a null element, while ArrayDistinct explicitly retains one null after collect_list and flatten, so array(1, NULL) produces different results despite the stated equivalence. I also checked the aggregate buffer serialization and stable binary/float keying against CollectSet.

[cloud-fan]

The documented equivalent expression preserves one null element, so this aggregate should do the same. For array(1, NULL), array_distinct(flatten(collect_list(...))) returns [1, NULL], but this branch drops the null and the result metadata declares containsNull = false. Please retain one null in the buffer, update the result nullability, and change the null-element test accordingly.

[ChuckLin2025]

It's a little tricky here:

┌────────────────────────────────────────────┬────────────────┬──────────────┬─────────────────────────────────┐
│                 Expression                 │     Input      │    Result    │           Keeps NULL?           │
├────────────────────────────────────────────┼────────────────┼──────────────┼─────────────────────────────────┤
│ array_distinct(array(1, NULL, NULL))       │ one row        │ [1, None]    │ yes — keeps one                 │
├────────────────────────────────────────────┼────────────────┼──────────────┼─────────────────────────────────┤
│ array_distinct(flatten(collect_list(...))) │ [[1,NULL],[2]] │ [1, 2, None] │ yes — the documented equivalent │
├────────────────────────────────────────────┼────────────────┼──────────────┼─────────────────────────────────┤
│ collect_union (mine)                       │ [[1,NULL],[2]] │ [1, 2]       │ no — drops it                   │
├────────────────────────────────────────────┼────────────────┼──────────────┼─────────────────────────────────┤
│ collect_set(1,NULL,2)                      │ scalars        │ [1, 2]       │ no                              │
└────────────────────────────────────────────┴────────────────┴──────────────┴─────────────────────────────────┘```

[ChuckLin2025]

Looks like array_distinct is already inconsistnt with collect_set. and we need to make a decision here:

• Pick array_distinct's so that collect_union can fully replace array_distinct(flatten(collect_list(...))). Like its document claims right now;
• Or pick collect_set's, as they are in the same funciton family. Different null hanlding will bring confusing easier.

[ChuckLin2025]

Thanks -- I went with supporting both IGNORE/RESPECT NULLS, defaulting to IGNORE, matching collect_list/collect_set:

• Default (IGNORE NULLS): null elements dropped -- consistent with collect_set and the rest of the collect family.
• RESPECT NULLS: keeps a single null element, and in that mode collect_union is exactly array_distinct(flatten(collect_list(...))) (verified: [[1,NULL],[2]] -> {1, 2, NULL}).

This fixes the inconsistency: the buffer retains one null under RESPECT NULLS and the result's containsNull is derived from ignoreNulls (!ignoreNulls), so nullability metadata is correct in both modes. Wired CollectUnion into FunctionResolution.applyIgnoreNulls next to CollectList/CollectSet; updated the docs so the array_distinct(flatten(collect_list)) equivalence is scoped to RESPECT NULLS; and split the null-element test into IGNORE (default) and RESPECT cases. Also applied your line-351 comment suggestion.

[cloud-fan]

2 addressed, 0 remaining, 1 new. (1 newly introduced, 0 late catches, 0 previously raised.)
0 blocking, 0 non-blocking, 1 nit.
The earlier null-semantics concerns are addressed; only one terminology nit remains in a schema comment.

Nits: 1 minor item (see inline comments).

Verification

I traced element insertion, buffer serialization, floating-point and binary normalization, and null handling against CollectSet. I also verified that FunctionResolution forwards the SQL null-handling flag and that the Scala, Python, and Connect entry points converge on the registered Catalyst expression.

[cloud-fan]

1 addressed, 0 remaining, 0 new.
0 blocking, 0 non-blocking, 0 nits.
The earlier schema-comment nit is addressed, and the current implementation and public entry points are ready from this review's perspective.

Verification

I traced element insertion, set merging and serialization, binary/floating-point normalization, null-clause propagation, result nullability, registry resolution, and the Scala, Python, and Connect entry points. I also compared the implementation with CollectSet and rechecked every prior review thread against the current head.

[cloud-fan]

Merge Summary:

• merged into master d5ec658
• merged into branch-4.x eeab9be

Posted by merge_spark_pr.py