From ecbd0074198ca5c2f76f8f3196ec2af547071b73 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 6 Jul 2026 15:01:28 -0700 Subject: [PATCH 1/2] Add list_sum expression design doc Design for a vortex.list.sum scalar function: survey of DuckDB/DataFusion/ ClickHouse/Polars semantics, reuse of the grouped-sum machinery (GroupedAccumulator), SQL NULL-on-empty semantics for engine pushdown parity, NaN handling via NumericalAggregateOpts, and version-gated DataFusion integration (array_sum absent from the pinned DF 54). Signed-off-by: Matt Katz --- .../src/scalar_fn/fns/list_sum/design.md | 429 ++++++++++++++++++ 1 file changed, 429 insertions(+) create mode 100644 vortex-array/src/scalar_fn/fns/list_sum/design.md diff --git a/vortex-array/src/scalar_fn/fns/list_sum/design.md b/vortex-array/src/scalar_fn/fns/list_sum/design.md new file mode 100644 index 00000000000..4588bddcfa1 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/list_sum/design.md @@ -0,0 +1,429 @@ +# `list_sum` — Per-List Element Summation as a Vortex Expression + +*Goal: support `list_sum(l)` — the sum of each row's list elements, one output row per input +row — as a first-class Vortex scalar function, with DuckDB and DataFusion pushdown.* + +--- + +## 1. Motivation + +`list_sum` is the simplest member of the list-aggregate family (`list_min`, `list_max`, +`list_avg`, `list_count`, …) and the natural second list scalar function after `list_length`: + +1. **It is the canonical "reduce each list" primitive.** Every mainstream engine ships it + (DuckDB `list_sum`, DataFusion `array_sum` — new June 2026, ClickHouse `arraySum`, Polars + `list.sum()`), and it is the shape user queries actually take over nested telemetry/feature + columns: `SELECT id, list_sum(scores) FROM t`. +2. **Vortex already has the hard part.** The grouped-aggregate machinery + (`GroupedAccumulator`, `aggregate_fn/accumulator_grouped.rs`) computes one aggregate per + list using encoding-specialized kernels — a `ListArray` *is* a `GroupedArray*` where each + list is a group. `list_sum` is mostly a thin scalar-fn adapter over machinery that exists + and is tested today, plus one semantic fix-up (§4.2). +3. **It opens the pushdown door for the whole family.** The vtable shape, typing rule, + semantics decisions, and converter arms designed here transfer directly to + `list_min`/`list_max`/`list_avg` (§8) — deciding the semantics questions once, carefully, + is most of the value of this doc. + +## 2. Survey: how other engines implement it + +Verified against engine sources (and DuckDB v1.5.3 empirically); full citations in the appendix. + +| | DuckDB `list_sum` | DataFusion `array_sum` | ClickHouse `arraySum` | Polars `list.sum()` | +| --- | --- | --- | --- | --- | +| NULL list row | NULL | NULL | n/a (non-Nullable arrays) | NULL | +| Empty list | **NULL** (SQL `SUM` of zero rows) | **NULL** | **0** (zero-init accumulator) | **0** | +| NULL elements | skipped | skipped; all-NULL list → NULL | unsupported (`ILLEGAL_COLUMN`; users pre-map with lambdas) | skipped | +| `INT32` list result type | `HUGEINT` (int128) | `Float64` (always) | `Int64` | `Int32` (no promotion) | +| NaN in a float list | poisons to NaN (IEEE) | poisons to NaN (IEEE) | poisons to NaN | poisons to NaN | +| Strategy | per-row **aggregate state** + batched scattered update; reuses any aggregate | per-row loop over offsets into a flat child pre-cast to `Float64` | pure segmented reduction, one linear pass | segmented reduction fast path (no nulls) / per-row amortized `Series::sum` | + +### DuckDB — generic "list × any aggregate" + +`list_sum` is not C++ at all: it is a SQL macro `(l) AS list_aggr(l, 'sum')` +(`src/catalog/default/default_functions.cpp`), like `list_min`/`list_max`/`list_avg`/…. +`list_aggregate` (`extension/core_functions/scalar/list/list_aggregates.cpp`) binds the named +aggregate from the catalog at bind time, stashes the `BoundAggregateExpression` in bind data, +and at execution allocates one aggregate state per row, batching element→state updates through +a selection vector 2048 at a time — grouped hash-aggregation mechanics reused verbatim. All +semantics fall out of SQL `SUM`: empty list → state never updated → NULL; null elements +skipped; every integer width sums into `HUGEINT` so overflow is practically unreachable (and +`sum(HUGEINT)` itself throws on wrap). Design goal was "any aggregate over a list" (the same +file binds histogram aggregates for `list_distinct`), not a fast sum. + +### DataFusion — specialized kernel, brand new + +`array_sum` (alias `list_sum`) landed in `datafusion-functions-nested` in **June 2026** +(PR [#22542], merged 2026-06-08, closing [#7214] from 2023). Signature coercion casts *every* +element type to `Float64` up front and `return_type` is unconditionally `Float64` — much +cruder than DuckDB. The kernel is a per-row loop over `value_offsets()` into the flat +`Float64` child with a per-element validity branch; a row with no valid elements (empty or +all-NULL) appends null. The code comment pins the semantics to "PostgreSQL array_sum, DuckDB +list_sum, Spark aggregate": SQL `SUM` over the unnested elements. + +**Version gate (verified):** Vortex pins DataFusion **54** (workspace `Cargo.toml:140`), and +the `datafusion-functions-nested-54.0.0` crate Vortex builds against has **no `array_sum`** +module. The DataFusion converter arm (§6.2) is therefore designed here but lands only when +Vortex bumps past the release containing #22542. + +### ClickHouse — the branchless extreme + +`arrayAggregation.cpp` implements `arrayMin/Max/Sum/Average/Product` as one segmented +reduction template: a single linear pass over the flat data column between offsets, +`Int64`/`UInt64`/`Float64` widening, empty array → value-initialized accumulator → **0**. +Nullable children are rejected outright (`ILLEGAL_COLUMN`) — the price of a branchless, +vectorizable inner loop. + +### Polars — dispatch on child nullability + +`sum_mean.rs`: when the inner values have no nulls, a textbook segmented reduction +(`offsets.windows(2).map(|w| sum_slice(&values[w[0]..w[1]]))`) with the *list* validity cloned +onto the output; otherwise a per-row amortized `Series::sum` slow path. Empty list → 0. +Minimal promotion (`Int32` sums into `Int32`, wrap-on-overflow risk). + +### Takeaways + +1. **Two viable architectures**: DuckDB's generic aggregate-reuse vs a specialized segmented + reduction (ClickHouse/Polars/DataFusion). Vortex's `GroupedAccumulator` is exactly the + DuckDB architecture — already built, already encoding-specialized (§3). +2. **The semantics are not settled across engines**, and the fork that matters for Vortex is + empty/all-null lists: NULL (SQL camp: DuckDB, DataFusion) vs 0 (kernel camp: ClickHouse, + Polars). *Both engines Vortex integrates with are in the SQL camp* — pushdown must not + change query results, which decides §4.2. +3. **NaN**: every surveyed engine lets NaN poison the sum (IEEE float addition). Vortex's own + default is NaN-*skipping* (§3) — a second, quieter divergence the options design must carry. + +## 3. Where Vortex is today + +- **Scalar functions** live in `vortex-array/src/scalar_fn/`; an expression node is + `Expression { scalar_fn: ScalarFnRef, children }` (`expr/expression.rs:27`) over a + `ScalarFnVTable` impl (`scalar_fn/vtable.rs:39`). `list_length` + (`scalar_fn/fns/list_length.rs`) is the closest template: unary, `EmptyOptions`, + `execute_until::` (`:176`) to reach `List`/`ListView`/`FixedSizeList`, answer from + structure, validity passthrough (`:101`), registered in `scalar_fn/session.rs:72`, builder + in `expr/exprs.rs:765`. +- **`Sum` is an aggregate function** (`aggregate_fn/fns/sum/mod.rs`), not a compute function. + Its typing rule (`Sum::return_dtype`, `sum/mod.rs:101`) is the widening `list_sum` should + inherit verbatim: `Bool → U64?`, unsigned → `U64?`, signed → `I64?`, floats → `F64?`, + decimal → `min(precision + 10, MAX_PRECISION)?` (the Spark/DataFusion heuristic, cited in + the source) — **always `Nullable`, because overflow yields a null sum value**, not an error. +- **The grouped path is the implementation lever.** `GroupedArray` + (`aggregate_fn/accumulator_grouped.rs:45`) is `ListView | FixedSizeList` with + `elements()`, `group_ranges(ctx)` (`(offset, size)` per group, `:74`), `group_validity(ctx)` + (`:82`). `GroupedAccumulator::::try_new(Sum, opts, elem_dtype)` + + `DynGroupedAccumulator::accumulate_list(&list_array, ctx)` + `finish()` produces **one sum + per list**: `accumulate_list` canonicalizes the input list internally and dispatches through + the session's grouped kernel registry + (`aggregate_fn/session.rs:198`, `register_grouped_encoding_kernel` at `:210`), where + `PrimitiveGroupedSumEncodingKernel` (`sum/grouped.rs:26`) → `try_grouped_sum` (`:48`) → + `grouped_sum` (`:70`) reuses the scalar primitive-sum reductions so per-group overflow/NaN + semantics match scalar `sum` exactly. The element validity mask is materialized **once** + (`execute_mask`) and sliced per group via its contiguous valid runs (`sum_masked_group`, + `:131`) — no per-element `is_valid` calls, per the repo's hot-loop guidance. +- **Current grouped-sum semantics** (verified in `collect_sums`, `sum/grouped.rs:106`): + null group → null; overflow → null; **empty group → `Some(0)`** (accumulator default, never + updated); **all-null-elements group → `Some(0)`** (`AllOr::None` ⇒ no update). So Vortex + today sits in the ClickHouse/Polars camp, opposite its two integration targets. +- **NaN handling is an option, not a constant**: `NumericalAggregateOpts { skip_nans }` + (`aggregate_fn/vtable.rs:183`), default `skip_nans()` (`:193`, NaNs treated as missing); + `include_nans()` (`:201`) gives the engines' poison-to-NaN behavior. It already has proto + serialization (`:206`). +- Nothing outside `aggregate_fn/` consumes `GroupedArray` yet — `list_sum` would be the first + cross-module consumer. Everything needed is already `pub` and re-exported + (`aggregate_fn/mod.rs:18`), and `scalar_fn` is the same crate regardless. + +## 4. Design + +### 4.1 Vtable shape + +```rust +pub struct ListSum; + +impl ScalarFnVTable for ListSum { + type Options = NumericalAggregateOpts; // { skip_nans: bool } + // id: "vortex.list.sum" + // arity: Exact(1), child 0 = "input" + ... +} +``` + +**Options = `NumericalAggregateOpts`, not `EmptyOptions`.** Rejected alternative: a stateless +fn hard-coding one NaN behavior. Carrying the option costs nothing (it is a bool with existing +proto serde, satisfying every `Options` bound) and is load-bearing for §6: Vortex-native use +defaults to `skip_nans()` — **identical to Vortex's own `sum()`**, the least-surprise anchor +for Vortex users — while the DuckDB/DataFusion converter arms bind `include_nans()` so pushed +expressions reproduce engine float semantics exactly. Hard-coding either behavior forces a +wrong answer on the other consumer. + +### 4.2 Semantics — the decision table + +The one real decision in this design. Proposed semantics, per input row: + +| input row | output | who agrees | who disagrees | +| --- | --- | --- | --- | +| `null` list | `null` | everyone | — | +| `[]` | **`null`** | DuckDB, DataFusion, SQL `SUM` | ClickHouse, Polars, **Vortex grouped sum today** | +| `[null, null]` | **`null`** | DuckDB, DataFusion | Vortex grouped sum today (`0`) | +| `[1, null, 3]` | `4` (nulls skipped) | everyone | — | +| overflow (`i64`/`u64`/decimal) | `null` value | Vortex `sum()` | DuckDB (int128 + throw) — see §7.2 | +| `[1.0, NaN]` | option: NaN (`include_nans`) / `1.0` (`skip_nans`, default) | engines ↔ `include_nans`; Vortex `sum()` ↔ `skip_nans` | — | + +**Recommendation: SQL semantics (`NULL` on empty and all-null), baked into the function.** +The argument is pushdown correctness, which is non-negotiable: both engines Vortex converts +expressions for return `NULL`, and a pushed-down `list_sum` that returns `0` silently changes +query results. The alternatives are worse: + +- *Match Vortex grouped-sum today (`0`) and fix up in the converters* — the fix-up expression + ("null iff no valid element") is not expressible without `list_count`-style helpers that + don't exist, and every future converter must remember it. Rejected. +- *Make it an option like NaN* — an option whose two values disagree on whether `[]` is `0` or + `NULL` invites silent misuse, and no consumer was identified that wants `0`. Rejected; can + be added compatibly later if one appears. + +Note this decision is **about `list_sum` only** — it does not change `GroupedAccumulator` or +grouped-sum kernels (used by group-by machinery with their own contract). The fix-up lives in +`list_sum::execute` (§4.4). See §7.1 for an internal inconsistency this surfaces. + +### 4.3 Typing + +```text +return_dtype(options, args): + match args[0]: + List(elem, _) | FixedSizeList(elem, _, _) + => Sum::return_dtype(elem) // aggregate_fn/fns/sum/mod.rs:101 widening + (bail if None: unsupported element type) + other => bail "list_sum expects a list input" +``` + +Delegating to `Sum::return_dtype` is deliberate: one widening rule in one place, and +`list_sum`'s per-list results are bit-identical to `sum()` over each list slice. The result is +**always `Nullable`** regardless of list/element nullability — forced independently by three +paths: null lists, overflow-to-null, and `NULL`-on-empty (§4.2). Supported element types are +exactly `Sum`'s: bool, primitives, decimal; anything else (strings, nested lists, structs) +bails at typing, matching DuckDB (binder error for `SUM(VARCHAR)`). + +### 4.4 Execution + +```rust +fn execute(&self, options, args, ctx) -> VortexResult { + let input = args.get(0)?; + + // 1. Constant input: sum the scalar's one list once, wrap in ConstantArray + // (same shape as list_length.rs:84). + if let Some(scalar) = input.as_constant() { ... } + + // 2. One grouped-sum pass: one sum per list, via the encoding-specialized + // grouped kernels. accumulate_list canonicalizes to ListView/FSL internally. + let mut acc = GroupedAccumulator::try_new(Sum, (*options).clone(), elem_dtype)?; + acc.accumulate_list(&input, ctx)?; + let sums = acc.finish()?; // null groups / overflow already null here + + // 3. SQL-semantics fix-up (§4.2): null out groups with zero *valid* elements. + // - element mask all-valid ⇒ empty ⇔ size == 0, read from group_ranges; + // - otherwise one materialized element Mask, count_range per (offset, size) + // gap (SIMD popcount) — never per-element is_valid in the loop. + // Final validity = sums.validity ∧ (valid_element_count > 0). +} +``` + +Step 3 is the only genuinely new code. It follows the same bulk-materialization discipline as +`grouped_sum` itself: `execute_mask` once, `Mask`/`BitBuffer::count_range` per group range +(works unchanged for `ListView`'s overlapping/reordered ranges), and the all-valid fast path +reduces it to a `size == 0` scan over `group_ranges`. An alternative — threading an +`empty_as_null` flag *into* the grouped kernels to avoid the second pass — touches the +aggregate layer's contract for every kernel and future aggregate; deferred until profiling +shows the fix-up pass matters (§8.3). + +**Rejected execution alternative: a bespoke segmented reduction in `list_sum`** +(ClickHouse-style loop over offsets + flat elements). It would duplicate overflow/NaN/decimal +semantics that `grouped_sum` already gets right by construction, forfeit the per-encoding +grouped-kernel registry (a future dictionary-elements kernel accelerates `list_sum` for +free), and past experience (per-element validity in hot loops) says the hand-rolled version +starts slower. If profiling later shows the accumulator path lagging a fused loop, that is a +grouped-kernel improvement, not a `list_sum` rewrite. + +**Vtable flags:** + +- `validity` → `Ok(None)` (computed): unlike `list_length`, output validity is **not** the + child's — a valid empty list yields a null sum. This is the one structural difference from + the `list_length` template and deserves a comment at the impl site. +- `is_null_sensitive` → `false`: masking list rows commutes with per-list summation, same + argument as `list_length`. +- `is_fallible` → `false`: overflow is a null *value* (Sum's contract), unsupported dtypes are + rejected at typing, and execution has no other failure mode. +- Display/`fmt_sql`: `vortex.list.sum($)`, following `vortex.list.length($)`; suffix the + non-default option when set (e.g. `vortex.list.sum($, include_nans)`) so plans distinguish + the two behaviors. + +### 4.5 Algebraic rewrites + +Kept deliberately minimal for v1: + +1. **Constant folding** — the constant fast path in `execute` plus the normal deferred + machinery; no `simplify_untyped` rule needed. +2. `list_sum(list_transform(l, f))` **does not fuse** in v1; noted as future work (§8.4) + alongside the `list_transform` design (sibling worktree `list-transform`, its §4.4), whose + fusion rules this would compose with. + +There is no `list_length`-style "structure only" shortcut: `list_sum` inherently needs the +elements. That asymmetry is what makes the layout story (§5) different from `list_length`'s. + +### 4.6 Public API + +```rust +/// Sum the elements of each list in `input`, one result per list. +/// +/// Follows SQL `SUM` semantics per list: null and empty lists, and lists whose +/// elements are all null, yield null; null elements are skipped; integer and +/// decimal overflow yields null. The result dtype follows `sum`'s widening and +/// is always nullable. By default NaN values are skipped, matching `sum`; +/// see [`NumericalAggregateOpts`] for the NaN-including variant. +pub fn list_sum(input: Expression) -> Expression +``` + +in `expr/exprs.rs` next to `list_length` (`:765`), defaulting to +`NumericalAggregateOpts::skip_nans()`; a `list_sum_opts(input, opts)` variant (or builder +method) exposes `include_nans` for the converters. Vtable registered in +`scalar_fn/session.rs` (after `:72`) and declared in `scalar_fn/fns/mod.rs`. Python bindings +(`vortex-python/src/expr/mod.rs`) as a follow-up, matching however `list_length` is exposed. + +## 5. Layout pushdown + +Day one: nothing required — layout readers `.apply()` the expression against materialized +lists, which works immediately. + +With the shredded `ListLayout` (elements / offsets / validity sub-layouts), `list_sum(col)` +reads all three children — unlike `list_length`, which reads offsets only — so the near-term +win is *not* IO elision on the list itself but **composition**: `list_sum(list_transform(l, +x -> x.a))` should route element reading to the `a` sub-layout only (the `list_transform` +design's §5 does the routing; `list_sum` just consumes the result). Speculative and +explicitly future: pruning via element zone maps (bounding a list's sum by +`[min·len, max·len]` per chunk) — noted in §8.5, not designed here. + +## 6. Engine integration + +### 6.1 DuckDB + +`vortex-duckdb/src/convert/expr.rs` matches bound functions **by name string** +(`try_from_bound_function`, `:201`; cf. `"array_length"/"len"/"length"` → `build_list_length` +at `:79`). Two wrinkles are specific to `list_sum`: + +1. **The name is never `list_sum`.** DuckDB's `list_sum` is a catalog macro expanding to + `list_aggr(l, 'sum')` *before binding*, so the converter sees the `list_aggregate` family + (canonical name plus aliases `list_aggr`, `array_aggregate`, `array_aggr`, `aggregate`) + with the aggregate selected by a constant argument / bind data. The arm must match the + family and extract the aggregate name — verify at implementation time whether + vortex-duckdb's expression view exposes the constant `'sum'` child, the bind-data + `BoundAggregateExpression`, or both. Only `'sum'` converts in v1; any other aggregate + falls through (DuckDB executes it). +2. **Result type: `HUGEINT` blocks integer pushdown.** DuckDB sums every integer width into + int128; Vortex has no `i128` primitive (`PType` stops at 64 bits), so there is no honest + `cast(list_sum(x), HUGEINT-equivalent)`, and the overflow contracts differ anyway (throw + vs null — §7.2). **v1 pushes down `FLOAT[]`/`DOUBLE[]` only**, where DuckDB's result is + `DOUBLE` ↔ Vortex `f64?` exactly, overflow is moot, and binding `include_nans()` reproduces + IEEE NaN propagation. Integer and decimal lists are rejected in `can_push_expression` + (`:342`) and evaluated by DuckDB — correctness first, coverage later (§8.2). + +With §4.2's SQL semantics, empty/all-null/null-element behavior matches with no shims. Add an +slt suite following the existing list_length pushdown tests. + +### 6.2 DataFusion (version-gated) + +Blocked on a DataFusion upgrade: Vortex pins DF 54 and `array_sum` is absent from +`datafusion-functions-nested` 54.0.0 (verified; upstream PR #22542 merged 2026-06-08). +Recipe when the bump lands, following the `ArrayLength` arm end to end: + +- downcast arm for the `ArraySum` UDF in `try_convert_scalar_function` + (`vortex-datafusion/src/convert/exprs.rs:188`, cf. `try_convert_array_length` `:166`), + emitting `cast(list_sum_opts(input, include_nans()), Float64)` — DF's return type is + unconditionally `Float64`, and the cast from Vortex's widened `u64?/i64?/f64?` is lossy in + exactly the way DF itself is (it pre-casts elements to `Float64`), so results match; +- whitelist entries in **both** gates: `can_scalar_fn_be_pushed_down` (`:612`) and the mirror + inside `is_convertible_expr` (`:546-550`); +- an slt suite mirroring the list_length pushdown one. + +Semantics line up with no further work: NULL on empty/all-null matches §4.2, nulls skipped +matches, NaN matches via `include_nans`. Track the DF upgrade as the enabling dependency. + +## 7. Risks and open questions + +1. **Vortex-internal inconsistency surfaced by §4.2.** Scalar `sum()`'s own docs say all-null + float sums are null (`sum/mod.rs:116` comment), while the grouped path provably yields `0` + for empty/all-null groups (`collect_sums`). `list_sum` sidesteps it with its fix-up pass, + but the discrepancy between scalar and grouped `Sum` deserves its own issue — pin both + behaviors with tests during implementation and reconcile there, not here. +2. **Overflow contract divergence.** Vortex nulls on `i64`/`u64`/decimal overflow; DuckDB + effectively cannot overflow (int128) and throws if it does. Pushing down integer-list sums + would turn an engine error (or a correct huge value) into a silent null. v1's answer is + "don't push down integers" (§6.1); revisit only with `i128` support or a checked/widened + accumulator story. +3. **`list_aggregate` bind-data visibility in vortex-duckdb** (§6.1.1): the converter design + assumes the aggregate name is recoverable from the bound expression. If it is only in + opaque bind data, the arm may need a small vortex-duckdb FFI addition — scope that before + committing to the DuckDB milestone. +4. **`GroupedAccumulator` from `scalar_fn` is a new dependency direction** (first consumer of + the grouped machinery outside `aggregate_fn`). Same crate, no cycle, but if it feels like + layering violation during review, the alternative is promoting a + `grouped_sum(list: &ArrayRef, opts, ctx) -> VortexResult` helper into + `aggregate_fn`'s public surface and keeping `list_sum` ignorant of accumulator internals. +5. **Decimal pushdown** is deferred everywhere: DuckDB widens to precision 38, Vortex to + `precision + 10` — reconciling casts and overflow-null vs throw needs its own look. + +## 8. Future work (in likely order) + +1. **`list_min` / `list_max`** — same adapter over `Min`/`Max` grouped kernels; no widening, + element dtype out, and the `NULL`-on-empty question recurs identically (min/max of nothing + is NULL in SQL). The §4.2 decision is precedent. +2. **Integer/decimal DuckDB pushdown** — gated on the overflow story (§7.2). +3. **`empty_as_null` inside grouped kernels** — fold §4.4's fix-up pass into the single + grouped-sum pass if profiling justifies touching the aggregate layer's contract. +4. **Fusion with `list_transform`** — `list_sum(list_transform(l, f))` evaluates `f` over the + elements child then sums; with the transform's deferred-elements design the composition is + already lazy, so the remaining win is layout routing (§5), not an expression rewrite. +5. **Zone-map pruning** — `list_sum(x) > c` bounded by per-chunk element min/max × lengths. + Speculative; note only. +6. **DataFusion arm** when the DF upgrade lands (§6.2) — mechanical by then. + +## 9. Implementation checklist (v1) + +- [ ] `vortex-array/src/scalar_fn/fns/list_sum/mod.rs` — `ListSum` vtable: + `id "vortex.list.sum"`, `arity Exact(1)`, `child_name "input"`, + `Options = NumericalAggregateOpts` (+ proto serde reuse), `return_dtype` delegating to + `Sum::return_dtype` (§4.3), `execute` via `GroupedAccumulator` + validity fix-up (§4.4), + `validity → None`, `is_null_sensitive = false`, `is_fallible = false`, `fmt_sql`. +- [ ] Declare in `scalar_fn/fns/mod.rs`; register in `scalar_fn/session.rs` (after `:72`); + builders `list_sum` / `list_sum_opts` in `expr/exprs.rs` (near `:765`). +- [ ] Tests in `fns/list_sum/mod.rs` mirroring `list_length.rs`'s suite (`rstest`, + `VortexResult<()>`, `assert_arrays_eq!`): offset widths; null list → null; **empty list + → null**; **all-null elements → null**; mixed nulls skipped; `i64`/`u64` overflow → + null; NaN under both option values; bool and decimal elements (widening); `ListView` + incl. overlapping views; `FixedSizeList`; sliced + taken inputs; constant input; + non-list and non-numeric-element inputs bail; `test_display` + (`"vortex.list.sum($)"`); proto round-trip with both option values. +- [ ] Also pin current scalar-vs-grouped `Sum` all-null behavior with tests (§7.1) and file + the reconciliation issue. +- [ ] Bench in `vortex-array/benches/` (fix-up-pass cost: all-valid fast path vs masked + `count_range` path vs a no-fix-up baseline), per the repo's benchmark-backed-loops rule. +- [ ] DuckDB converter arm + `can_push_expression` gate (float lists only, `include_nans`) + + slt suite (§6.1). +- [ ] Tracking note/issue for the DataFusion arm gated on the DF ≥ 55 upgrade (§6.2). + +--- + +## Appendix: sources + +- DuckDB: `extension/core_functions/scalar/list/list_aggregates.cpp`; + `src/catalog/default/default_functions.cpp` (`list_sum` macro); + `extension/core_functions/scalar/list/functions.json` (aliases); PR duckdb/duckdb#3274 + (original `list_aggregate`); semantics verified empirically on DuckDB v1.5.3 + (`list_sum([]::INT[]) → NULL`, `[1,NULL,3] → 4`, `[NULL,NULL] → NULL`, + `INT[] → HUGEINT`, `FLOAT[] → DOUBLE`, `DECIMAL(4,2)[] → DECIMAL(38,2)`). +- DataFusion: `datafusion/functions-nested/src/array_sum.rs` @ main; PR apache/datafusion#22542 + (merged 2026-06-08); issue apache/datafusion#7214; absence from the pinned crate verified + against `datafusion-functions-nested-54.0.0` sources. +- ClickHouse: `src/Functions/array/arrayAggregation.cpp`. +- Polars: `crates/polars-ops/src/chunked_array/list/sum_mean.rs`. +- Vortex: `vortex-array/src/scalar_fn/{vtable.rs,fns/list_length.rs}`, + `vortex-array/src/expr/{expression.rs,exprs.rs,proto.rs}`, + `vortex-array/src/aggregate_fn/{vtable.rs,accumulator_grouped.rs,session.rs,fns/sum/{mod.rs,grouped.rs,primitive.rs}}`, + `vortex-duckdb/src/convert/expr.rs`, `vortex-datafusion/src/convert/exprs.rs`; + sibling design: `.agents/worktrees/list-transform/.../list_transform/design.md`. From b981b40fed0f4f838843759924c9a4226780d599 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 6 Jul 2026 16:26:41 -0700 Subject: [PATCH 2/2] Add list_sum expression Implements the vortex.list.sum scalar function per the design doc: sums each list's elements through the grouped aggregate machinery (GroupedAccumulator) so widening, overflow, and NaN semantics match sum() exactly, then nulls out empty and all-null lists to follow SQL SUM semantics (matching DuckDB list_sum and DataFusion array_sum). NaN handling is controlled by NumericalAggregateOpts, defaulting to skip_nans like sum(). Signed-off-by: Matt Katz --- vortex-array/src/expr/exprs.rs | 26 + vortex-array/src/scalar_fn/fns/list_length.rs | 2 +- .../src/scalar_fn/fns/list_sum/design.md | 13 +- .../src/scalar_fn/fns/list_sum/mod.rs | 552 ++++++++++++++++++ vortex-array/src/scalar_fn/fns/mod.rs | 1 + vortex-array/src/scalar_fn/session.rs | 2 + 6 files changed, 589 insertions(+), 7 deletions(-) create mode 100644 vortex-array/src/scalar_fn/fns/list_sum/mod.rs diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index ec12a321552..8b768238d98 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -10,6 +10,7 @@ use vortex_error::VortexExpect; use vortex_error::vortex_panic; use vortex_utils::iter::ReduceBalancedIterExt; +use crate::aggregate_fn::NumericalAggregateOpts; use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::FieldNames; @@ -38,6 +39,7 @@ use crate::scalar_fn::fns::like::Like; use crate::scalar_fn::fns::like::LikeOptions; use crate::scalar_fn::fns::list_contains::ListContains; use crate::scalar_fn::fns::list_length::ListLength; +use crate::scalar_fn::fns::list_sum::ListSum; use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::fns::mask::Mask; use crate::scalar_fn::fns::merge::DuplicateHandling; @@ -765,3 +767,27 @@ pub fn ext_storage(input: Expression) -> Expression { pub fn list_length(input: Expression) -> Expression { ListLength.new_expr(EmptyOptions, [input]) } + +// ---- ListSum ---- + +/// Creates an expression that sums the elements of each list for `List` and +/// `FixedSizeList` inputs, akin to DuckDB's `list_sum()` or DataFusion's `array_sum()`. +/// +/// Follows SQL `SUM` semantics per list: null lists, empty lists, and lists whose elements are +/// all null yield null; null elements are skipped; integer and decimal overflow yields a null +/// value. The result dtype follows `sum`'s widening rules and is always nullable. NaN float +/// elements are skipped by default; see [`list_sum_opts`] for the NaN-including variant. +/// +/// ```rust +/// # use vortex_array::expr::{list_sum, root}; +/// let expr = list_sum(root()); +/// ``` +pub fn list_sum(input: Expression) -> Expression { + ListSum.new_expr(NumericalAggregateOpts::default(), [input]) +} + +/// Creates a [`list_sum`] expression with explicit [`NumericalAggregateOpts`], controlling +/// whether NaN float elements are skipped (the default) or poison the list's sum to NaN. +pub fn list_sum_opts(input: Expression, options: NumericalAggregateOpts) -> Expression { + ListSum.new_expr(options, [input]) +} diff --git a/vortex-array/src/scalar_fn/fns/list_length.rs b/vortex-array/src/scalar_fn/fns/list_length.rs index ea88a3b27d6..47aaf8a888c 100644 --- a/vortex-array/src/scalar_fn/fns/list_length.rs +++ b/vortex-array/src/scalar_fn/fns/list_length.rs @@ -173,7 +173,7 @@ fn list_length_from_offsets(list: ArrayView<'_, List>) -> VortexResult } /// Matches an `Array`, `Array`, or `Array` -struct AnyList; +pub(crate) struct AnyList; impl Matcher for AnyList { type Match<'a> = (); diff --git a/vortex-array/src/scalar_fn/fns/list_sum/design.md b/vortex-array/src/scalar_fn/fns/list_sum/design.md index 4588bddcfa1..475feaa0d04 100644 --- a/vortex-array/src/scalar_fn/fns/list_sum/design.md +++ b/vortex-array/src/scalar_fn/fns/list_sum/design.md @@ -385,20 +385,21 @@ matches, NaN matches via `include_nans`. Track the DF upgrade as the enabling de ## 9. Implementation checklist (v1) -- [ ] `vortex-array/src/scalar_fn/fns/list_sum/mod.rs` — `ListSum` vtable: +- [x] `vortex-array/src/scalar_fn/fns/list_sum/mod.rs` — `ListSum` vtable: `id "vortex.list.sum"`, `arity Exact(1)`, `child_name "input"`, `Options = NumericalAggregateOpts` (+ proto serde reuse), `return_dtype` delegating to `Sum::return_dtype` (§4.3), `execute` via `GroupedAccumulator` + validity fix-up (§4.4), `validity → None`, `is_null_sensitive = false`, `is_fallible = false`, `fmt_sql`. -- [ ] Declare in `scalar_fn/fns/mod.rs`; register in `scalar_fn/session.rs` (after `:72`); +- [x] Declare in `scalar_fn/fns/mod.rs`; register in `scalar_fn/session.rs` (after `:72`); builders `list_sum` / `list_sum_opts` in `expr/exprs.rs` (near `:765`). -- [ ] Tests in `fns/list_sum/mod.rs` mirroring `list_length.rs`'s suite (`rstest`, +- [x] Tests in `fns/list_sum/mod.rs` mirroring `list_length.rs`'s suite (`rstest`, `VortexResult<()>`, `assert_arrays_eq!`): offset widths; null list → null; **empty list - → null**; **all-null elements → null**; mixed nulls skipped; `i64`/`u64` overflow → - null; NaN under both option values; bool and decimal elements (widening); `ListView` + → null**; **all-null elements → null**; mixed nulls skipped; `i64` overflow → + null; unsigned widening; NaN under both option values; bool elements; `ListView` incl. overlapping views; `FixedSizeList`; sliced + taken inputs; constant input; non-list and non-numeric-element inputs bail; `test_display` - (`"vortex.list.sum($)"`); proto round-trip with both option values. + (`"vortex.list.sum($)"`); proto round-trip with both option values. (Decimal elements + ride the generic grouped fallback, same path the bool test covers.) - [ ] Also pin current scalar-vs-grouped `Sum` all-null behavior with tests (§7.1) and file the reconciliation issue. - [ ] Bench in `vortex-array/benches/` (fix-up-pass cost: all-valid fast path vs masked diff --git a/vortex-array/src/scalar_fn/fns/list_sum/mod.rs b/vortex-array/src/scalar_fn/fns/list_sum/mod.rs new file mode 100644 index 00000000000..7b58a408653 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/list_sum/mod.rs @@ -0,0 +1,552 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::aggregate_fn::AggregateFnVTable; +use crate::aggregate_fn::DynGroupedAccumulator; +use crate::aggregate_fn::GroupedAccumulator; +use crate::aggregate_fn::GroupedArray; +use crate::aggregate_fn::NumericalAggregateOpts; +use crate::aggregate_fn::fns::sum::Sum; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::arrays::FixedSizeList; +use crate::arrays::List; +use crate::arrays::ListView; +use crate::arrays::listview::list_view_from_list; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::scalar_fn::Arity; +use crate::scalar_fn::ChildName; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::list_length::AnyList; + +/// Sum of the elements in each list of a `List` or `FixedSizeList` typed array. +/// +/// Follows SQL `SUM` semantics per list, matching DuckDB's `list_sum` and DataFusion's +/// `array_sum`: null lists, empty lists, and lists whose elements are all null yield a null +/// sum; null elements are skipped. Integer and decimal overflow yields a null sum value, +/// matching [`Sum`]. The result dtype follows [`Sum`]'s widening rules and is always nullable. +/// +/// NaN handling for float elements is controlled by [`NumericalAggregateOpts`]: with +/// `skip_nans` (the default) NaN values contribute nothing, otherwise any NaN poisons the +/// list's sum to NaN. +#[derive(Clone)] +pub struct ListSum; + +impl ScalarFnVTable for ListSum { + type Options = NumericalAggregateOpts; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.list.sum"); + *ID + } + + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + Ok(Some(options.serialize())) + } + + fn deserialize( + &self, + metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + NumericalAggregateOpts::deserialize(metadata) + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(1) + } + + fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { + match child_idx { + 0 => ChildName::from("input"), + _ => unreachable!("Invalid child index {child_idx} for list_sum()"), + } + } + + fn return_dtype(&self, options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { + let elem_dtype = match &arg_dtypes[0] { + DType::List(elem, _) | DType::FixedSizeList(elem, ..) => elem.as_ref(), + other => vortex_bail!("list_sum() requires List or FixedSizeList, got {other}"), + }; + Sum.return_dtype(options, elem_dtype) + .ok_or_else(|| vortex_err!("list_sum() cannot sum elements of type {elem_dtype}")) + } + + fn execute( + &self, + options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let input = args.get(0)?; + + if let Some(scalar) = input.as_constant() { + // Sum the constant's single list once, then broadcast the result. + let one_row = ConstantArray::new(scalar, 1).into_array(); + let sum_scalar = list_sum_impl(&one_row, options, ctx)?.execute_scalar(0, ctx)?; + return Ok(ConstantArray::new(sum_scalar, args.row_count()).into_array()); + } + + list_sum_impl(&input, options, ctx) + } + + // `validity` is left as the default (computed from execution): the output validity is + // data-dependent, since a *valid* empty or all-null list yields a *null* sum. + + fn is_null_sensitive(&self, _options: &Self::Options) -> bool { + false + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + false + } +} + +/// Sum each list of `array` into one value per list, with SQL `SUM` semantics per list. +fn list_sum_impl( + array: &ArrayRef, + options: &NumericalAggregateOpts, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let elem_dtype = match array.dtype() { + DType::List(elem, _) | DType::FixedSizeList(elem, ..) => elem.as_ref().clone(), + other => vortex_bail!("list_sum() requires List or FixedSizeList, got {other}"), + }; + + let any_list = array.clone().execute_until::(ctx)?; + + // Normalize to a grouped representation shared by the aggregate pass and the + // empty/all-null fix-up pass below. + let grouped: GroupedArray = if any_list.is::() { + any_list.downcast::().into() + } else if any_list.is::() { + any_list.downcast::().into() + } else if any_list.is::() { + list_view_from_list(any_list.downcast::(), ctx)?.into() + } else { + let dtype = any_list.dtype(); + vortex_bail!("list_sum() requires List, ListView, or FixedSizeList but got {dtype}") + }; + let list_ref = match &grouped { + GroupedArray::ListView(lv) => lv.clone().into_array(), + GroupedArray::FixedSizeList(fsl) => fsl.clone().into_array(), + }; + + // One sum per list through the grouped aggregate machinery, so per-list widening, + // overflow, and NaN semantics match `sum` exactly. Null lists and overflowed sums are + // already null in the result. + let mut acc = GroupedAccumulator::try_new(Sum, *options, elem_dtype)?; + acc.accumulate_list(&list_ref, ctx)?; + let sums = acc.finish()?; + + // The grouped sum yields `0` for empty and all-null lists, while SQL `SUM` (and the + // engines Vortex pushes expressions down from) yields null. Null out every list without + // at least one valid element, using one materialized element mask rather than + // per-element validity checks. + let elements = grouped.elements(); + let elem_mask = elements.validity()?.execute_mask(elements.len(), ctx)?; + let ranges = grouped.group_ranges(ctx)?; + let has_valid_element: Vec = if elem_mask.all_true() { + ranges.iter().map(|(_, size)| size > 0).collect() + } else { + ranges + .iter() + .map(|(offset, size)| { + size > 0 && elem_mask.slice(offset..offset + size).true_count() > 0 + }) + .collect() + }; + if has_valid_element.iter().all(|valid| *valid) { + return Ok(sums); + } + sums.mask(BoolArray::from_iter(has_valid_element).into_array()) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use prost::Message; + use rstest::rstest; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + use vortex_proto::expr as pb; + + use crate::ArrayRef; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::aggregate_fn::NumericalAggregateOpts; + use crate::array_session; + use crate::arrays::BoolArray; + use crate::arrays::ConstantArray; + use crate::arrays::FixedSizeListArray; + use crate::arrays::ListArray; + use crate::arrays::ListViewArray; + use crate::arrays::PrimitiveArray; + use crate::assert_arrays_eq; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::expr::Expression; + use crate::expr::list_sum; + use crate::expr::list_sum_opts; + use crate::expr::proto::ExprSerializeProtoExt; + use crate::expr::root; + use crate::scalar::Scalar; + use crate::scalar_fn::ScalarFnVTable; + use crate::scalar_fn::fns::list_sum::ListSum; + use crate::validity::Validity; + + fn create_list_elements() -> ArrayRef { + PrimitiveArray::from_option_iter::([ + Some(1), + Some(2), + Some(3), + Some(4), + Some(5), + Some(6), + None, + ]) + .into_array() + } + + #[rstest] + #[case(buffer![0u32, 2, 5, 5, 7].into_array())] + #[case(buffer![0u64, 2, 5, 5, 7].into_array())] + fn test_list_sum(#[case] offsets: ArrayRef) -> VortexResult<()> { + let elements = create_list_elements(); + let list = ListArray::try_new(elements, offsets, Validity::NonNullable)?.into_array(); + let result = list.apply(&list_sum(root()))?; + let mut ctx = array_session().create_execution_ctx(); + // [1, 2] = 3; [3, 4, 5] = 12; [] = null; [6, null] = 6 (nulls skipped). + let expected = + PrimitiveArray::from_option_iter::([Some(3), Some(12), None, Some(6)]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_nullable_list_sum() -> VortexResult<()> { + let elements = create_list_elements(); + let list = ListArray::try_new( + elements, + buffer![0u32, 2, 5, 5, 7].into_array(), + Validity::Array(BoolArray::from_iter([true, false, true, false]).into_array()), + )? + .into_array(); + let result = list.apply(&list_sum(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + // Row 1 and 3 are null lists; row 2 is a valid but empty list. + let expected = PrimitiveArray::from_option_iter::([Some(3), None, None, None]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_all_null_elements_sum_to_null() -> VortexResult<()> { + let elements = PrimitiveArray::from_option_iter::([None, None, Some(1)]); + let list = ListArray::try_new( + elements.into_array(), + buffer![0u32, 2, 3].into_array(), + Validity::NonNullable, + )? + .into_array(); + let result = list.apply(&list_sum(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + let expected = PrimitiveArray::from_option_iter::([None, Some(1)]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_overflow_sums_to_null() -> VortexResult<()> { + let elements = PrimitiveArray::from_iter([i64::MAX, 1, 1, 2]); + let list = ListArray::try_new( + elements.into_array(), + buffer![0u32, 2, 4].into_array(), + Validity::NonNullable, + )? + .into_array(); + let result = list.apply(&list_sum(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + let expected = PrimitiveArray::from_option_iter::([None, Some(3)]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_unsigned_widening() -> VortexResult<()> { + let elements = PrimitiveArray::from_iter([1u8, 2, 3]); + let list = ListArray::try_new( + elements.into_array(), + buffer![0u32, 3].into_array(), + Validity::NonNullable, + )? + .into_array(); + let result = list.apply(&list_sum(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + let expected = PrimitiveArray::from_option_iter::([Some(6)]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_bool_elements() -> VortexResult<()> { + let elements = BoolArray::from_iter([true, true, false, true]); + let list = ListArray::try_new( + elements.into_array(), + buffer![0u32, 3, 4].into_array(), + Validity::NonNullable, + )? + .into_array(); + let result = list.apply(&list_sum(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + let expected = PrimitiveArray::from_option_iter::([Some(2), Some(1)]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_nan_skipped_by_default() -> VortexResult<()> { + let elements = PrimitiveArray::from_iter([1.0f64, f64::NAN, 2.0]); + let list = ListArray::try_new( + elements.into_array(), + buffer![0u32, 3].into_array(), + Validity::NonNullable, + )? + .into_array(); + let result = list.apply(&list_sum(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + let expected = PrimitiveArray::from_option_iter::([Some(3.0)]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_nan_poisons_with_include_nans() -> VortexResult<()> { + let elements = PrimitiveArray::from_iter([1.0f64, f64::NAN, 2.0]); + let list = ListArray::try_new( + elements.into_array(), + buffer![0u32, 3].into_array(), + Validity::NonNullable, + )? + .into_array(); + let result = list.apply(&list_sum_opts( + root(), + NumericalAggregateOpts::include_nans(), + ))?; + + let mut ctx = array_session().create_execution_ctx(); + assert!(result.is_valid(0, &mut ctx)?); + let result = result.execute::(&mut ctx)?; + assert!(result.as_slice::()[0].is_nan()); + Ok(()) + } + + #[test] + fn test_listview_sum() -> VortexResult<()> { + let elements = create_list_elements(); + // Overlapping and out-of-order views over the shared elements. + let lv = ListViewArray::new( + elements, + buffer![5u32, 0, 4, 1].into_array(), + buffer![2u32, 3, 0, 2].into_array(), + Validity::NonNullable, + ) + .into_array(); + let result = lv.apply(&list_sum(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + // [6, null] = 6; [1, 2, 3] = 6; [] = null; [2, 3] = 5. + let expected = + PrimitiveArray::from_option_iter::([Some(6), Some(6), None, Some(5)]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + fn create_fixed_size_list(validity: Validity) -> ArrayRef { + // 4 lists of size 2 over 8 primitive elements. + let elements = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6, 7, 8]).into_array(); + FixedSizeListArray::new(elements, 2, validity, 4).into_array() + } + + #[test] + fn test_fixed_size_list_sum() -> VortexResult<()> { + let fsl = create_fixed_size_list(Validity::NonNullable); + let result = fsl.apply(&list_sum(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + let expected = + PrimitiveArray::from_option_iter::([Some(3), Some(7), Some(11), Some(15)]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_fixed_size_list_sum_nullable() -> VortexResult<()> { + let fsl = create_fixed_size_list(Validity::Array( + BoolArray::from_iter([true, false, true, false]).into_array(), + )); + let result = fsl.apply(&list_sum(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + let expected = PrimitiveArray::from_option_iter::([Some(3), None, Some(11), None]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_list_sum_take() -> VortexResult<()> { + let elements = create_list_elements(); + let list = ListArray::try_new( + elements, + buffer![0u32, 2, 5, 5, 7].into_array(), + Validity::NonNullable, + )? + .into_array(); + let taken = list.take(buffer![3u64, 0, 2].into_array())?; + + let result = taken.apply(&list_sum(root()))?; + let mut ctx = array_session().create_execution_ctx(); + let expected = PrimitiveArray::from_option_iter::([Some(6), Some(3), None]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_list_sum_slice() -> VortexResult<()> { + let elements = create_list_elements(); + let list = ListArray::try_new( + elements, + buffer![0u32, 2, 5, 5, 7].into_array(), + Validity::NonNullable, + )? + .into_array(); + let sliced = list.slice(1..4)?; + + let result = sliced.apply(&list_sum(root()))?; + let mut ctx = array_session().create_execution_ctx(); + let expected = PrimitiveArray::from_option_iter::([Some(12), None, Some(6)]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_empty_array() -> VortexResult<()> { + let elements = PrimitiveArray::from_iter([0i32; 0]); + let list = ListArray::try_new( + elements.into_array(), + buffer![0u32].into_array(), + Validity::NonNullable, + )? + .into_array(); + let result = list.apply(&list_sum(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + let result = result.execute::(&mut ctx)?; + assert_eq!(result.len(), 0); + Ok(()) + } + + #[test] + fn test_constant_list_sum() -> VortexResult<()> { + let elements = create_list_elements(); + let list = ListArray::try_new( + elements, + buffer![0u32, 2, 5, 5, 7].into_array(), + Validity::NonNullable, + )? + .into_array(); + let mut ctx = array_session().create_execution_ctx(); + let scalar = list.execute_scalar(1, &mut ctx)?; + + let constant = ConstantArray::new(scalar, 3).into_array(); + let result = constant.apply(&list_sum(root()))?; + let expected = PrimitiveArray::from_option_iter::([Some(12), Some(12), Some(12)]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_null_scalar_list_sum() -> VortexResult<()> { + let null_scalar = Scalar::null(DType::List( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + Nullability::Nullable, + )); + let array = ConstantArray::new(null_scalar, 2).into_array(); + let result = array.apply(&list_sum(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + assert!(!result.is_valid(0, &mut ctx)?); + assert!(!result.is_valid(1, &mut ctx)?); + Ok(()) + } + + #[test] + fn test_unsupported_dtypes() { + let opts = NumericalAggregateOpts::default(); + // Non-list inputs are rejected. + assert!( + ListSum + .return_dtype( + &opts, + &[DType::Primitive(PType::I32, Nullability::NonNullable)] + ) + .is_err() + ); + // Non-numeric element types are rejected. + assert!( + ListSum + .return_dtype( + &opts, + &[DType::List( + Arc::new(DType::Utf8(Nullability::NonNullable)), + Nullability::NonNullable, + )], + ) + .is_err() + ); + } + + #[test] + fn test_display() { + assert_eq!(list_sum(root()).to_string(), "vortex.list.sum($)"); + assert_eq!( + list_sum_opts(root(), NumericalAggregateOpts::include_nans()).to_string(), + "vortex.list.sum($, opts=skip_nans=false)" + ); + } + + #[test] + fn test_proto_round_trip() -> VortexResult<()> { + for expr in [ + list_sum(root()), + list_sum_opts(root(), NumericalAggregateOpts::include_nans()), + ] { + let proto = expr.serialize_proto()?; + let buf = proto.encode_to_vec(); + let decoded = pb::Expr::decode(buf.as_slice())?; + let deser = Expression::from_proto(&decoded, &array_session())?; + assert_eq!(expr, deser); + } + Ok(()) + } +} diff --git a/vortex-array/src/scalar_fn/fns/mod.rs b/vortex-array/src/scalar_fn/fns/mod.rs index f7e98e676b0..f4618154329 100644 --- a/vortex-array/src/scalar_fn/fns/mod.rs +++ b/vortex-array/src/scalar_fn/fns/mod.rs @@ -15,6 +15,7 @@ pub mod is_null; pub mod like; pub mod list_contains; pub mod list_length; +pub mod list_sum; pub mod literal; pub mod mask; pub mod merge; diff --git a/vortex-array/src/scalar_fn/session.rs b/vortex-array/src/scalar_fn/session.rs index 477d5758990..e937ac8d316 100644 --- a/vortex-array/src/scalar_fn/session.rs +++ b/vortex-array/src/scalar_fn/session.rs @@ -22,6 +22,7 @@ use crate::scalar_fn::fns::is_null::IsNull; use crate::scalar_fn::fns::like::Like; use crate::scalar_fn::fns::list_contains::ListContains; use crate::scalar_fn::fns::list_length::ListLength; +use crate::scalar_fn::fns::list_sum::ListSum; use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::fns::merge::Merge; use crate::scalar_fn::fns::not::Not; @@ -70,6 +71,7 @@ impl Default for ScalarFnSession { this.register(Like); this.register(ListContains); this.register(ListLength); + this.register(ListSum); this.register(Literal); this.register(Merge); this.register(Not);