Skip to content

Commit b672f3c

Browse files
committed
list_transform scalar fn
Signed-off-by: Matt Katz <mhkatz97@gmail.com>
1 parent aafd2d3 commit b672f3c

5 files changed

Lines changed: 786 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};
@@ -1382,6 +1384,22 @@ pub fn bound_list_sum_opts(
13821384
.vortex_expect("list-sum expressions require a numeric list child")
13831385
}
13841386

1387+
// ---- ListTransform ----
1388+
1389+
/// Creates an expression that applies a lambda to every element of each list.
1390+
///
1391+
/// The lambda takes one parameter, the element.
1392+
///
1393+
/// ```rust
1394+
/// # use vortex_array::expr::{Lambda, checked_add, col, list_transform, lit, var};
1395+
/// let lambda = Lambda::new(["x"], checked_add(var("x"), lit(1_i32))).unwrap();
1396+
/// let expr = list_transform(col("l"), lambda);
1397+
/// assert_eq!(expr.to_string(), "vortex.list.transform($.l, (x) -> (@x + 1i32))");
1398+
/// ```
1399+
pub fn list_transform(list: Expression, lambda: Lambda) -> Expression {
1400+
ListTransform.new_expr(ListTransformOptions::Lambda(lambda), [list])
1401+
}
1402+
13851403
/// Reinterpret timestamp wall times in `options.time_zone`, preserving the time unit.
13861404
///
13871405
/// `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
@@ -92,6 +92,7 @@ pub use exprs::list_contains_opts;
9292
pub use exprs::list_length;
9393
pub use exprs::list_sum;
9494
pub use exprs::list_sum_opts;
95+
pub use exprs::list_transform;
9596
pub use exprs::lit;
9697
pub use exprs::lt;
9798
pub use exprs::lt_eq;
Lines changed: 221 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,221 @@
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::ListArray;
23+
use crate::arrays::ListViewArray;
24+
use crate::arrays::PrimitiveArray;
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::list_from_list_view;
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::match_each_integer_ptype;
36+
use crate::scalar_fn::Arity;
37+
use crate::scalar_fn::ChildName;
38+
use crate::scalar_fn::ExecutionArgs;
39+
use crate::scalar_fn::ScalarFnId;
40+
use crate::scalar_fn::ScalarFnVTable;
41+
use crate::validity::Validity;
42+
43+
/// Apply a lambda to every element of each list, keeping each list's length and validity.
44+
///
45+
/// The lambda takes one parameter, the element. The only child is the list, of `List` dtype in any
46+
/// encoding, and the result is a `List`. The lambda runs only on the elements of valid lists.
47+
#[derive(Clone)]
48+
pub struct ListTransform;
49+
50+
/// The lambda of a [`ListTransform`], as written or bound to the element's dtype.
51+
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
52+
pub enum ListTransformOptions {
53+
/// The lambda as written, in an [`Expression`](crate::expr::Expression).
54+
Lambda(Lambda),
55+
/// The lambda bound to the element's dtype, in a
56+
/// [`BoundExpression`](crate::expr::BoundExpression) or an array.
57+
Bound(BoundLambda),
58+
}
59+
60+
impl Display for ListTransformOptions {
61+
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
62+
match self {
63+
Self::Lambda(lambda) => lambda.fmt(f),
64+
Self::Bound(lambda) => lambda.fmt(f),
65+
}
66+
}
67+
}
68+
69+
impl ScalarFnVTable for ListTransform {
70+
type Options = ListTransformOptions;
71+
72+
fn id(&self) -> ScalarFnId {
73+
static ID: CachedId = CachedId::new("vortex.list.transform");
74+
*ID
75+
}
76+
77+
fn arity(&self, _options: &ListTransformOptions) -> Arity {
78+
Arity::Exact(1)
79+
}
80+
81+
fn child_name(&self, _options: &ListTransformOptions, child_idx: usize) -> ChildName {
82+
match child_idx {
83+
0 => ChildName::from("list"),
84+
_ => unreachable!("Invalid child index {child_idx} for list_transform()"),
85+
}
86+
}
87+
88+
fn fmt_sql(
89+
&self,
90+
options: &ListTransformOptions,
91+
expr: &dyn ExprDisplay,
92+
f: &mut Formatter<'_>,
93+
) -> fmt::Result {
94+
write!(f, "{}({}, {options})", self.id(), expr.display_child(0))
95+
}
96+
97+
fn return_dtype(&self, options: &ListTransformOptions, args: &[DType]) -> VortexResult<DType> {
98+
let element = match options {
99+
ListTransformOptions::Lambda(lambda) => {
100+
lambda.return_dtype(&element_dtype(&args[0], lambda.params())?)?
101+
}
102+
ListTransformOptions::Bound(lambda) => {
103+
check_bound_to(lambda, &args[0])?;
104+
lambda.dtype().clone()
105+
}
106+
};
107+
Ok(DType::List(Arc::new(element), args[0].nullability()))
108+
}
109+
110+
fn bind_options(
111+
&self,
112+
options: &ListTransformOptions,
113+
args: &[DType],
114+
) -> VortexResult<Option<ListTransformOptions>> {
115+
match options {
116+
ListTransformOptions::Lambda(lambda) => {
117+
let element = element_dtype(&args[0], lambda.params())?;
118+
Ok(Some(ListTransformOptions::Bound(lambda.bind(&element)?)))
119+
}
120+
ListTransformOptions::Bound(lambda) => {
121+
check_bound_to(lambda, &args[0])?;
122+
Ok(None)
123+
}
124+
}
125+
}
126+
127+
fn execute(
128+
&self,
129+
options: &ListTransformOptions,
130+
args: &dyn ExecutionArgs,
131+
ctx: &mut ExecutionCtx,
132+
) -> VortexResult<ArrayRef> {
133+
let ListTransformOptions::Bound(lambda) = options else {
134+
vortex_bail!("list_transform()'s lambda is not bound");
135+
};
136+
// Every element belongs to exactly one valid list, in order: nothing comes before the
137+
// first list or after the last, and a null list has no elements.
138+
let list = args.get(0)?.execute::<ListViewArray>(ctx)?;
139+
let list = list_from_list_view(empty_null_lists(list, ctx)?, ctx)?;
140+
let elements = list.elements();
141+
let values = lambda.evaluate(vec![elements.clone()], elements.len())?;
142+
vortex_ensure_eq!(values.len(), elements.len());
143+
// SAFETY: the offsets are the list's, over one value per element.
144+
Ok(unsafe {
145+
ListArray::new_unchecked(values, list.offsets().clone(), list.list_validity())
146+
}
147+
.into_array())
148+
}
149+
150+
fn is_strict(&self, _options: &ListTransformOptions) -> bool {
151+
true
152+
}
153+
154+
fn is_infallible(&self, options: &ListTransformOptions) -> bool {
155+
match options {
156+
ListTransformOptions::Lambda(lambda) => lambda.is_infallible(),
157+
ListTransformOptions::Bound(lambda) => lambda.is_infallible(),
158+
}
159+
}
160+
}
161+
162+
/// Check that a bound lambda's parameter has the element dtype of a list of this dtype.
163+
fn check_bound_to(lambda: &BoundLambda, list: &DType) -> VortexResult<()> {
164+
let element = element_dtype(list, lambda.params().names().as_ref())?;
165+
vortex_ensure!(
166+
lambda.is_bound_to(&element),
167+
"list_transform()'s lambda {lambda} is bound to element dtype {}, not {}",
168+
lambda.params().fields().format(", "),
169+
element[0]
170+
);
171+
Ok(())
172+
}
173+
174+
/// The dtype of the lambda's one parameter, the element, for a list of this dtype.
175+
fn element_dtype(list: &DType, params: &[FieldName]) -> VortexResult<[DType; 1]> {
176+
let DType::List(element, _) = list else {
177+
vortex_bail!("list_transform() takes a List, got {list}");
178+
};
179+
vortex_ensure!(
180+
params.len() == 1,
181+
"list_transform() takes a lambda of one parameter, the element, not ({})",
182+
params.iter().format(", ")
183+
);
184+
Ok([element.as_ref().clone()])
185+
}
186+
187+
/// The lists with every null list emptied, so that nothing under a null list reaches the lambda.
188+
fn empty_null_lists(list: ListViewArray, ctx: &mut ExecutionCtx) -> VortexResult<ListViewArray> {
189+
let validity = list.listview_validity();
190+
if matches!(validity, Validity::NonNullable | Validity::AllValid) {
191+
return Ok(list);
192+
}
193+
let valid = validity.execute_mask(list.len(), ctx)?;
194+
let sizes = list.sizes().clone().execute::<PrimitiveArray>(ctx)?;
195+
let sizes = match_each_integer_ptype!(sizes.ptype(), |T| {
196+
let sizes = sizes.as_slice::<T>();
197+
let null = |index: usize| !valid.value(index);
198+
// Writers usually leave null lists empty already.
199+
if !(0..sizes.len()).any(|index| null(index) && sizes[index] != T::default()) {
200+
return Ok(list);
201+
}
202+
PrimitiveArray::from_iter((0..sizes.len()).map(|index| {
203+
if null(index) {
204+
T::default()
205+
} else {
206+
sizes[index]
207+
}
208+
}))
209+
});
210+
// SAFETY: only sizes changed, and only to 0, so every view stays within the elements. The new
211+
// array is not marked zero-copy to a list: an emptied list leaves its elements behind as a gap,
212+
// which the rebuild must drop.
213+
Ok(unsafe {
214+
ListViewArray::new_unchecked(
215+
list.elements().clone(),
216+
list.offsets().clone(),
217+
sizes.into_array(),
218+
validity,
219+
)
220+
})
221+
}

0 commit comments

Comments
 (0)