Skip to content

Commit 1bcaf3f

Browse files
authored
Add list_sum expression (#8676)
Adds a `vortex.list.sum` scalar function: the sum of each row's list elements, one output row per input row — the equivalent of DuckDB's `list_sum()` and DataFusion's `array_sum()`. Note that list entries that are empty or all null need to be nullified. However, `GroupedAccumulator` returns 0 in these cases, so we need to either: 1. Change the behavior of `GroupedAccumulator` and grouped sum (or even sum) kernels or 2. Do a second pass of list element validity and sizes, construct a mask for the sums, and apply it. The first option is probably best but can be applied in a followup PR. --------- Signed-off-by: Matt Katz <mhkatz97@gmail.com>
1 parent 24ad8fd commit 1bcaf3f

6 files changed

Lines changed: 810 additions & 0 deletions

File tree

vortex-array/Cargo.toml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -248,6 +248,10 @@ harness = false
248248
name = "list_length"
249249
harness = false
250250

251+
[[bench]]
252+
name = "list_sum"
253+
harness = false
254+
251255
[[bench]]
252256
name = "listview_rebuild"
253257
harness = false

vortex-array/benches/list_sum.rs

Lines changed: 192 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,192 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
//! Benchmarks for the `list_sum` scalar function over `List`, `ListView`, and
5+
//! `FixedSizeList` inputs.
6+
7+
#![expect(clippy::unwrap_used)]
8+
#![expect(clippy::cast_possible_truncation)]
9+
10+
use std::sync::LazyLock;
11+
12+
use divan::Bencher;
13+
use rand::RngExt;
14+
use rand::SeedableRng;
15+
use rand::distr::Uniform;
16+
use rand::rngs::StdRng;
17+
use vortex_array::ArrayRef;
18+
use vortex_array::Canonical;
19+
use vortex_array::IntoArray;
20+
use vortex_array::VortexSessionExecute;
21+
use vortex_array::arrays::BoolArray;
22+
use vortex_array::arrays::FixedSizeListArray;
23+
use vortex_array::arrays::ListArray;
24+
use vortex_array::arrays::ListViewArray;
25+
use vortex_array::arrays::PrimitiveArray;
26+
use vortex_array::expr::list_sum;
27+
use vortex_array::expr::root;
28+
use vortex_array::validity::Validity;
29+
use vortex_buffer::Buffer;
30+
use vortex_session::VortexSession;
31+
32+
fn main() {
33+
divan::main();
34+
}
35+
36+
static SESSION: LazyLock<VortexSession> = LazyLock::new(vortex_array::array_session);
37+
38+
const BASE_LIST_SIZE: usize = 8;
39+
40+
const SMALL: usize = 100;
41+
const MEDIUM: usize = 10_000;
42+
const LARGE: usize = 1_000_000;
43+
44+
/// A uniformly-random partition of `num_lists * BASE_LIST_SIZE` elements into `num_lists`
45+
/// lists, plus a validity mask with ~1/8 of lists null at random positions.
46+
fn random_lists(num_lists: usize) -> (Vec<i32>, Validity) {
47+
let mut rng = StdRng::seed_from_u64(num_lists as u64);
48+
let total = (num_lists * BASE_LIST_SIZE) as i32;
49+
50+
let cut_dist = Uniform::new_inclusive(0i32, total).unwrap();
51+
let mut cuts: Vec<i32> = (0..num_lists - 1).map(|_| rng.sample(cut_dist)).collect();
52+
cuts.sort_unstable();
53+
let mut sizes = Vec::with_capacity(num_lists);
54+
let mut prev = 0i32;
55+
for cut in cuts {
56+
sizes.push(cut - prev);
57+
prev = cut;
58+
}
59+
sizes.push(total - prev);
60+
61+
let null_dist = Uniform::new(0u32, 8).unwrap();
62+
let valid = (0..num_lists).map(|_| rng.sample(null_dist) != 0);
63+
(
64+
sizes,
65+
Validity::Array(BoolArray::from_iter(valid).into_array()),
66+
)
67+
}
68+
69+
/// `0..total` as `i32` elements, with ~1/8 nulled at random positions when `nullable`.
70+
fn make_elements(total: i32, nullable: bool) -> ArrayRef {
71+
if !nullable {
72+
return PrimitiveArray::from_iter(0..total).into_array();
73+
}
74+
let mut rng = StdRng::seed_from_u64(total as u64);
75+
let null_dist = Uniform::new(0u32, 8).unwrap();
76+
PrimitiveArray::from_option_iter((0..total).map(|v| (rng.sample(null_dist) != 0).then_some(v)))
77+
.into_array()
78+
}
79+
80+
/// A canonical `List<i32>` of `num_lists` variable-length lists, ~1/8 of them null.
81+
fn make_list(num_lists: usize, nullable_elements: bool) -> ArrayRef {
82+
let (sizes, validity) = random_lists(num_lists);
83+
let total: i32 = sizes.iter().sum();
84+
let elements = make_elements(total, nullable_elements);
85+
let offsets: Buffer<i32> = std::iter::once(0)
86+
.chain(sizes.iter().scan(0i32, |acc, &s| {
87+
*acc += s;
88+
Some(*acc)
89+
}))
90+
.collect();
91+
ListArray::try_new(elements, offsets.into_array(), validity)
92+
.unwrap()
93+
.into_array()
94+
}
95+
96+
/// A gapless `ListView<i32>` of `num_lists` variable-length lists, ~1/8 of them null.
97+
fn make_listview(num_lists: usize) -> ArrayRef {
98+
let (sizes, validity) = random_lists(num_lists);
99+
let total: i32 = sizes.iter().sum();
100+
let elements = make_elements(total, false);
101+
let offsets: Buffer<i32> = sizes
102+
.iter()
103+
.scan(0i32, |acc, &s| {
104+
let start = *acc;
105+
*acc += s;
106+
Some(start)
107+
})
108+
.collect();
109+
let sizes: Buffer<i32> = sizes.into_iter().collect();
110+
ListViewArray::new(elements, offsets.into_array(), sizes.into_array(), validity).into_array()
111+
}
112+
113+
/// A `FixedSizeList<i32, 8>` of `num_lists` lists, ~1/8 of them null.
114+
fn make_fsl(num_lists: usize) -> ArrayRef {
115+
let mut rng = StdRng::seed_from_u64(num_lists as u64);
116+
let total = (num_lists * BASE_LIST_SIZE) as i32;
117+
let elements = make_elements(total, false);
118+
let null_dist = Uniform::new(0u32, 8).unwrap();
119+
let valid = (0..num_lists).map(|_| rng.sample(null_dist) != 0);
120+
let validity = Validity::Array(BoolArray::from_iter(valid).into_array());
121+
FixedSizeListArray::new(elements, BASE_LIST_SIZE as u32, validity, num_lists).into_array()
122+
}
123+
124+
/// Apply `list_sum(root())` and materialize the result.
125+
fn run(bencher: Bencher, array: ArrayRef) {
126+
let expr = list_sum(root());
127+
bencher
128+
.with_inputs(|| (&array, SESSION.create_execution_ctx()))
129+
.bench_refs(|(array, ctx)| {
130+
array
131+
.clone()
132+
.apply(&expr)
133+
.unwrap()
134+
.execute::<Canonical>(ctx)
135+
.unwrap()
136+
});
137+
}
138+
139+
#[divan::bench]
140+
fn list_sum_small(bencher: Bencher) {
141+
run(bencher, make_list(SMALL, false));
142+
}
143+
144+
#[divan::bench]
145+
fn list_sum_medium(bencher: Bencher) {
146+
run(bencher, make_list(MEDIUM, false));
147+
}
148+
149+
#[divan::bench]
150+
fn list_sum_large(bencher: Bencher) {
151+
run(bencher, make_list(LARGE, false));
152+
}
153+
154+
#[divan::bench]
155+
fn list_sum_nullable_elements_medium(bencher: Bencher) {
156+
run(bencher, make_list(MEDIUM, true));
157+
}
158+
159+
#[divan::bench]
160+
fn list_sum_nullable_elements_large(bencher: Bencher) {
161+
run(bencher, make_list(LARGE, true));
162+
}
163+
164+
#[divan::bench]
165+
fn listview_sum_small(bencher: Bencher) {
166+
run(bencher, make_listview(SMALL));
167+
}
168+
169+
#[divan::bench]
170+
fn listview_sum_medium(bencher: Bencher) {
171+
run(bencher, make_listview(MEDIUM));
172+
}
173+
174+
#[divan::bench]
175+
fn listview_sum_large(bencher: Bencher) {
176+
run(bencher, make_listview(LARGE));
177+
}
178+
179+
#[divan::bench]
180+
fn fsl_sum_small(bencher: Bencher) {
181+
run(bencher, make_fsl(SMALL));
182+
}
183+
184+
#[divan::bench]
185+
fn fsl_sum_medium(bencher: Bencher) {
186+
run(bencher, make_fsl(MEDIUM));
187+
}
188+
189+
#[divan::bench]
190+
fn fsl_sum_large(bencher: Bencher) {
191+
run(bencher, make_fsl(LARGE));
192+
}

vortex-array/src/expr/exprs.rs

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ use vortex_error::VortexExpect;
1010
use vortex_error::vortex_panic;
1111
use vortex_utils::iter::ReduceBalancedIterExt;
1212

13+
use crate::aggregate_fn::NumericalAggregateOpts;
1314
use crate::dtype::DType;
1415
use crate::dtype::FieldName;
1516
use crate::dtype::FieldNames;
@@ -38,6 +39,7 @@ use crate::scalar_fn::fns::like::Like;
3839
use crate::scalar_fn::fns::like::LikeOptions;
3940
use crate::scalar_fn::fns::list_contains::ListContains;
4041
use crate::scalar_fn::fns::list_length::ListLength;
42+
use crate::scalar_fn::fns::list_sum::ListSum;
4143
use crate::scalar_fn::fns::literal::Literal;
4244
use crate::scalar_fn::fns::mask::Mask;
4345
use crate::scalar_fn::fns::merge::DuplicateHandling;
@@ -781,3 +783,27 @@ pub fn ext_storage(input: Expression) -> Expression {
781783
pub fn list_length(input: Expression) -> Expression {
782784
ListLength.new_expr(EmptyOptions, [input])
783785
}
786+
787+
// ---- ListSum ----
788+
789+
/// Creates an expression that sums the elements of each list for `List` and
790+
/// `FixedSizeList` inputs, akin to DuckDB's `list_sum()`.
791+
///
792+
/// Follows SQL `SUM` semantics per list: null lists, empty lists, and lists whose elements are
793+
/// all null yield null; null elements are skipped; integer and decimal overflow yields a null
794+
/// value. The result dtype follows `sum`'s widening rules and is always nullable. NaN float
795+
/// elements are skipped by default; see [`list_sum_opts`] for the NaN-including variant.
796+
///
797+
/// ```rust
798+
/// # use vortex_array::expr::{list_sum, root};
799+
/// let expr = list_sum(root());
800+
/// ```
801+
pub fn list_sum(input: Expression) -> Expression {
802+
ListSum.new_expr(NumericalAggregateOpts::default(), [input])
803+
}
804+
805+
/// Creates a [`list_sum`] expression with explicit [`NumericalAggregateOpts`], controlling
806+
/// whether NaN float elements are skipped (the default) or poison the list's sum to NaN.
807+
pub fn list_sum_opts(input: Expression, options: NumericalAggregateOpts) -> Expression {
808+
ListSum.new_expr(options, [input])
809+
}

0 commit comments

Comments
 (0)