Skip to content

Commit f62c12f

Browse files
authored
Add list_transform (#10503)
## Summary `list_transform(list, lambda)` applies a one-parameter lambda to every element of a `List` (in any encoding), keeping each list's length and validity. ```rust list_transform(col("l"), Lambda::new(["x"], checked_add(var("x"), lit(1_i32)))?) // vortex.list.transform($.l, (x) -> (@x + 1i32)) ``` ## Changes - `ListTransform`'s options are `ListTransformOptions`: the lambda as written (`Lambda`) in an expression, or bound to the element's dtype (`BoundLambda`) in a bound expression or an array. Its only child is the list. - `bind_options` binds the lambda to the element's dtype when the node is bound, so `execute` only runs it. A `FixedSizeList`, or a lambda with anything but exactly one parameter, is rejected there and in `return_dtype`. A lambda already bound to another element dtype is an error. - The result keeps the input's shape: a `ListArray` stays a `ListArray` and anything else runs as a list view, with the same offsets (and sizes) and validity around the new elements. - A lambda that can fail runs only on the elements valid lists reference. One that can't runs on the whole elements buffer, cut to the lists' range for a `ListArray`. - The function is strict (a null list gives a null result), it can fail only if the lambda can, and it is not serializable yet. ## Benchmark 100k lists, ~800k `i32` elements, ~1/8 of the lists null where nullable, result exported to Arrow, median µs, local scratch divan benchmark against DataFusion 55.2's `array_transform` on the same lists: | Lists | Lambda | DataFusion | This PR | |---|---|---|---| | Exact | `x + 1` | 49 | 71 | | Null lists already empty | `x + 1` | 60 | 73 | | Null lists holding elements | `x + 1` | 224 | 274 | | Null lists holding elements | `x = 3` | 249 | 81 | | Sliced | `x = 3` | 85 | 99 | | List view, back to back | `x + 1` | 903 | 75 | Most of the `x + 1` gap is Vortex's checked add against DataFusion's wrapping add; with `x = 3` the two kernels are even. DataFusion casts a list view to a list before transforming it. --------- Signed-off-by: Matt Katz <mhkatz97@gmail.com>
1 parent db0c700 commit f62c12f

6 files changed

Lines changed: 910 additions & 2 deletions

File tree

‎vortex-array/src/expr/exprs.rs‎

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ use crate::dtype::FieldNames;
1818
use crate::dtype::Nullability;
1919
use crate::expr::BoundExpression;
2020
use crate::expr::Expression;
21+
use crate::expr::Lambda;
2122
use crate::scalar::Scalar;
2223
use crate::scalar::ScalarValue;
2324
use crate::scalar_fn::EmptyOptions;
@@ -45,6 +46,8 @@ use crate::scalar_fn::fns::list_contains::ListContains;
4546
use crate::scalar_fn::fns::list_contains::ListContainsOptions;
4647
use crate::scalar_fn::fns::list_length::ListLength;
4748
use crate::scalar_fn::fns::list_sum::ListSum;
49+
use crate::scalar_fn::fns::list_transform::ListTransform;
50+
use crate::scalar_fn::fns::list_transform::ListTransformOptions;
4851
use crate::scalar_fn::fns::literal::Literal;
4952
use crate::scalar_fn::fns::mask::Mask;
5053
use crate::scalar_fn::fns::merge::DuplicateHandling;
@@ -153,8 +156,7 @@ pub fn bound_get_item(field: impl Into<FieldName>, child: BoundExpression) -> Bo
153156

154157
/// Creates an expression that reads the parameter `name` of the enclosing lambda.
155158
///
156-
/// It is only meaningful in the body of a [`Lambda`](crate::expr::Lambda). Anywhere else it fails
157-
/// to bind.
159+
/// It is only meaningful in the body of a [`Lambda`]. Anywhere else it fails to bind.
158160
///
159161
/// ```rust
160162
/// # use vortex_array::dtype::{DType, Nullability, PType};
@@ -1391,6 +1393,22 @@ pub fn bound_list_sum_opts(
13911393
.vortex_expect("list-sum expressions require a numeric list child")
13921394
}
13931395

1396+
// ---- ListTransform ----
1397+
1398+
/// Creates an expression that applies a lambda to every element of each list.
1399+
///
1400+
/// The lambda takes one parameter, the element.
1401+
///
1402+
/// ```rust
1403+
/// # use vortex_array::expr::{Lambda, checked_add, col, list_transform, lit, var};
1404+
/// let lambda = Lambda::new(["x"], checked_add(var("x"), lit(1_i32))).unwrap();
1405+
/// let expr = list_transform(col("l"), lambda);
1406+
/// assert_eq!(expr.to_string(), "vortex.list.transform($.l, (x) -> (@x + 1i32))");
1407+
/// ```
1408+
pub fn list_transform(list: Expression, lambda: Lambda) -> Expression {
1409+
ListTransform.new_expr(ListTransformOptions::Lambda(lambda), [list])
1410+
}
1411+
13941412
/// Reinterpret timestamp wall times in `options.time_zone`, preserving the time unit.
13951413
///
13961414
/// `ambiguous` is a UTF-8 expression containing `raise`, `earliest`, `latest`, or `null`.

‎vortex-array/src/expr/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@ pub use exprs::list_contains_opts;
8888
pub use exprs::list_length;
8989
pub use exprs::list_sum;
9090
pub use exprs::list_sum_opts;
91+
pub use exprs::list_transform;
9192
pub use exprs::lit;
9293
pub use exprs::lt;
9394
pub use exprs::lt_eq;
Lines changed: 236 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,236 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
#[cfg(test)]
5+
mod tests;
6+
7+
use std::fmt;
8+
use std::fmt::Display;
9+
use std::fmt::Formatter;
10+
use std::sync::Arc;
11+
12+
use itertools::Itertools;
13+
use vortex_error::VortexResult;
14+
use vortex_error::vortex_bail;
15+
use vortex_error::vortex_ensure;
16+
use vortex_error::vortex_ensure_eq;
17+
use vortex_session::registry::CachedId;
18+
19+
use crate::ArrayRef;
20+
use crate::ExecutionCtx;
21+
use crate::IntoArray;
22+
use crate::arrays::List;
23+
use crate::arrays::ListArray;
24+
use crate::arrays::ListViewArray;
25+
use crate::arrays::list::ListArrayExt;
26+
use crate::arrays::list::ListArraySlotsExt;
27+
use crate::arrays::listview::ListViewArrayExt;
28+
use crate::arrays::listview::ListViewArraySlotsExt;
29+
use crate::arrays::listview::ListViewRebuildMode;
30+
use crate::dtype::DType;
31+
use crate::dtype::FieldName;
32+
use crate::expr::BoundLambda;
33+
use crate::expr::Lambda;
34+
use crate::expr::display::ExprDisplay;
35+
use crate::scalar_fn::Arity;
36+
use crate::scalar_fn::ChildName;
37+
use crate::scalar_fn::ExecutionArgs;
38+
use crate::scalar_fn::ScalarFnId;
39+
use crate::scalar_fn::ScalarFnVTable;
40+
41+
/// Apply a lambda to every element of each list, keeping each list's length and validity.
42+
///
43+
/// The lambda takes one parameter, the element. The only child is the list, of `List` dtype in any
44+
/// encoding, and the result is a `List` of the same shape. A lambda that can fail runs only on the
45+
/// elements of valid lists.
46+
#[derive(Clone)]
47+
pub struct ListTransform;
48+
49+
/// The lambda of a [`ListTransform`], as written or bound to the element's dtype.
50+
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
51+
pub enum ListTransformOptions {
52+
/// The lambda as written, in an [`Expression`](crate::expr::Expression).
53+
Lambda(Lambda),
54+
/// The lambda bound to the element's dtype, in a
55+
/// [`BoundExpression`](crate::expr::BoundExpression) or an array.
56+
Bound(BoundLambda),
57+
}
58+
59+
impl Display for ListTransformOptions {
60+
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
61+
match self {
62+
Self::Lambda(lambda) => lambda.fmt(f),
63+
Self::Bound(lambda) => lambda.fmt(f),
64+
}
65+
}
66+
}
67+
68+
impl ScalarFnVTable for ListTransform {
69+
type Options = ListTransformOptions;
70+
71+
fn id(&self) -> ScalarFnId {
72+
static ID: CachedId = CachedId::new("vortex.list.transform");
73+
*ID
74+
}
75+
76+
fn arity(&self, _options: &ListTransformOptions) -> Arity {
77+
Arity::Exact(1)
78+
}
79+
80+
fn child_name(&self, _options: &ListTransformOptions, child_idx: usize) -> ChildName {
81+
match child_idx {
82+
0 => ChildName::from("list"),
83+
_ => unreachable!("Invalid child index {child_idx} for list_transform()"),
84+
}
85+
}
86+
87+
fn fmt_sql(
88+
&self,
89+
options: &ListTransformOptions,
90+
expr: &dyn ExprDisplay,
91+
f: &mut Formatter<'_>,
92+
) -> fmt::Result {
93+
write!(f, "{}({}, {options})", self.id(), expr.display_child(0))
94+
}
95+
96+
fn return_dtype(&self, options: &ListTransformOptions, args: &[DType]) -> VortexResult<DType> {
97+
let element = match options {
98+
ListTransformOptions::Lambda(lambda) => {
99+
lambda.return_dtype(&element_dtype(&args[0], lambda.params())?)?
100+
}
101+
ListTransformOptions::Bound(lambda) => {
102+
check_bound_to(lambda, &args[0])?;
103+
lambda.dtype().clone()
104+
}
105+
};
106+
Ok(DType::List(Arc::new(element), args[0].nullability()))
107+
}
108+
109+
fn bind_options(
110+
&self,
111+
options: &ListTransformOptions,
112+
args: &[DType],
113+
) -> VortexResult<Option<ListTransformOptions>> {
114+
match options {
115+
ListTransformOptions::Lambda(lambda) => {
116+
let element = element_dtype(&args[0], lambda.params())?;
117+
Ok(Some(ListTransformOptions::Bound(lambda.bind(&element)?)))
118+
}
119+
ListTransformOptions::Bound(lambda) => {
120+
check_bound_to(lambda, &args[0])?;
121+
Ok(None)
122+
}
123+
}
124+
}
125+
126+
fn execute(
127+
&self,
128+
options: &ListTransformOptions,
129+
args: &dyn ExecutionArgs,
130+
ctx: &mut ExecutionCtx,
131+
) -> VortexResult<ArrayRef> {
132+
let ListTransformOptions::Bound(lambda) = options else {
133+
vortex_bail!("list_transform()'s lambda is not bound");
134+
};
135+
match args.get(0)?.try_downcast::<List>() {
136+
Ok(list) => transform_list(list, lambda, ctx),
137+
Err(list) => transform_list_view(list.execute::<ListViewArray>(ctx)?, lambda, ctx),
138+
}
139+
}
140+
141+
fn is_strict(&self, _options: &ListTransformOptions) -> bool {
142+
true
143+
}
144+
145+
fn is_infallible(&self, options: &ListTransformOptions) -> bool {
146+
match options {
147+
ListTransformOptions::Lambda(lambda) => lambda.is_infallible(),
148+
ListTransformOptions::Bound(lambda) => lambda.is_infallible(),
149+
}
150+
}
151+
}
152+
153+
/// Check that a bound lambda's parameter has the element dtype of a list of this dtype.
154+
fn check_bound_to(lambda: &BoundLambda, list: &DType) -> VortexResult<()> {
155+
let element = element_dtype(list, lambda.params().names().as_ref())?;
156+
vortex_ensure!(
157+
lambda.is_bound_to(&element),
158+
"list_transform()'s lambda {lambda} is bound to element dtype {}, not {}",
159+
lambda.params().fields().format(", "),
160+
element[0]
161+
);
162+
Ok(())
163+
}
164+
165+
/// The dtype of the lambda's one parameter, the element, for a list of this dtype.
166+
fn element_dtype(list: &DType, params: &[FieldName]) -> VortexResult<[DType; 1]> {
167+
let DType::List(element, _) = list else {
168+
vortex_bail!("list_transform() takes a List, got {list}");
169+
};
170+
vortex_ensure!(
171+
params.len() == 1,
172+
"list_transform() takes a lambda of one parameter, the element, not ({})",
173+
params.iter().format(", ")
174+
);
175+
Ok([element.as_ref().clone()])
176+
}
177+
178+
/// The lambda over a list's elements, inside the list's offsets and validity.
179+
///
180+
/// A lambda that can fail runs only on the elements valid lists reference. One that can't runs on
181+
/// every element from the first list to the last: results under null lists are never read.
182+
fn transform_list(
183+
list: ListArray,
184+
lambda: &BoundLambda,
185+
ctx: &mut ExecutionCtx,
186+
) -> VortexResult<ArrayRef> {
187+
let list = if lambda.is_infallible() {
188+
trim_ends(list, ctx)?
189+
} else {
190+
list.compact(ctx)?
191+
};
192+
let values = lambda.evaluate(vec![list.elements().clone()], list.elements().len())?;
193+
vortex_ensure_eq!(values.len(), list.elements().len());
194+
// SAFETY: the offsets and validity are the list's, over one value per element.
195+
Ok(
196+
unsafe { ListArray::new_unchecked(values, list.offsets().clone(), list.list_validity()) }
197+
.into_array(),
198+
)
199+
}
200+
201+
/// The lambda over a list view's elements, inside the view's offsets, sizes and validity.
202+
///
203+
/// A lambda that can fail runs only on the elements valid lists reference.
204+
fn transform_list_view(
205+
view: ListViewArray,
206+
lambda: &BoundLambda,
207+
ctx: &mut ExecutionCtx,
208+
) -> VortexResult<ArrayRef> {
209+
let view = if lambda.is_infallible() {
210+
view
211+
} else {
212+
view.rebuild(ListViewRebuildMode::MakeExactWithEmptyNulls, ctx)?
213+
};
214+
let values = lambda.evaluate(vec![view.elements().clone()], view.elements().len())?;
215+
vortex_ensure_eq!(values.len(), view.elements().len());
216+
// SAFETY: the offsets, sizes and validity are the view's, over one value per element, so the
217+
// views keep their bounds and their layout.
218+
Ok(unsafe {
219+
ListViewArray::new_unchecked(
220+
values,
221+
view.offsets().clone(),
222+
view.sizes().clone(),
223+
view.listview_validity(),
224+
)
225+
.with_zero_copy_to_list(view.is_zero_copy_to_list())
226+
}
227+
.into_array())
228+
}
229+
230+
/// The list without what comes before its first list and after its last.
231+
fn trim_ends(list: ListArray, ctx: &mut ExecutionCtx) -> VortexResult<ListArray> {
232+
if list.offset_at(0, ctx)? == 0 && list.offset_at(list.len(), ctx)? == list.elements().len() {
233+
return Ok(list);
234+
}
235+
list.reset_offsets(false, ctx)
236+
}

0 commit comments

Comments
 (0)