vortex_array/arrays/scalar_fn/
rules.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::Any;
5use std::sync::Arc;
6
7use itertools::Itertools;
8use vortex_error::VortexExpect;
9use vortex_error::VortexResult;
10
11use crate::ArrayRef;
12use crate::IntoArray;
13use crate::array::ArrayView;
14use crate::arrays::Constant;
15use crate::arrays::ConstantArray;
16use crate::arrays::Filter;
17use crate::arrays::ScalarFn;
18use crate::arrays::ScalarFnArray;
19use crate::arrays::Slice;
20use crate::arrays::StructArray;
21use crate::arrays::scalar_fn::ScalarFnArrayExt;
22use crate::dtype::DType;
23use crate::optimizer::rules::ArrayParentReduceRule;
24use crate::optimizer::rules::ArrayReduceRule;
25use crate::optimizer::rules::ParentRuleSet;
26use crate::optimizer::rules::ReduceRuleSet;
27use crate::scalar_fn::ReduceCtx;
28use crate::scalar_fn::ReduceNode;
29use crate::scalar_fn::ReduceNodeRef;
30use crate::scalar_fn::ScalarFnRef;
31use crate::scalar_fn::fns::pack::Pack;
32use crate::validity::Validity;
33
34pub(super) const RULES: ReduceRuleSet<ScalarFn> =
35    ReduceRuleSet::new(&[&ScalarFnPackToStructRule, &ScalarFnAbstractReduceRule]);
36
37pub(super) const PARENT_RULES: ParentRuleSet<ScalarFn> = ParentRuleSet::new(&[
38    ParentRuleSet::lift(&ScalarFnUnaryFilterPushDownRule),
39    ParentRuleSet::lift(&ScalarFnSliceReduceRule),
40]);
41
42/// Converts a ScalarFnArray with Pack into a StructArray directly.
43#[derive(Debug)]
44struct ScalarFnPackToStructRule;
45impl ArrayReduceRule<ScalarFn> for ScalarFnPackToStructRule {
46    fn reduce(&self, array: ArrayView<'_, ScalarFn>) -> VortexResult<Option<ArrayRef>> {
47        let Some(pack_options) = array.scalar_fn().as_opt::<Pack>() else {
48            return Ok(None);
49        };
50
51        let validity = match pack_options.nullability {
52            crate::dtype::Nullability::NonNullable => Validity::NonNullable,
53            crate::dtype::Nullability::Nullable => Validity::AllValid,
54        };
55
56        Ok(Some(
57            StructArray::try_new(
58                pack_options.names.clone(),
59                array.children(),
60                array.len(),
61                validity,
62            )?
63            .into_array(),
64        ))
65    }
66}
67
68#[derive(Debug)]
69struct ScalarFnSliceReduceRule;
70impl ArrayParentReduceRule<ScalarFn> for ScalarFnSliceReduceRule {
71    type Parent = Slice;
72
73    fn reduce_parent(
74        &self,
75        array: ArrayView<'_, ScalarFn>,
76        parent: ArrayView<'_, Slice>,
77        _child_idx: usize,
78    ) -> VortexResult<Option<ArrayRef>> {
79        let range = parent.slice_range();
80
81        let children: Vec<_> = array
82            .iter_children()
83            .map(|c| c.slice(range.clone()))
84            .collect::<VortexResult<_>>()?;
85
86        Ok(Some(
87            ScalarFnArray::try_new_with_len(array.scalar_fn().clone(), children, range.len())?
88                .into_array(),
89        ))
90    }
91}
92
93#[derive(Debug)]
94struct ScalarFnAbstractReduceRule;
95impl ArrayReduceRule<ScalarFn> for ScalarFnAbstractReduceRule {
96    fn reduce(&self, array: ArrayView<'_, ScalarFn>) -> VortexResult<Option<ArrayRef>> {
97        if let Some(reduced) = array
98            .scalar_fn()
99            .reduce(array.as_ref(), &ArrayReduceCtx { len: array.len() })?
100        {
101            return Ok(Some(
102                reduced
103                    .as_any()
104                    .downcast_ref::<ArrayRef>()
105                    .vortex_expect("ReduceNode is not an ArrayRef")
106                    .clone(),
107            ));
108        }
109        Ok(None)
110    }
111}
112
113impl ReduceNode for ArrayRef {
114    fn as_any(&self) -> &dyn Any {
115        self
116    }
117
118    fn node_dtype(&self) -> VortexResult<DType> {
119        Ok(self.dtype().clone())
120    }
121
122    fn scalar_fn(&self) -> Option<&ScalarFnRef> {
123        self.as_opt::<ScalarFn>().map(|a| a.data().scalar_fn())
124    }
125
126    fn child(&self, idx: usize) -> ReduceNodeRef {
127        Arc::new(self.nth_child(idx).vortex_expect("child idx out of bounds"))
128    }
129
130    fn child_count(&self) -> usize {
131        self.nchildren()
132    }
133}
134
135struct ArrayReduceCtx {
136    // The length of the array being reduced
137    len: usize,
138}
139impl ReduceCtx for ArrayReduceCtx {
140    fn new_node(
141        &self,
142        scalar_fn: ScalarFnRef,
143        children: &[ReduceNodeRef],
144    ) -> VortexResult<ReduceNodeRef> {
145        Ok(Arc::new(
146            ScalarFnArray::try_new_with_len(
147                scalar_fn,
148                children
149                    .iter()
150                    .map(|c| {
151                        c.as_any()
152                            .downcast_ref::<ArrayRef>()
153                            .vortex_expect("ReduceNode is not an ArrayRef")
154                            .clone()
155                    })
156                    .collect(),
157                self.len,
158            )?
159            .into_array(),
160        ))
161    }
162}
163
164#[derive(Debug)]
165struct ScalarFnUnaryFilterPushDownRule;
166
167impl ArrayParentReduceRule<ScalarFn> for ScalarFnUnaryFilterPushDownRule {
168    type Parent = Filter;
169
170    fn reduce_parent(
171        &self,
172        child: ArrayView<'_, ScalarFn>,
173        parent: ArrayView<'_, Filter>,
174        _child_idx: usize,
175    ) -> VortexResult<Option<ArrayRef>> {
176        // If we only have one non-constant child, then it is _always_ cheaper to push down the
177        // filter over the children of the scalar function array.
178        if child
179            .iter_children()
180            .filter(|c| !c.is::<Constant>())
181            .count()
182            == 1
183        {
184            let new_children: Vec<_> = child
185                .iter_children()
186                .map(|c| match c.as_opt::<Constant>() {
187                    Some(array) => {
188                        Ok(ConstantArray::new(array.scalar().clone(), parent.len()).into_array())
189                    }
190                    None => c.filter(parent.filter_mask().clone()),
191                })
192                .try_collect()?;
193
194            let new_array =
195                ScalarFnArray::try_new(child.scalar_fn().clone(), new_children)?.into_array();
196
197            return Ok(Some(new_array));
198        }
199
200        Ok(None)
201    }
202}
203
204#[cfg(test)]
205mod tests {
206    use vortex_error::VortexExpect;
207
208    use crate::array::IntoArray;
209    use crate::arrays::ChunkedArray;
210    use crate::arrays::PrimitiveArray;
211    use crate::arrays::scalar_fn::rules::ConstantArray;
212    use crate::dtype::DType;
213    use crate::dtype::Nullability;
214    use crate::dtype::PType;
215    use crate::expr::cast;
216    use crate::expr::is_null;
217    use crate::expr::root;
218
219    #[test]
220    fn test_empty_constants() {
221        let array = ChunkedArray::try_new(
222            vec![
223                ConstantArray::new(Some(1u64), 0).into_array(),
224                PrimitiveArray::from_iter(vec![2u64])
225                    .into_array()
226                    .apply(&cast(
227                        root(),
228                        DType::Primitive(PType::U64, Nullability::Nullable),
229                    ))
230                    .vortex_expect("casted"),
231            ],
232            DType::Primitive(PType::U64, Nullability::Nullable),
233        )
234        .vortex_expect("construction")
235        .into_array();
236
237        let expr = is_null(root());
238        array.apply(&expr).vortex_expect("expr evaluation");
239    }
240}
vortex_array/arrays/scalar_fn/rules.rs

vortex_array/arrays/scalar_fn/
rules.rs
