Reveal Arrow exporter fast-path encodings before conversion

vortex-array/src/arrays/dict/compute/mask.rs

@@ -4,14 +4,17 @@
 use vortex_error::VortexResult;
 
 use crate::ArrayRef;
+use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::array::ArrayView;
 use crate::arrays::Dict;
 use crate::arrays::DictArray;
+use crate::arrays::Primitive;
 use crate::arrays::dict::DictArraySlotsExt;
 use crate::arrays::scalar_fn::ScalarFnFactoryExt;
 use crate::scalar_fn::EmptyOptions;
 use crate::scalar_fn::fns::mask::Mask as MaskExpr;
+use crate::scalar_fn::fns::mask::MaskKernel;
 use crate::scalar_fn::fns::mask::MaskReduce;
 
 impl MaskReduce for Dict {
@@ -27,3 +30,65 @@ impl MaskReduce for Dict {
         }))
     }
 }
+
+impl MaskKernel for Dict {
+    fn mask(
+        array: ArrayView<'_, Dict>,
+        mask: &ArrayRef,
+        _ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        // Keep Primitive codes probe-able: AND the mask into their validity instead of hiding
+        // them behind a scalar-fn wrapper, which would defeat downstream fast paths that
+        // downcast the codes (constant-code slicing, specialized primitive slicing, ...).
+        if let Some(codes) = array.codes().as_typed::<Primitive>()
+            && let Some(masked_codes) = <Primitive as MaskReduce>::mask(codes, mask)?
+        {
+            // SAFETY: masking codes doesn't change dict invariants
+            return Ok(Some(unsafe {
+                DictArray::new_unchecked(masked_codes, array.values().clone()).into_array()
+            }));
+        }
+
+        <Dict as MaskReduce>::mask(array, mask)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_buffer::buffer;
+    use vortex_error::VortexResult;
+    use vortex_error::vortex_err;
+
+    use crate::IntoArray;
+    use crate::VortexSessionExecute;
+    use crate::array_session;
+    use crate::arrays::BoolArray;
+    use crate::arrays::Dict;
+    use crate::arrays::DictArray;
+    use crate::arrays::Primitive;
+    use crate::arrays::VarBinViewArray;
+    use crate::arrays::dict::DictArraySlotsExt;
+    use crate::arrays::scalar_fn::ScalarFnFactoryExt;
+    use crate::scalar_fn::EmptyOptions;
+    use crate::scalar_fn::fns::mask::Mask;
+
+    #[test]
+    fn mask_kernel_keeps_primitive_codes_probeable() -> VortexResult<()> {
+        let mut ctx = array_session().create_execution_ctx();
+        let dict = DictArray::try_new(
+            buffer![0u8, 1, 0].into_array(),
+            VarBinViewArray::from_iter_str(["a", "b"]).into_array(),
+        )?;
+        let mask = BoolArray::from_iter([true, false, true]);
+        let masked = Mask.try_new_array(3, EmptyOptions, [dict.into_array(), mask.into_array()])?;
+
+        // The mask execute-parent kernel reveals a Dict whose codes remain Primitive (with
+        // masked validity) rather than being hidden behind a scalar-fn wrapper.
+        let revealed = masked.execute_until::<Dict>(&mut ctx)?;
+        let dict = revealed
+            .try_downcast::<Dict>()
+            .map_err(|a| vortex_err!("expected Dict, got {}", a.encoding_id()))?;
+        assert!(dict.codes().is::<Primitive>());
+        Ok(())
+    }
+}

vortex-array/src/arrays/dict/compute/slice.rs

@@ -7,6 +7,7 @@ use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 
 use crate::ArrayRef;
+use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::array::ArrayView;
 use crate::arrays::Constant;
@@ -15,12 +16,23 @@ use crate::arrays::Dict;
 use crate::arrays::DictArray;
 use crate::arrays::Primitive;
 use crate::arrays::dict::DictArraySlotsExt;
+use crate::arrays::slice::SliceKernel;
 use crate::arrays::slice::SliceReduce;
 use crate::expr::stats::Precision;
 use crate::expr::stats::Stat;
 use crate::scalar::Scalar;
 use crate::scalar::ScalarValue;
 
+impl SliceKernel for Dict {
+    fn slice(
+        array: ArrayView<'_, Self>,
+        range: Range<usize>,
+        _ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        <Dict as SliceReduce>::slice(array, range)
+    }
+}
+
 impl SliceReduce for Dict {
     fn slice(array: ArrayView<'_, Self>, range: Range<usize>) -> VortexResult<Option<ArrayRef>> {
         if let Some(code) = array.codes().as_opt::<Constant>() {

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

@@ -6,18 +6,24 @@ use vortex_session::VortexSession;
 use crate::ArrayVTable;
 use crate::arrays::Chunked;
 use crate::arrays::Dict;
+use crate::arrays::Slice;
 use crate::arrays::dict::TakeExecuteAdaptor;
+use crate::arrays::slice::SliceExecuteAdaptor;
 use crate::optimizer::kernels::ArrayKernelsExt;
 use crate::scalar_fn::ScalarFnVTable;
 use crate::scalar_fn::fns::binary::Binary;
 use crate::scalar_fn::fns::binary::CompareExecuteAdaptor;
 use crate::scalar_fn::fns::fill_null::FillNull;
 use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor;
+use crate::scalar_fn::fns::mask::Mask;
+use crate::scalar_fn::fns::mask::MaskExecuteAdaptor;
 
 pub(crate) fn initialize(session: &VortexSession) {
     let kernels = session.kernels();
     kernels.register_execute_parent_kernel(Binary.id(), Dict, CompareExecuteAdaptor(Dict));
     kernels.register_execute_parent_kernel(Dict.id(), Chunked, TakeExecuteAdaptor(Chunked));
     kernels.register_execute_parent_kernel(Dict.id(), Dict, TakeExecuteAdaptor(Dict));
     kernels.register_execute_parent_kernel(FillNull.id(), Dict, FillNullExecuteAdaptor(Dict));
+    kernels.register_execute_parent_kernel(Mask.id(), Dict, MaskExecuteAdaptor(Dict));
+    kernels.register_execute_parent_kernel(Slice.id(), Dict, SliceExecuteAdaptor(Dict));
 }

vortex-array/src/arrays/list/compute/kernels.rs

@@ -7,16 +7,22 @@ use crate::ArrayVTable;
 use crate::arrays::Dict;
 use crate::arrays::Filter;
 use crate::arrays::List;
+use crate::arrays::Slice;
 use crate::arrays::dict::TakeExecuteAdaptor;
 use crate::arrays::filter::FilterExecuteAdaptor;
+use crate::arrays::slice::SliceExecuteAdaptor;
 use crate::optimizer::kernels::ArrayKernelsExt;
 use crate::scalar_fn::ScalarFnVTable;
 use crate::scalar_fn::fns::cast::Cast;
 use crate::scalar_fn::fns::cast::CastExecuteAdaptor;
+use crate::scalar_fn::fns::mask::Mask;
+use crate::scalar_fn::fns::mask::MaskExecuteAdaptor;
 
 pub(crate) fn initialize(session: &VortexSession) {
     let kernels = session.kernels();
     kernels.register_execute_parent_kernel(Cast.id(), List, CastExecuteAdaptor(List));
     kernels.register_execute_parent_kernel(Filter.id(), List, FilterExecuteAdaptor(List));
     kernels.register_execute_parent_kernel(Dict.id(), List, TakeExecuteAdaptor(List));
+    kernels.register_execute_parent_kernel(Mask.id(), List, MaskExecuteAdaptor(List));
+    kernels.register_execute_parent_kernel(Slice.id(), List, SliceExecuteAdaptor(List));
 }

vortex-array/src/arrays/list/compute/mask.rs

@@ -4,11 +4,13 @@
 use vortex_error::VortexResult;
 
 use crate::ArrayRef;
+use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::array::ArrayView;
 use crate::arrays::List;
 use crate::arrays::ListArray;
 use crate::arrays::list::ListArraySlotsExt;
+use crate::scalar_fn::fns::mask::MaskKernel;
 use crate::scalar_fn::fns::mask::MaskReduce;
 use crate::validity::Validity;
 
@@ -22,3 +24,13 @@ impl MaskReduce for List {
         .map(|a| Some(a.into_array()))
     }
 }
+
+impl MaskKernel for List {
+    fn mask(
+        array: ArrayView<'_, List>,
+        mask: &ArrayRef,
+        _ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        <List as MaskReduce>::mask(array, mask)
+    }
+}

vortex-array/src/arrays/list/compute/slice.rs

@@ -6,11 +6,13 @@ use std::ops::Range;
 use vortex_error::VortexResult;
 
 use crate::ArrayRef;
+use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::array::ArrayView;
 use crate::arrays::List;
 use crate::arrays::ListArray;
 use crate::arrays::list::ListArraySlotsExt;
+use crate::arrays::slice::SliceKernel;
 use crate::arrays::slice::SliceReduce;
 
 impl SliceReduce for List {
@@ -25,3 +27,13 @@ impl SliceReduce for List {
         ))
     }
 }
+
+impl SliceKernel for List {
+    fn slice(
+        array: ArrayView<'_, Self>,
+        range: Range<usize>,
+        _ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        <List as SliceReduce>::slice(array, range)
+    }
+}

vortex-array/src/arrays/varbin/compute/mask.rs

@@ -4,11 +4,13 @@
 use vortex_error::VortexResult;
 
 use crate::ArrayRef;
+use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::array::ArrayView;
 use crate::arrays::VarBin;
 use crate::arrays::VarBinArray;
 use crate::arrays::varbin::VarBinArraySlotsExt;
+use crate::scalar_fn::fns::mask::MaskKernel;
 use crate::scalar_fn::fns::mask::MaskReduce;
 use crate::validity::Validity;
 
@@ -26,6 +28,16 @@ impl MaskReduce for VarBin {
     }
 }
 
+impl MaskKernel for VarBin {
+    fn mask(
+        array: ArrayView<'_, VarBin>,
+        mask: &ArrayRef,
+        _ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        <VarBin as MaskReduce>::mask(array, mask)
+    }
+}
+
 #[cfg(test)]
 mod test {
     use crate::IntoArray;

vortex-array/src/arrays/varbin/compute/slice.rs

@@ -6,10 +6,12 @@ use std::ops::Range;
 use vortex_error::VortexResult;
 
 use crate::ArrayRef;
+use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::array::ArrayView;
 use crate::arrays::VarBin;
 use crate::arrays::VarBinArray;
+use crate::arrays::slice::SliceKernel;
 use crate::arrays::slice::SliceReduce;
 use crate::arrays::varbin::VarBinArraySlotsExt;
 
@@ -19,6 +21,16 @@ impl SliceReduce for VarBin {
     }
 }
 
+impl SliceKernel for VarBin {
+    fn slice(
+        array: ArrayView<'_, Self>,
+        range: Range<usize>,
+        _ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        <VarBin as SliceReduce>::slice(array, range)
+    }
+}
+
 impl VarBin {
     pub fn _slice(array: ArrayView<'_, VarBin>, range: Range<usize>) -> VortexResult<ArrayRef> {
         Ok(unsafe {

vortex-array/src/arrays/varbin/vtable/kernel.rs

@@ -6,20 +6,26 @@ use vortex_session::VortexSession;
 use crate::ArrayVTable;
 use crate::arrays::Dict;
 use crate::arrays::Filter;
+use crate::arrays::Slice;
 use crate::arrays::VarBin;
 use crate::arrays::dict::TakeExecuteAdaptor;
 use crate::arrays::filter::FilterExecuteAdaptor;
+use crate::arrays::slice::SliceExecuteAdaptor;
 use crate::optimizer::kernels::ArrayKernelsExt;
 use crate::scalar_fn::ScalarFnVTable;
 use crate::scalar_fn::fns::binary::Binary;
 use crate::scalar_fn::fns::binary::CompareExecuteAdaptor;
 use crate::scalar_fn::fns::cast::Cast;
 use crate::scalar_fn::fns::cast::CastExecuteAdaptor;
+use crate::scalar_fn::fns::mask::Mask;
+use crate::scalar_fn::fns::mask::MaskExecuteAdaptor;
 
 pub(crate) fn initialize(session: &VortexSession) {
     let kernels = session.kernels();
     kernels.register_execute_parent_kernel(Cast.id(), VarBin, CastExecuteAdaptor(VarBin));
     kernels.register_execute_parent_kernel(Binary.id(), VarBin, CompareExecuteAdaptor(VarBin));
     kernels.register_execute_parent_kernel(Filter.id(), VarBin, FilterExecuteAdaptor(VarBin));
     kernels.register_execute_parent_kernel(Dict.id(), VarBin, TakeExecuteAdaptor(VarBin));
+    kernels.register_execute_parent_kernel(Mask.id(), VarBin, MaskExecuteAdaptor(VarBin));
+    kernels.register_execute_parent_kernel(Slice.id(), VarBin, SliceExecuteAdaptor(VarBin));
 }