From e39ab1a4a2bf2fe8cf7610db8dfbed1ce43e9d5f Mon Sep 17 00:00:00 2001 From: comphead Date: Wed, 1 Jul 2026 09:54:33 -0700 Subject: [PATCH 1/2] faet: support `collect_list` for `windows` --- .../spark/src/function/aggregate/collect.rs | 31 +- .../spark/aggregate/collect_window.slt | 398 ++++++++++++++++++ 2 files changed, 428 insertions(+), 1 deletion(-) create mode 100644 datafusion/sqllogictest/test_files/spark/aggregate/collect_window.slt diff --git a/datafusion/spark/src/function/aggregate/collect.rs b/datafusion/spark/src/function/aggregate/collect.rs index 5af0fd39cca07..2920dea980cc2 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 { + 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()) +} + // #[derive(Debug, PartialEq, Eq, Hash)] pub struct SparkCollectList { @@ -89,6 +102,10 @@ impl AggregateUDFImpl for SparkCollectList { data_type, ))) } + + fn default_value(&self, data_type: &DataType) -> Result { + empty_list_scalar(data_type) + } } // @@ -147,6 +164,10 @@ impl AggregateUDFImpl for SparkCollectSet { data_type, ))) } + + fn default_value(&self, data_type: &DataType) -> Result { + empty_list_scalar(data_type) + } } /// Wrapper accumulator that returns an empty list instead of NULL when all inputs are NULL. @@ -186,6 +207,14 @@ impl Accumulator for NullToEmptyListAccumulator { } } + 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() } 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 0000000000000..5661cb9432427 --- /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; From d4c2cc733eee0a46f75389d6daa665ae1f41918f Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 2 Jul 2026 17:07:25 -0700 Subject: [PATCH 2/2] faet: support `collect_list` for `windows` --- .../spark/src/function/aggregate/collect.rs | 25 ++++++++----------- 1 file changed, 11 insertions(+), 14 deletions(-) diff --git a/datafusion/spark/src/function/aggregate/collect.rs b/datafusion/spark/src/function/aggregate/collect.rs index 2920dea980cc2..310bc1c890657 100644 --- a/datafusion/spark/src/function/aggregate/collect.rs +++ b/datafusion/spark/src/function/aggregate/collect.rs @@ -94,12 +94,11 @@ impl AggregateUDFImpl for SparkCollectList { } fn accumulator(&self, acc_args: AccumulatorArgs) -> Result> { - 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(), ))) } @@ -156,12 +155,11 @@ impl AggregateUDFImpl for SparkCollectSet { } fn accumulator(&self, acc_args: AccumulatorArgs) -> Result> { - 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(), ))) } @@ -175,12 +173,12 @@ impl AggregateUDFImpl for SparkCollectSet { #[derive(Debug)] struct NullToEmptyListAccumulator { inner: T, - data_type: DataType, + list_type: DataType, } impl NullToEmptyListAccumulator { - 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 } } } @@ -200,8 +198,7 @@ impl Accumulator for NullToEmptyListAccumulator { fn evaluate(&mut self) -> Result { 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) } @@ -216,6 +213,6 @@ impl Accumulator for NullToEmptyListAccumulator { } fn size(&self) -> usize { - self.inner.size() + self.data_type.size() + self.inner.size() + self.list_type.size() } }