You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/07/03 18:47:55 UTC
(datafusion) branch gh-readonly-queue/main/pr-23281-b7e7b514e9f93c0a013286511a36e86ffcf2d662 created (now 01c7527be0)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-23281-b7e7b514e9f93c0a013286511a36e86ffcf2d662
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 01c7527be0 spark: support `collect_list` `collect_set` for `windows` execution (#23281)
This branch includes the following new commits:
new 01c7527be0 spark: support `collect_list` `collect_set` for `windows` execution (#23281)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: spark: support `collect_list` `collect_set` for `windows` execution (#23281)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23281-b7e7b514e9f93c0a013286511a36e86ffcf2d662
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 01c7527be0b318e3517f4cebf150638724e7e5db
Author: Oleks V <co...@users.noreply.github.com>
AuthorDate: Fri Jul 3 11:47:33 2026 -0700
spark: support `collect_list` `collect_set` for `windows` execution (#23281)
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->
- Closes #23261 .
## Rationale for this change
Expand windows coverage to support `collect_set` `collect_list`
functions.
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->
## What changes are included in this PR?
<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
## Are these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->
## Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
-->
<!--
If there are any breaking changes to public APIs, please add the `api
change` label.
-->
---
datafusion/spark/src/function/aggregate/collect.rs | 56 ++-
.../test_files/spark/aggregate/collect_window.slt | 398 +++++++++++++++++++++
2 files changed, 439 insertions(+), 15 deletions(-)
diff --git a/datafusion/spark/src/function/aggregate/collect.rs b/datafusion/spark/src/function/aggregate/collect.rs
index 5af0fd39cc..310bc1c890 100644
--- a/datafusion/spark/src/function/aggregate/collect.rs
+++ b/datafusion/spark/src/function/aggregate/collect.rs
@@ -18,7 +18,7 @@
use arrow::array::ArrayRef;
use arrow::datatypes::{DataType, Field, FieldRef};
use datafusion_common::utils::SingleRowListArrayBuilder;
-use datafusion_common::{Result, ScalarValue};
+use datafusion_common::{Result, ScalarValue, internal_err};
use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
use datafusion_expr::utils::format_state_name;
use datafusion_expr::{Accumulator, AggregateUDFImpl, Signature, Volatility};
@@ -33,6 +33,19 @@ use std::sync::Arc;
// - returns an empty list when all inputs are NULL
// - does not support ordering
+/// Build an empty list `ScalarValue` for a `List(element_type)` data type.
+/// Used as the result for empty window frames and for groups whose inputs
+/// were all NULL, matching Spark's `collect_list` / `collect_set` semantics.
+fn empty_list_scalar(list_type: &DataType) -> Result<ScalarValue> {
+ let DataType::List(field) = list_type else {
+ return internal_err!(
+ "collect_list/collect_set expected List return type, got {list_type:?}"
+ );
+ };
+ let empty = arrow::array::new_empty_array(field.data_type());
+ Ok(SingleRowListArrayBuilder::new(empty).build_list_scalar())
+}
+
// <https://spark.apache.org/docs/latest/api/sql/index.html#collect_list>
#[derive(Debug, PartialEq, Eq, Hash)]
pub struct SparkCollectList {
@@ -81,14 +94,17 @@ impl AggregateUDFImpl for SparkCollectList {
}
fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
- let field = &acc_args.expr_fields[0];
- let data_type = field.data_type().clone();
+ let element_type = acc_args.expr_fields[0].data_type().clone();
let ignore_nulls = true;
Ok(Box::new(NullToEmptyListAccumulator::new(
- ArrayAggAccumulator::try_new(&data_type, ignore_nulls)?,
- data_type,
+ ArrayAggAccumulator::try_new(&element_type, ignore_nulls)?,
+ acc_args.return_type().clone(),
)))
}
+
+ fn default_value(&self, data_type: &DataType) -> Result<ScalarValue> {
+ empty_list_scalar(data_type)
+ }
}
// <https://spark.apache.org/docs/latest/api/sql/index.html#collect_set>
@@ -139,14 +155,17 @@ impl AggregateUDFImpl for SparkCollectSet {
}
fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
- let field = &acc_args.expr_fields[0];
- let data_type = field.data_type().clone();
+ let element_type = acc_args.expr_fields[0].data_type().clone();
let ignore_nulls = true;
Ok(Box::new(NullToEmptyListAccumulator::new(
- DistinctArrayAggAccumulator::try_new(&data_type, None, ignore_nulls)?,
- data_type,
+ DistinctArrayAggAccumulator::try_new(&element_type, None, ignore_nulls)?,
+ acc_args.return_type().clone(),
)))
}
+
+ fn default_value(&self, data_type: &DataType) -> Result<ScalarValue> {
+ empty_list_scalar(data_type)
+ }
}
/// Wrapper accumulator that returns an empty list instead of NULL when all inputs are NULL.
@@ -154,12 +173,12 @@ impl AggregateUDFImpl for SparkCollectSet {
#[derive(Debug)]
struct NullToEmptyListAccumulator<T: Accumulator> {
inner: T,
- data_type: DataType,
+ list_type: DataType,
}
impl<T: Accumulator> NullToEmptyListAccumulator<T> {
- pub fn new(inner: T, data_type: DataType) -> Self {
- Self { inner, data_type }
+ pub fn new(inner: T, list_type: DataType) -> Self {
+ Self { inner, list_type }
}
}
@@ -179,14 +198,21 @@ impl<T: Accumulator> Accumulator for NullToEmptyListAccumulator<T> {
fn evaluate(&mut self) -> Result<ScalarValue> {
let result = self.inner.evaluate()?;
if result.is_null() {
- let empty_array = arrow::array::new_empty_array(&self.data_type);
- Ok(SingleRowListArrayBuilder::new(empty_array).build_list_scalar())
+ empty_list_scalar(&self.list_type)
} else {
Ok(result)
}
}
+ fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
+ self.inner.retract_batch(values)
+ }
+
+ fn supports_retract_batch(&self) -> bool {
+ self.inner.supports_retract_batch()
+ }
+
fn size(&self) -> usize {
- self.inner.size() + self.data_type.size()
+ self.inner.size() + self.list_type.size()
}
}
diff --git a/datafusion/sqllogictest/test_files/spark/aggregate/collect_window.slt b/datafusion/sqllogictest/test_files/spark/aggregate/collect_window.slt
new file mode 100644
index 0000000000..5661cb9432
--- /dev/null
+++ b/datafusion/sqllogictest/test_files/spark/aggregate/collect_window.slt
@@ -0,0 +1,398 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+
+# http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+#######
+# Tests for Spark-compat collect_list / collect_set as WINDOW functions.
+# Spark semantics:
+# - NULL inputs are skipped (Hive collect_list/collect_set behavior).
+# - An empty frame (or one where all inputs were NULL) evaluates to []
+# rather than NULL (nullable = false in Spark's Collect aggregate).
+# - collect_list preserves frame order; collect_set deduplicates.
+# Validates that NullToEmptyListAccumulator forwards retract_batch
+# so the wrapped ArrayAggAccumulator / DistinctArrayAggAccumulator can
+# drive sliding window frames.
+#######
+
+statement ok
+CREATE TABLE t(ts INT, val TEXT) AS VALUES
+ (1, 'A'), (2, 'B'), (3, 'C'), (4, 'D'), (5, 'E');
+
+# Unbounded preceding frame — accumulator only sees update_batch.
+query ?
+SELECT collect_list(val)
+ OVER (ORDER BY ts ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)
+FROM t;
+----
+[A]
+[A, B]
+[A, B, C]
+[A, B, C, D]
+[A, B, C, D, E]
+
+# Bounded sliding ROWS frame — requires retract_batch.
+query ?
+SELECT collect_list(val)
+ OVER (ORDER BY ts ROWS BETWEEN 1 PRECEDING AND CURRENT ROW)
+FROM t;
+----
+[A]
+[A, B]
+[B, C]
+[C, D]
+[D, E]
+
+# Wider sliding window.
+query ?
+SELECT collect_list(val)
+ OVER (ORDER BY ts ROWS BETWEEN 2 PRECEDING AND CURRENT ROW)
+FROM t;
+----
+[A]
+[A, B]
+[A, B, C]
+[B, C, D]
+[C, D, E]
+
+# Centered sliding window with PRECEDING + FOLLOWING.
+query ?
+SELECT collect_list(val)
+ OVER (ORDER BY ts ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING)
+FROM t;
+----
+[A, B]
+[A, B, C]
+[B, C, D]
+[C, D, E]
+[D, E]
+
+# Unbounded both sides — every row sees the full input.
+query ?
+SELECT collect_list(val)
+ OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING)
+FROM t;
+----
+[A, B, C, D, E]
+[A, B, C, D, E]
+[A, B, C, D, E]
+[A, B, C, D, E]
+[A, B, C, D, E]
+
+# Single-row frame.
+query ?
+SELECT collect_list(val)
+ OVER (ORDER BY ts ROWS BETWEEN CURRENT ROW AND CURRENT ROW)
+FROM t;
+----
+[A]
+[B]
+[C]
+[D]
+[E]
+
+# Empty leading frame on the first row — Spark returns [].
+query ?
+SELECT collect_list(val)
+ OVER (ORDER BY ts ROWS BETWEEN 2 PRECEDING AND 1 PRECEDING)
+FROM t;
+----
+[]
+[A]
+[A, B]
+[B, C]
+[C, D]
+
+# Empty trailing frame on the last row — Spark returns [].
+query ?
+SELECT collect_list(val)
+ OVER (ORDER BY ts ROWS BETWEEN 1 FOLLOWING AND 2 FOLLOWING)
+FROM t;
+----
+[B, C]
+[C, D]
+[D, E]
+[E]
+[]
+
+#######
+# NULL handling — Spark's collect_list skips NULL inputs.
+#######
+
+statement ok
+CREATE TABLE t_nulls(ts INT, val TEXT) AS VALUES
+ (1, 'A'), (2, NULL), (3, 'C'), (4, NULL), (5, 'E');
+
+# NULLs filtered out of the materialized list, but the row still emits one entry.
+query ?
+SELECT collect_list(val)
+ OVER (ORDER BY ts ROWS BETWEEN 1 PRECEDING AND CURRENT ROW)
+FROM t_nulls;
+----
+[A]
+[A]
+[C]
+[C]
+[E]
+
+# Wider frame.
+query ?
+SELECT collect_list(val)
+ OVER (ORDER BY ts ROWS BETWEEN 2 PRECEDING AND CURRENT ROW)
+FROM t_nulls;
+----
+[A]
+[A]
+[A, C]
+[C]
+[C, E]
+
+# All-NULL frame collapses to [].
+statement ok
+CREATE TABLE t_allnull(ts INT, val TEXT) AS VALUES
+ (1, NULL), (2, NULL), (3, NULL);
+
+query ?
+SELECT collect_list(val)
+ OVER (ORDER BY ts ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)
+FROM t_allnull;
+----
+[]
+[]
+[]
+
+#######
+# PARTITION BY — each partition starts with fresh accumulator state.
+#######
+
+statement ok
+CREATE TABLE t_parts(grp INT, ts INT, val TEXT) AS VALUES
+ (1, 1, 'A'), (1, 2, 'B'), (1, 3, 'C'),
+ (2, 1, 'X'), (2, 2, 'Y'), (2, 3, 'Z');
+
+query I?
+SELECT grp, collect_list(val)
+ OVER (PARTITION BY grp ORDER BY ts
+ ROWS BETWEEN 1 PRECEDING AND CURRENT ROW)
+FROM t_parts
+ORDER BY grp, ts;
+----
+1 [A]
+1 [A, B]
+1 [B, C]
+2 [X]
+2 [X, Y]
+2 [Y, Z]
+
+#######
+# RANGE frame with value gaps — exercises multi-row retract.
+#######
+
+statement ok
+CREATE TABLE t_range(ts INT, val TEXT) AS VALUES
+ (1, 'A'), (2, 'B'), (3, 'C'), (4, 'D'), (100, 'E');
+
+query ?
+SELECT collect_list(val)
+ OVER (ORDER BY ts RANGE BETWEEN 2 PRECEDING AND 2 FOLLOWING)
+FROM t_range;
+----
+[A, B, C]
+[A, B, C, D]
+[A, B, C, D]
+[B, C, D]
+[E]
+
+#######
+# GROUPS frame — rows tied on ORDER BY are processed together.
+#######
+
+statement ok
+CREATE TABLE t_groups(ts INT, val TEXT) AS VALUES
+ (1, 'A'), (1, 'B'), (2, 'C'), (2, 'D'), (3, 'E');
+
+query ?
+SELECT collect_list(val)
+ OVER (ORDER BY ts GROUPS BETWEEN 1 PRECEDING AND CURRENT ROW)
+FROM t_groups;
+----
+[A, B]
+[A, B]
+[A, B, C, D]
+[A, B, C, D]
+[C, D, E]
+
+#######
+# Integer-typed input — guards against type-specific regressions.
+#######
+
+statement ok
+CREATE TABLE t_int(ts INT, val INT) AS VALUES
+ (1, 10), (2, 20), (3, 30), (4, 40), (5, 50);
+
+query ?
+SELECT collect_list(val)
+ OVER (ORDER BY ts ROWS BETWEEN 1 PRECEDING AND CURRENT ROW)
+FROM t_int;
+----
+[10]
+[10, 20]
+[20, 30]
+[30, 40]
+[40, 50]
+
+#######
+# collect_set as a WINDOW function.
+# array_sort wraps the result because the underlying HashMap iteration
+# order is not deterministic.
+#######
+
+statement ok
+CREATE TABLE t_set(ts INT, val TEXT) AS VALUES
+ (1,'A'),(2,'A'),(3,'B'),(4,'C'),(5,'B');
+
+# Sliding ROWS frame, 2 PRECEDING.
+# Frame contents per row:
+# [A] -> {A}
+# [A,A] -> {A}
+# [A,A,B] -> {A,B}
+# [A,B,C] -> {A,B,C}
+# [B,C,B] -> {B,C}
+query ?
+SELECT array_sort(collect_set(val)
+ OVER (ORDER BY ts ROWS BETWEEN 2 PRECEDING AND CURRENT ROW))
+FROM t_set;
+----
+[A]
+[A]
+[A, B]
+[A, B, C]
+[B, C]
+
+# Narrower frame.
+query ?
+SELECT array_sort(collect_set(val)
+ OVER (ORDER BY ts ROWS BETWEEN 1 PRECEDING AND CURRENT ROW))
+FROM t_set;
+----
+[A]
+[A]
+[A, B]
+[B, C]
+[B, C]
+
+# Unbounded preceding — every distinct seen so far stays in.
+query ?
+SELECT array_sort(collect_set(val)
+ OVER (ORDER BY ts ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW))
+FROM t_set;
+----
+[A]
+[A]
+[A, B]
+[A, B, C]
+[A, B, C]
+
+# collect_set with NULLs — NULL never enters the set.
+statement ok
+CREATE TABLE t_set_nulls(ts INT, val TEXT) AS VALUES
+ (1,'A'),(2,NULL),(3,'A'),(4,NULL),(5,'B');
+
+query ?
+SELECT array_sort(collect_set(val)
+ OVER (ORDER BY ts ROWS BETWEEN 1 PRECEDING AND CURRENT ROW))
+FROM t_set_nulls;
+----
+[A]
+[A]
+[A]
+[A]
+[B]
+
+# collect_set with PARTITION BY — partition isolation.
+statement ok
+CREATE TABLE t_set_parts(grp INT, ts INT, val TEXT) AS VALUES
+ (1, 1, 'A'), (1, 2, 'A'), (1, 3, 'B'),
+ (2, 1, 'B'), (2, 2, 'C'), (2, 3, 'C');
+
+query I?
+SELECT grp, array_sort(collect_set(val)
+ OVER (PARTITION BY grp ORDER BY ts
+ ROWS BETWEEN 1 PRECEDING AND CURRENT ROW))
+FROM t_set_parts
+ORDER BY grp, ts;
+----
+1 [A]
+1 [A]
+1 [A, B]
+2 [B]
+2 [B, C]
+2 [C]
+
+# Empty leading frame on the first row — Spark returns [].
+query ?
+SELECT array_sort(collect_set(val)
+ OVER (ORDER BY ts ROWS BETWEEN 2 PRECEDING AND 1 PRECEDING))
+FROM t_set;
+----
+[]
+[A]
+[A]
+[A, B]
+[B, C]
+
+# All-NULL window — set is empty.
+query ?
+SELECT array_sort(collect_set(val)
+ OVER (ORDER BY ts ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW))
+FROM t_allnull;
+----
+[]
+[]
+[]
+
+#######
+# Cleanup
+#######
+
+statement ok
+DROP TABLE t;
+
+statement ok
+DROP TABLE t_nulls;
+
+statement ok
+DROP TABLE t_allnull;
+
+statement ok
+DROP TABLE t_parts;
+
+statement ok
+DROP TABLE t_range;
+
+statement ok
+DROP TABLE t_groups;
+
+statement ok
+DROP TABLE t_int;
+
+statement ok
+DROP TABLE t_set;
+
+statement ok
+DROP TABLE t_set_nulls;
+
+statement ok
+DROP TABLE t_set_parts;
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org