vortex_array/
optimizer.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use vortex_error::VortexResult;
7
8use crate::vtable::VTable;
9use crate::{Array, ArrayRef};
10
11impl dyn Array + '_ {
12    /// Optimize this array by applying optimization rules recursively to its children in a single
13    /// bottom-up pass.
14    pub fn optimize(&self) -> VortexResult<ArrayRef> {
15        let slf = self.to_array();
16        let children = self.children();
17
18        let mut new_children = Vec::with_capacity(children.len());
19        let mut children_modified = false;
20        for (idx, child) in children.iter().enumerate() {
21            let child = child.optimize()?;
22
23            // Check if the child can reduce us (its parent), and if so bail early.
24            if let Some(reduced) = child.reduce_parent(&slf, idx)? {
25                return Ok(reduced);
26            }
27
28            if !Arc::ptr_eq(&child, &children[idx]) {
29                children_modified = true;
30            }
31            new_children.push(child);
32        }
33
34        if children_modified {
35            return self.with_children(&new_children);
36        }
37
38        Ok(slf)
39    }
40}
41
42/// An optimizer rule that tries to reduce/replace a parent array where the implementer is a
43/// child array in the `CHILD_IDX` position of the parent array.
44pub trait ReduceParent<Parent: VTable, const CHILD_IDX: usize>: VTable {
45    /// Try to reduce/replace the given parent array based on this child array.
46    ///
47    /// If no reduction is possible, return None.
48    fn reduce_parent(array: &Self::Array, parent: &Parent::Array)
49    -> VortexResult<Option<ArrayRef>>;
50}
51
52/// A generic optimizer rule that can be applied to an array to try to optimize it.
53pub trait OptimizerRule {
54    /// Try to optimize the given array, returning a replacement if successful.
55    ///
56    /// If no optimization is possible, return None.
57    fn optimize(&self, array: &ArrayRef) -> VortexResult<Option<ArrayRef>>;
58}
59
60#[cfg(test)]
61mod tests {
62    use vortex_buffer::{bitbuffer, buffer};
63    use vortex_dtype::PTypeDowncast;
64    use vortex_vector::VectorOps;
65
66    use crate::IntoArray;
67    use crate::arrays::{BoolArray, MaskedArray, PrimitiveArray};
68    use crate::validity::Validity;
69
70    #[test]
71    fn test_masked_pushdown() {
72        let array = PrimitiveArray::from_iter([0u32, 1, 2, 3]);
73        assert!(!array.dtype().is_nullable());
74
75        let masked = MaskedArray::try_new(
76            array.into_array(),
77            Validity::Array(BoolArray::from(bitbuffer![0 1 0 1]).into_array()),
78        )
79        .unwrap();
80
81        let result = masked.optimize().unwrap();
82        assert_eq!(masked.dtype(), result.dtype());
83        assert!(result.dtype().is_nullable());
84
85        let vector = result.execute().unwrap().into_primitive().into_u32();
86        assert_eq!(vector.elements(), &buffer![0, 1, 2, 3]);
87        assert_eq!(vector.validity().to_bit_buffer(), bitbuffer![0 1 0 1]);
88    }
89}
vortex_array/optimizer.rs

vortex_array/
optimizer.rs
