|
| 1 | +// SPDX-License-Identifier: Apache-2.0 |
| 2 | +// SPDX-FileCopyrightText: Copyright the Vortex contributors |
| 3 | + |
| 4 | +use std::sync::Arc; |
| 5 | + |
| 6 | +use vortex_error::VortexResult; |
| 7 | + |
| 8 | +use crate::vtable::VTable; |
| 9 | +use crate::{Array, ArrayRef}; |
| 10 | + |
| 11 | +impl 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. |
| 44 | +pub 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. |
| 53 | +pub 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)] |
| 61 | +mod tests { |
| 62 | + use vortex_buffer::{bitbuffer, buffer}; |
| 63 | + use vortex_dtype::PTypeDowncast; |
| 64 | + use vortex_vector::VectorOps; |
| 65 | + |
| 66 | + use crate::arrays::{BoolArray, MaskedArray, PrimitiveArray}; |
| 67 | + use crate::validity::Validity; |
| 68 | + use crate::{ArrayOperator, IntoArray}; |
| 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 | +} |
0 commit comments