More operator rules

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

encodings/fastlanes/src/bitpacking/vtable/kernels/filter.rs

@@ -22,9 +22,8 @@ use vortex_mask::Mask;
 use vortex_mask::MaskValues;
 use vortex_vector::Vector;
 use vortex_vector::VectorMutOps;
-use vortex_vector::primitive::PVector;
 use vortex_vector::primitive::PVectorMut;
-use vortex_vector::primitive::PrimitiveVector;
+use vortex_vector::primitive::PrimitiveVectorMut;
 
 use crate::BitPackedArray;
 use crate::BitPackedVTable;
@@ -83,36 +82,46 @@ impl ExecuteParentKernel<BitPackedVTable> for BitPackingFilterKernel {
             }
         });
 
-        let primitive_vector: PrimitiveVector = match array.ptype() {
-            PType::U8 => filter_primitive::<u8>(array, values)?.into(),
-            PType::U16 => filter_primitive::<u16>(array, values)?.into(),
-            PType::U32 => filter_primitive::<u32>(array, values)?.into(),
-            PType::U64 => filter_primitive::<u64>(array, values)?.into(),
+        let mut primitive_vector: PrimitiveVectorMut = match array.ptype() {
+            PType::U8 => filter_primitive_without_patches::<u8>(array, values)?.into(),
+            PType::U16 => filter_primitive_without_patches::<u16>(array, values)?.into(),
+            PType::U32 => filter_primitive_without_patches::<u32>(array, values)?.into(),
+            PType::U64 => filter_primitive_without_patches::<u64>(array, values)?.into(),
 
             // Since the fastlanes crate only supports unsigned integers, and since we know that all
             // numbers are going to be non-negative, we can safely "cast" to unsigned and back.
             PType::I8 => {
-                let pvector = filter_primitive::<u8>(array, values)?;
+                let pvector = filter_primitive_without_patches::<u8>(array, values)?;
                 unsafe { pvector.transmute::<i8>() }.into()
             }
             PType::I16 => {
-                let pvector = filter_primitive::<u16>(array, values)?;
+                let pvector = filter_primitive_without_patches::<u16>(array, values)?;
                 unsafe { pvector.transmute::<i16>() }.into()
             }
             PType::I32 => {
-                let pvector = filter_primitive::<u32>(array, values)?;
+                let pvector = filter_primitive_without_patches::<u32>(array, values)?;
                 unsafe { pvector.transmute::<i32>() }.into()
             }
             PType::I64 => {
-                let pvector = filter_primitive::<u64>(array, values)?;
+                let pvector = filter_primitive_without_patches::<u64>(array, values)?;
                 unsafe { pvector.transmute::<i64>() }.into()
             }
             other => {
                 unreachable!("Unsupported ptype {other} for bitpacking, we also checked this above")
             }
         };
 
-        Ok(Some(primitive_vector.into()))
+        // TODO(connor): We want a `PatchesArray` or patching compute functions instead of this.
+        let patches = array
+            .patches()
+            .map(|patches| patches.filter(&Mask::Values(values.clone())))
+            .transpose()?
+            .flatten();
+        if let Some(patches) = patches {
+            primitive_vector = patches.apply_to_primitive_vector(primitive_vector);
+        }
+
+        Ok(Some(primitive_vector.freeze().into()))
     }
 }
 
@@ -125,10 +134,10 @@ impl ExecuteParentKernel<BitPackedVTable> for BitPackingFilterKernel {
 /// This function fully decompresses the array for all but the most selective masks because the
 /// FastLanes decompression is so fast and the bookkeepping necessary to decompress individual
 /// elements is relatively slow.
-fn filter_primitive<U: UnsignedPType + BitPacking>(
+fn filter_primitive_without_patches<U: UnsignedPType + BitPacking>(
     array: &BitPackedArray,
     selection: &Arc<MaskValues>,
-) -> VortexResult<PVector<U>> {
+) -> VortexResult<PVectorMut<U>> {
     let values = filter_with_indices(array, selection.indices());
     let validity = array
         .validity_mask()
@@ -141,19 +150,7 @@ fn filter_primitive<U: UnsignedPType + BitPacking>(
         "`filter_with_indices` was somehow incorrect"
     );
 
-    let mut pvector = unsafe { PVectorMut::new_unchecked(values, validity) };
-
-    // TODO(connor): We want a `PatchesArray` or patching compute functions instead of this.
-    let patches = array
-        .patches()
-        .map(|patches| patches.filter(&Mask::Values(selection.clone())))
-        .transpose()?
-        .flatten();
-    if let Some(patches) = patches {
-        pvector = patches.apply_to_pvector(pvector);
-    }
-
-    Ok(pvector.freeze())
+    Ok(unsafe { PVectorMut::new_unchecked(values, validity) })
 }
 
 fn filter_with_indices<T: NativePType + BitPacking>(

vortex-array/src/arrays/dict/vtable/rules.rs

@@ -14,6 +14,7 @@ use crate::arrays::ConstantVTable;
 use crate::arrays::DictArray;
 use crate::arrays::DictVTable;
 use crate::arrays::ScalarFnArray;
+use crate::builtins::ArrayBuiltins;
 use crate::optimizer::ArrayOptimizer;
 use crate::optimizer::rules::ArrayParentReduceRule;
 use crate::optimizer::rules::ParentRuleSet;
@@ -94,10 +95,19 @@ impl ArrayParentReduceRule<DictVTable> for DictionaryScalarFnValuesPushDownRule
                 .into_array()
                 .optimize()?;
 
-        let new_dict =
-            unsafe { DictArray::new_unchecked(array.codes().clone(), new_values) }.into_array();
+        // We can only push down null-sensitive functions when we have all-valid codes.
+        // In these cases, we cannot have the codes influence the nullability of the output DType.
+        // Therefore, we cast the codes to be non-nullable and then cast the dictionary output
+        // back to nullable if needed.
+        if sig.is_null_sensitive() && array.codes().dtype().is_nullable() {
+            let new_codes = array.codes().cast(array.codes().dtype().as_nonnullable())?;
+            let new_dict = unsafe { DictArray::new_unchecked(new_codes, new_values) }.into_array();
+            return Ok(Some(new_dict.cast(parent.dtype().clone())?));
+        }
 
-        Ok(Some(new_dict))
+        Ok(Some(
+            unsafe { DictArray::new_unchecked(array.codes().clone(), new_values) }.into_array(),
+        ))
     }
 }
 
@@ -117,13 +127,27 @@ impl ArrayParentReduceRule<DictVTable> for DictionaryScalarFnCodesPullUpRule {
         parent: &ScalarFnArray,
         child_idx: usize,
     ) -> VortexResult<Option<ArrayRef>> {
-        // Check that all siblings are dictionaries with the same codes as us.
+        // Don't attempt to pull up if there are less than 2 siblings.
+        if parent.children().len() < 2 {
+            return Ok(None);
+        }
+
+        // Check that all siblings are dictionaries, and have the same number of values as us.
+        // This is a cheap first loop.
+        if !parent.children().iter().enumerate().all(|(idx, c)| {
+            idx == child_idx
+                || c.as_opt::<DictVTable>()
+                    .is_some_and(|c| c.values().len() == array.values().len())
+        }) {
+            return Ok(None);
+        }
+
+        // Now run the slightly more expensive check that all siblings have the same codes as us.
+        // We use the cheaper Precision::Ptr to avoid doing data comparisons.
         if !parent.children().iter().enumerate().all(|(idx, c)| {
             idx == child_idx
-                || c.as_opt::<DictVTable>().is_some_and(|c| {
-                    c.values().len() == array.values().len()
-                        && c.codes().array_eq(array.codes(), Precision::Value)
-                })
+                || c.as_opt::<DictVTable>()
+                    .is_some_and(|c| c.codes().array_eq(array.codes(), Precision::Value))
         }) {
             return Ok(None);
         }

vortex-array/src/optimizer/rules.rs

@@ -153,6 +153,25 @@ impl<V: VTable> ParentRuleSet<V> {
                 continue;
             }
             if let Some(reduced) = rule.reduce_parent(child, parent, child_idx)? {
+                // Debug assertions because these checks are already run elsewhere.
+                #[cfg(debug_assertions)]
+                {
+                    vortex_error::vortex_ensure!(
+                        reduced.len() == parent.len(),
+                        "Reduced array length mismatch from {:?}\nFrom:\n{}\nTo:\n{}",
+                        rule,
+                        parent.display_tree(),
+                        reduced.display_tree()
+                    );
+                    vortex_error::vortex_ensure!(
+                        reduced.dtype() == parent.dtype(),
+                        "Reduced array dtype mismatch from {:?}\nFrom:\n{}\nTo:\n{}",
+                        rule,
+                        parent.display_tree(),
+                        reduced.display_tree()
+                    );
+                }
+
                 return Ok(Some(reduced));
             }
         }

vortex-array/src/patches.rs

@@ -15,8 +15,10 @@ use vortex_dtype::IntegerPType;
 use vortex_dtype::NativePType;
 use vortex_dtype::Nullability::NonNullable;
 use vortex_dtype::PType;
+use vortex_dtype::PTypeDowncastExt;
 use vortex_dtype::UnsignedPType;
 use vortex_dtype::match_each_integer_ptype;
+use vortex_dtype::match_each_native_ptype;
 use vortex_dtype::match_each_unsigned_integer_ptype;
 use vortex_error::VortexError;
 use vortex_error::VortexExpect;
@@ -31,6 +33,7 @@ use vortex_scalar::PValue;
 use vortex_scalar::Scalar;
 use vortex_utils::aliases::hash_map::HashMap;
 use vortex_vector::primitive::PVectorMut;
+use vortex_vector::primitive::PrimitiveVectorMut;
 
 use crate::Array;
 use crate::ArrayRef;
@@ -823,6 +826,13 @@ impl Patches {
         }))
     }
 
+    /// Applies patches to a primitive vector, returning the patched vector.
+    pub fn apply_to_primitive_vector(&self, vector: PrimitiveVectorMut) -> PrimitiveVectorMut {
+        match_each_native_ptype!(vector.ptype(), |T| {
+            self.apply_to_pvector(vector.downcast::<T>()).into()
+        })
+    }
+
     /// Applies patches to a [`PVectorMut<T>`], returning the patched vector.
     ///
     /// This function modifies the elements buffer in-place at the positions specified by the patch

vortex-array/src/vtable/dyn_.rs

@@ -160,11 +160,15 @@ impl<V: VTable> DynVTable for ArrayVTableAdapter<V> {
         };
         vortex_ensure!(
             reduced.len() == array.len(),
-            "Reduced array length mismatch"
+            "Reduced array length mismatch from {} to {}",
+            array.display_tree(),
+            reduced.display_tree()
         );
         vortex_ensure!(
             reduced.dtype() == array.dtype(),
-            "Reduced array dtype mismatch"
+            "Reduced array dtype mismatch from {} to {}",
+            array.display_tree(),
+            reduced.display_tree()
         );
         Ok(Some(reduced))
     }
@@ -181,11 +185,15 @@ impl<V: VTable> DynVTable for ArrayVTableAdapter<V> {
 
         vortex_ensure!(
             reduced.len() == parent.len(),
-            "Reduced array length mismatch"
+            "Reduced array length mismatch from {} to {}",
+            parent.display_tree(),
+            reduced.display_tree()
         );
         vortex_ensure!(
             reduced.dtype() == parent.dtype(),
-            "Reduced array dtype mismatch"
+            "Reduced array dtype mismatch from {} to {}",
+            parent.display_tree(),
+            reduced.display_tree()
         );
 
         Ok(Some(reduced))

vortex-buffer/src/buffer_mut.rs

@@ -444,6 +444,33 @@ impl<T> BufferMut<T> {
             Self::copy_from_aligned(self, alignment)
         }
     }
+
+    /// Transmute a `Buffer<T>` into a `Buffer<U>`.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that all possible bit representations of type `T` are valid when
+    /// interpreted as type `U`.
+    /// See [`std::mem::transmute`] for more details.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the type `U` does not have the same size and alignment as `T`.
+    pub unsafe fn transmute<U>(self) -> BufferMut<U> {
+        assert_eq!(size_of::<T>(), size_of::<U>(), "Buffer type size mismatch");
+        assert_eq!(
+            align_of::<T>(),
+            align_of::<U>(),
+            "Buffer type alignment mismatch"
+        );
+
+        BufferMut {
+            bytes: self.bytes,
+            length: self.length,
+            alignment: self.alignment,
+            _marker: std::marker::PhantomData,
+        }
+    }
 }
 
 impl<T> Clone for BufferMut<T> {

vortex-vector/src/primitive/generic_mut.rs

@@ -125,6 +125,25 @@ impl<T> PVectorMut<T> {
         self.elements.push_n(value, n);
         self.validity.append_n(true, n);
     }
+
+    /// Transmute a `PVectorMut<T>` into a `PVectorMut<U>`.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that all values of type `T` in this vector are valid as type `U`.
+    /// See [`std::mem::transmute`] for more details.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the type `U` does not have the same size and alignment as `T`.
+    pub unsafe fn transmute<U: NativePType>(self) -> PVectorMut<U> {
+        let (buffer, mask) = self.into_parts();
+
+        // SAFETY: same guarantees as this function.
+        let buffer = unsafe { buffer.transmute::<U>() };
+
+        PVectorMut::new(buffer, mask)
+    }
 }
 
 impl<T: NativePType> VectorMutOps for PVectorMut<T> {