vortex_array/arrays/scalar_fn/
rules.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use itertools::Itertools;
5use vortex_error::VortexResult;
6
7use crate::ArrayRef;
8use crate::IntoArray;
9use crate::array::ArrayView;
10use crate::arrays::Constant;
11use crate::arrays::ConstantArray;
12use crate::arrays::Filter;
13use crate::arrays::ScalarFn;
14use crate::arrays::ScalarFnArray;
15use crate::arrays::Slice;
16use crate::arrays::StructArray;
17use crate::arrays::scalar_fn::ScalarFnArrayExt;
18use crate::optimizer::rules::ArrayParentReduceRule;
19use crate::optimizer::rules::ArrayReduceRule;
20use crate::optimizer::rules::ParentRuleSet;
21use crate::optimizer::rules::ReduceRuleSet;
22use crate::scalar_fn::ArrayReduceNode;
23use crate::scalar_fn::fns::pack::Pack;
24use crate::validity::Validity;
25
26pub(super) const RULES: ReduceRuleSet<ScalarFn> =
27    ReduceRuleSet::new(&[&ScalarFnPackToStructRule, &ScalarFnAbstractReduceRule]);
28
29pub(super) const PARENT_RULES: ParentRuleSet<ScalarFn> = ParentRuleSet::new(&[
30    ParentRuleSet::lift(&ScalarFnUnaryFilterPushDownRule),
31    ParentRuleSet::lift(&ScalarFnSliceReduceRule),
32]);
33
34/// Converts a ScalarFnArray with Pack into a StructArray directly.
35#[derive(Debug)]
36struct ScalarFnPackToStructRule;
37impl ArrayReduceRule<ScalarFn> for ScalarFnPackToStructRule {
38    fn reduce(&self, array: ArrayView<'_, ScalarFn>) -> VortexResult<Option<ArrayRef>> {
39        let Some(pack_options) = array.scalar_fn().as_opt::<Pack>() else {
40            return Ok(None);
41        };
42
43        let validity = match pack_options.nullability {
44            crate::dtype::Nullability::NonNullable => Validity::NonNullable,
45            crate::dtype::Nullability::Nullable => Validity::AllValid,
46        };
47
48        Ok(Some(
49            StructArray::try_new(
50                pack_options.names.clone(),
51                array.children(),
52                array.len(),
53                validity,
54            )?
55            .into_array(),
56        ))
57    }
58}
59
60#[derive(Debug)]
61struct ScalarFnSliceReduceRule;
62impl ArrayParentReduceRule<ScalarFn> for ScalarFnSliceReduceRule {
63    type Parent = Slice;
64
65    fn reduce_parent(
66        &self,
67        array: ArrayView<'_, ScalarFn>,
68        parent: ArrayView<'_, Slice>,
69        _child_idx: usize,
70    ) -> VortexResult<Option<ArrayRef>> {
71        let range = parent.slice_range();
72
73        let children: Vec<_> = array
74            .iter_children()
75            .map(|c| c.slice(range.clone()))
76            .collect::<VortexResult<_>>()?;
77
78        Ok(Some(
79            ScalarFnArray::try_new_with_len(array.scalar_fn().clone(), children, range.len())?
80                .into_array(),
81        ))
82    }
83}
84
85#[derive(Debug)]
86struct ScalarFnAbstractReduceRule;
87impl ArrayReduceRule<ScalarFn> for ScalarFnAbstractReduceRule {
88    fn reduce(&self, array: ArrayView<'_, ScalarFn>) -> VortexResult<Option<ArrayRef>> {
89        let node = ArrayReduceNode::new(array.as_ref());
90        if let Some(reduced) = array.scalar_fn().reduce_array(&node)? {
91            return Ok(Some(reduced.into_array()));
92        }
93        Ok(None)
94    }
95}
96
97#[derive(Debug)]
98struct ScalarFnUnaryFilterPushDownRule;
99
100impl ArrayParentReduceRule<ScalarFn> for ScalarFnUnaryFilterPushDownRule {
101    type Parent = Filter;
102
103    fn reduce_parent(
104        &self,
105        child: ArrayView<'_, ScalarFn>,
106        parent: ArrayView<'_, Filter>,
107        _child_idx: usize,
108    ) -> VortexResult<Option<ArrayRef>> {
109        // If we only have one non-constant child, then it is _always_ cheaper to push down the
110        // filter over the children of the scalar function array.
111        if child
112            .iter_children()
113            .filter(|c| !c.is::<Constant>())
114            .count()
115            == 1
116        {
117            let new_children: Vec<_> = child
118                .iter_children()
119                .map(|c| match c.as_opt::<Constant>() {
120                    Some(array) => {
121                        Ok(ConstantArray::new(array.scalar().clone(), parent.len()).into_array())
122                    }
123                    None => c.filter(parent.filter_mask().clone()),
124                })
125                .try_collect()?;
126
127            let new_array =
128                ScalarFnArray::try_new(child.scalar_fn().clone(), new_children)?.into_array();
129
130            return Ok(Some(new_array));
131        }
132
133        Ok(None)
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    use vortex_error::VortexExpect;
140
141    use crate::array::IntoArray;
142    use crate::arrays::ChunkedArray;
143    use crate::arrays::PrimitiveArray;
144    use crate::arrays::scalar_fn::rules::ConstantArray;
145    use crate::dtype::DType;
146    use crate::dtype::Nullability;
147    use crate::dtype::PType;
148    use crate::expr::cast;
149    use crate::expr::is_null;
150    use crate::expr::root;
151
152    #[test]
153    fn test_empty_constants() {
154        let array = ChunkedArray::try_new(
155            vec![
156                ConstantArray::new(Some(1u64), 0).into_array(),
157                PrimitiveArray::from_iter(vec![2u64])
158                    .into_array()
159                    .apply(&cast(
160                        root(),
161                        DType::Primitive(PType::U64, Nullability::Nullable),
162                    ))
163                    .vortex_expect("casted"),
164            ],
165            DType::Primitive(PType::U64, Nullability::Nullable),
166        )
167        .vortex_expect("construction")
168        .into_array();
169
170        let expr = is_null(root());
171        array.apply(&expr).vortex_expect("expr evaluation");
172    }
173}
vortex_array/arrays/scalar_fn/rules.rs

vortex_array/arrays/scalar_fn/
rules.rs
