Skip to content

Commit 58db9d3

Browse files
mhk197mhk197
committed
higher-order functions as their own expression and array kind
Signed-off-by: Matt Katz <[email protected]>
1 parent dbaf5f8 commit 58db9d3

21 files changed

Lines changed: 2548 additions & 22 deletions

File tree

‎vortex-array/src/arrays/higher_order/mod.rs‎

Lines changed: 299 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,299 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
//! An array that applies a higher-order function to its children, with its lambdas bound.
5+
6+
mod rules;
7+
8+
use std::fmt::Display;
9+
use std::fmt::Formatter;
10+
use std::hash::Hash;
11+
use std::hash::Hasher;
12+
use std::sync::Arc;
13+
14+
use itertools::Itertools;
15+
use vortex_error::VortexExpect;
16+
use vortex_error::VortexResult;
17+
use vortex_error::vortex_bail;
18+
use vortex_error::vortex_ensure;
19+
use vortex_error::vortex_ensure_eq;
20+
use vortex_error::vortex_panic;
21+
use vortex_session::VortexSession;
22+
23+
use crate::ArrayEq;
24+
use crate::ArrayHash;
25+
use crate::ArrayRef;
26+
use crate::ArraySlots;
27+
use crate::EqMode;
28+
use crate::ExecutionCtx;
29+
use crate::IntoArray;
30+
use crate::ProbeState;
31+
use crate::array::Array;
32+
use crate::array::ArrayId;
33+
use crate::array::ArrayParts;
34+
use crate::array::ArrayView;
35+
use crate::array::OperationsVTable;
36+
use crate::array::TypedArrayRef;
37+
use crate::array::VTable;
38+
use crate::array::ValidityVTable;
39+
use crate::array::with_empty_buffers;
40+
use crate::arrays::ConstantArray;
41+
use crate::buffer::BufferHandle;
42+
use crate::columnar::Columnar;
43+
use crate::dtype::DType;
44+
use crate::executor::ExecutionResult;
45+
use crate::expr::BoundLambda;
46+
use crate::higher_order_fn::HigherOrderFnRef;
47+
use crate::scalar::Scalar;
48+
use crate::scalar_fn::ScalarFnId;
49+
use crate::scalar_fn::VecExecutionArgs;
50+
use crate::serde::ArrayChildren;
51+
use crate::validity::Validity;
52+
53+
/// An array that applies a higher-order function to its children.
54+
pub type HigherOrderArray = Array<HigherOrder>;
55+
56+
/// The array vtable of [`HigherOrderArray`]. Its id is the function's.
57+
#[derive(Clone, Debug)]
58+
pub struct HigherOrder {
59+
id: ScalarFnId,
60+
}
61+
62+
/// The function and its bound lambdas.
63+
#[derive(Clone, Debug)]
64+
pub struct HigherOrderData {
65+
function: HigherOrderFnRef,
66+
lambdas: Arc<Vec<BoundLambda>>,
67+
}
68+
69+
impl HigherOrderData {
70+
/// The function.
71+
pub fn function(&self) -> &HigherOrderFnRef {
72+
&self.function
73+
}
74+
75+
/// The bound lambdas.
76+
pub fn lambdas(&self) -> &[BoundLambda] {
77+
&self.lambdas
78+
}
79+
}
80+
81+
impl Display for HigherOrderData {
82+
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
83+
write!(
84+
f,
85+
"higher_order_fn: {} {}",
86+
self.function,
87+
self.lambdas.iter().format(", ")
88+
)
89+
}
90+
}
91+
92+
impl ArrayHash for HigherOrderData {
93+
fn array_hash<H: Hasher>(&self, state: &mut H, _accuracy: EqMode) {
94+
self.function.hash(state);
95+
self.lambdas.hash(state);
96+
}
97+
}
98+
99+
impl ArrayEq for HigherOrderData {
100+
fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool {
101+
self.function == other.function && self.lambdas == other.lambdas
102+
}
103+
}
104+
105+
impl Array<HigherOrder> {
106+
/// Apply a function, with its lambdas bound, to these children of `len` rows.
107+
pub fn try_new(
108+
function: HigherOrderFnRef,
109+
lambdas: Arc<Vec<BoundLambda>>,
110+
children: impl IntoIterator<Item = ArrayRef>,
111+
len: usize,
112+
) -> VortexResult<Self> {
113+
let slots: ArraySlots = children.into_iter().map(Some).collect();
114+
vortex_ensure!(
115+
function.arity().matches(slots.len()),
116+
"{} requires {} children, got {}",
117+
function.id(),
118+
function.arity(),
119+
slots.len()
120+
);
121+
vortex_ensure!(
122+
slots.iter().flatten().all(|child| child.len() == len),
123+
"HigherOrderArray requires children of length {len}"
124+
);
125+
let arg_dtypes = slots
126+
.iter()
127+
.flatten()
128+
.map(|child| child.dtype().clone())
129+
.collect_vec();
130+
let dtype = function.return_dtype(&arg_dtypes, &lambdas)?;
131+
let vtable = HigherOrder { id: function.id() };
132+
let data = HigherOrderData { function, lambdas };
133+
Ok(
134+
unsafe {
135+
Array::from_parts_unchecked(ArrayParts::new(vtable, dtype, len, data, slots))
136+
},
137+
)
138+
}
139+
}
140+
141+
/// Accessors for a [`HigherOrderArray`].
142+
pub trait HigherOrderArrayExt: TypedArrayRef<HigherOrder> {
143+
/// The value children.
144+
fn children(&self) -> Vec<ArrayRef> {
145+
self.slots()
146+
.iter()
147+
.map(|slot| {
148+
slot.as_ref()
149+
.vortex_expect("HigherOrderArray child slot")
150+
.clone()
151+
})
152+
.collect()
153+
}
154+
}
155+
156+
impl<T: TypedArrayRef<HigherOrder>> HigherOrderArrayExt for T {}
157+
158+
impl VTable for HigherOrder {
159+
type TypedArrayData = HigherOrderData;
160+
type OperationsVTable = Self;
161+
type ValidityVTable = Self;
162+
163+
fn id(&self) -> ArrayId {
164+
self.id
165+
}
166+
167+
fn validate(
168+
&self,
169+
data: &HigherOrderData,
170+
dtype: &DType,
171+
len: usize,
172+
slots: &[Option<ArrayRef>],
173+
) -> VortexResult<()> {
174+
vortex_ensure_eq!(data.function.id(), self.id);
175+
vortex_ensure!(slots.iter().all(Option::is_some));
176+
vortex_ensure!(slots.iter().flatten().all(|child| child.len() == len));
177+
let arg_dtypes = slots
178+
.iter()
179+
.flatten()
180+
.map(|child| child.dtype().clone())
181+
.collect_vec();
182+
vortex_ensure_eq!(
183+
data.function.return_dtype(&arg_dtypes, &data.lambdas)?,
184+
*dtype
185+
);
186+
Ok(())
187+
}
188+
189+
fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
190+
0
191+
}
192+
193+
fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
194+
vortex_panic!("HigherOrderArray buffer index {idx} out of bounds")
195+
}
196+
197+
fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option<String> {
198+
None
199+
}
200+
201+
fn with_buffers(
202+
&self,
203+
array: ArrayView<'_, Self>,
204+
buffers: &[BufferHandle],
205+
) -> VortexResult<ArrayParts<Self>> {
206+
with_empty_buffers(self, array, buffers)
207+
}
208+
209+
fn serialize(
210+
_array: ArrayView<'_, Self>,
211+
_session: &VortexSession,
212+
) -> VortexResult<Option<Vec<u8>>> {
213+
Ok(None)
214+
}
215+
216+
fn deserialize(
217+
&self,
218+
_dtype: &DType,
219+
_len: usize,
220+
_metadata: &[u8],
221+
_buffers: &[BufferHandle],
222+
_children: &dyn ArrayChildren,
223+
_session: &VortexSession,
224+
) -> VortexResult<ArrayParts<Self>> {
225+
vortex_bail!("HigherOrderArray is not serializable")
226+
}
227+
228+
fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String {
229+
array.function.child_name(idx).as_ref().to_string()
230+
}
231+
232+
fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
233+
let args = VecExecutionArgs::new(array.children(), array.len());
234+
array
235+
.function
236+
.execute(&args, &array.lambdas, ctx)
237+
.map(ExecutionResult::done)
238+
}
239+
240+
fn reduce_parent(
241+
array: ArrayView<'_, Self>,
242+
parent: &ArrayRef,
243+
child_idx: usize,
244+
) -> VortexResult<Option<ArrayRef>> {
245+
rules::PARENT_RULES.evaluate(array, parent, child_idx)
246+
}
247+
}
248+
249+
impl OperationsVTable<HigherOrder> for HigherOrder {
250+
type ProbeState = ();
251+
252+
fn probe_scalar(
253+
state: &mut ProbeState<'_, HigherOrder>,
254+
index: usize,
255+
ctx: &mut ExecutionCtx,
256+
) -> VortexResult<Scalar> {
257+
Self::scalar_at(state.array(), index, ctx)
258+
}
259+
260+
fn scalar_at(
261+
array: ArrayView<'_, HigherOrder>,
262+
index: usize,
263+
ctx: &mut ExecutionCtx,
264+
) -> VortexResult<Scalar> {
265+
let inputs: Vec<_> = array
266+
.children()
267+
.iter()
268+
.map(|child| Ok(ConstantArray::new(child.execute_scalar(index, ctx)?, 1).into_array()))
269+
.collect::<VortexResult<_>>()?;
270+
let result =
271+
array
272+
.function
273+
.execute(&VecExecutionArgs::new(inputs, 1), &array.lambdas, ctx)?;
274+
Ok(match result.execute::<Columnar>(ctx)? {
275+
Columnar::Canonical(canonical) => canonical.into_array().execute_scalar(0, ctx)?,
276+
Columnar::Constant(constant) => constant.scalar().clone(),
277+
})
278+
}
279+
}
280+
281+
impl ValidityVTable<HigherOrder> for HigherOrder {
282+
fn validity(view: ArrayView<'_, HigherOrder>) -> VortexResult<Validity> {
283+
if view.function.is_strict() {
284+
// Valid where every child is.
285+
return view
286+
.children()
287+
.iter()
288+
.try_fold(Validity::from(view.dtype().nullability()), |acc, child| {
289+
acc.and(child.validity()?)
290+
});
291+
}
292+
// Prototype: a function that isn't strict would need a lazy validity, like the one
293+
// scalar function arrays build.
294+
vortex_bail!(
295+
"validity of a non-strict {} is not supported yet",
296+
view.function.id()
297+
)
298+
}
299+
}

‎vortex-array/src/arrays/higher_order/rules.rs‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
//! Slice and filter pushdown for [`HigherOrderArray`], as scalar function
5+
//! arrays have: selection comes before evaluation, so a lambda never runs on unselected rows.
6+
7+
use std::sync::Arc;
8+
9+
use itertools::Itertools;
10+
use vortex_error::VortexResult;
11+
12+
use super::HigherOrder;
13+
use super::HigherOrderArray;
14+
use super::HigherOrderArrayExt;
15+
use crate::ArrayRef;
16+
use crate::IntoArray;
17+
use crate::array::ArrayView;
18+
use crate::arrays::Filter;
19+
use crate::arrays::Slice;
20+
use crate::optimizer::rules::ArrayParentReduceRule;
21+
use crate::optimizer::rules::ParentRuleSet;
22+
23+
pub(super) const PARENT_RULES: ParentRuleSet<HigherOrder> = ParentRuleSet::new(&[
24+
ParentRuleSet::lift(&HigherOrderFilterPushdownRule),
25+
ParentRuleSet::lift(&HigherOrderSliceReduceRule),
26+
]);
27+
28+
#[derive(Debug)]
29+
struct HigherOrderSliceReduceRule;
30+
31+
impl ArrayParentReduceRule<HigherOrder> for HigherOrderSliceReduceRule {
32+
type Parent = Slice;
33+
34+
fn reduce_parent(
35+
&self,
36+
array: ArrayView<'_, HigherOrder>,
37+
parent: ArrayView<'_, Slice>,
38+
_child_idx: usize,
39+
) -> VortexResult<Option<ArrayRef>> {
40+
let range = parent.slice_range();
41+
let children: Vec<_> = array
42+
.children()
43+
.iter()
44+
.map(|child| child.slice(range.clone()))
45+
.try_collect()?;
46+
Ok(Some(
47+
HigherOrderArray::try_new(
48+
array.function.clone(),
49+
Arc::clone(&array.lambdas),
50+
children,
51+
range.len(),
52+
)?
53+
.into_array(),
54+
))
55+
}
56+
}
57+
58+
#[derive(Debug)]
59+
struct HigherOrderFilterPushdownRule;
60+
61+
impl ArrayParentReduceRule<HigherOrder> for HigherOrderFilterPushdownRule {
62+
type Parent = Filter;
63+
64+
fn reduce_parent(
65+
&self,
66+
array: ArrayView<'_, HigherOrder>,
67+
parent: ArrayView<'_, Filter>,
68+
_child_idx: usize,
69+
) -> VortexResult<Option<ArrayRef>> {
70+
let children: Vec<_> = array
71+
.children()
72+
.iter()
73+
.map(|child| child.filter(parent.filter_mask().clone()))
74+
.try_collect()?;
75+
Ok(Some(
76+
HigherOrderArray::try_new(
77+
array.function.clone(),
78+
Arc::clone(&array.lambdas),
79+
children,
80+
parent.len(),
81+
)?
82+
.into_array(),
83+
))
84+
}
85+
}

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,9 @@ pub use primitive::PrimitiveArray;
105105
pub mod scalar_fn;
106106
pub use scalar_fn::ScalarFn;
107107
pub use scalar_fn::ScalarFnArray;
108+
pub mod higher_order;
109+
pub use higher_order::HigherOrder;
110+
pub use higher_order::HigherOrderArray;
108111

109112
pub mod slice;
110113
pub use slice::Slice;

0 commit comments

Comments
 (0)