ChuckLin2025 opened a new pull request, #57592:
URL: https://github.com/apache/spark/pull/57592
### 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.
### 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`
(passes end-to-end through the PySpark runtime).
- Spark Connect parity check in `test_connect_function.py`.
- `ExpressionsSchemaSuite` regenerated `sql-expression-schema.md`.
--
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]
