Optimise take on dictionaries with constant codes (#10643)

If the codes are constant we can just forward the value from the
dictionary as a constant

Author: robert3005

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

@@ -0,0 +1,106 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+use vortex_mask::Mask;
+
+use super::slice::inherit_selection_stats;
+use super::slice::slice_constant_code;
+use crate::ArrayRef;
+use crate::IntoArray;
+use crate::array::ArrayView;
+use crate::arrays::Constant;
+use crate::arrays::Dict;
+use crate::arrays::DictArray;
+use crate::arrays::dict::DictArraySlotsExt;
+use crate::arrays::filter::FilterReduce;
+
+impl FilterReduce for Dict {
+    fn filter(array: ArrayView<'_, Dict>, mask: &Mask) -> VortexResult<Option<ArrayRef>> {
+        if let Some(code) = array.codes().as_opt::<Constant>() {
+            return slice_constant_code(array, code.scalar(), mask.true_count());
+        }
+
+        let codes = array.codes().filter(mask.clone())?;
+        if let Some(code) = codes.as_opt::<Constant>() {
+            return slice_constant_code(array, code.scalar(), codes.len());
+        }
+        inherit_selection_stats(array.codes(), &codes);
+
+        // SAFETY: Selection preserves the integer code type and non-null code bounds, so every code
+        // indexes the unchanged values. `new_unchecked` resets `all_values_referenced` because
+        // selection can remove a value's last reference.
+        let filtered = unsafe { DictArray::new_unchecked(codes, array.values().clone()) };
+
+        Ok(Some(filtered.into_array()))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_buffer::buffer;
+    use vortex_error::VortexResult;
+    use vortex_mask::Mask;
+
+    use crate::IntoArray;
+    use crate::VortexSessionExecute;
+    use crate::array_session;
+    use crate::arrays::ConstantArray;
+    use crate::arrays::Dict;
+    use crate::arrays::DictArray;
+    use crate::arrays::PrimitiveArray;
+    use crate::arrays::dict::DictArraySlotsExt;
+    use crate::assert_arrays_eq;
+    use crate::expr::stats::Precision;
+    use crate::expr::stats::Stat;
+    use crate::expr::stats::StatsProviderExt;
+    use crate::scalar::ScalarValue;
+
+    #[test]
+    fn filter_constant_valid_code_keeps_only_referenced_value() -> VortexResult<()> {
+        let mut ctx = array_session().create_execution_ctx();
+        let dict = DictArray::new(
+            ConstantArray::new(1u8, 5).into_array(),
+            buffer![10i32, 20, 30].into_array(),
+        );
+        let filtered = dict.filter(Mask::from_iter([true, false, true, true, false]))?;
+
+        let filtered_dict = filtered.as_::<Dict>();
+        assert_arrays_eq!(
+            filtered_dict.values(),
+            buffer![20i32].into_array(),
+            &mut ctx
+        );
+        assert_arrays_eq!(
+            filtered,
+            PrimitiveArray::from_iter([20i32, 20, 20]),
+            &mut ctx
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn filter_inherits_true_order_stats_from_codes() -> VortexResult<()> {
+        let codes = buffer![0u8, 0, 1, 2].into_array();
+        codes
+            .statistics()
+            .set(Stat::IsSorted, Precision::exact(ScalarValue::Bool(true)));
+        codes
+            .statistics()
+            .set(Stat::IsConstant, Precision::exact(ScalarValue::Bool(false)));
+        let dict = DictArray::new(codes, buffer![10i32, 20, 30].into_array());
+
+        let filtered = dict.filter(Mask::from_iter([true, false, true, true]))?;
+
+        let filtered_codes = filtered.as_::<Dict>().codes().clone();
+        assert_eq!(
+            filtered_codes.statistics().get_as::<bool>(Stat::IsSorted),
+            Precision::exact(true)
+        );
+        assert_eq!(
+            filtered_codes.statistics().get_as::<bool>(Stat::IsConstant),
+            Precision::Absent
+        );
+        Ok(())
+    }
+}

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

@@ -4,6 +4,7 @@
 mod cast;
 mod compare;
 mod fill_null;
+mod filter;
 pub(crate) mod is_constant;
 pub(crate) mod is_sorted;
 mod like;
@@ -13,7 +14,6 @@ pub(crate) mod rules;
 mod slice;
 
 use vortex_error::VortexResult;
-use vortex_mask::Mask;
 
 use super::Dict;
 use super::DictArray;
@@ -23,7 +23,6 @@ use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::array::ArrayView;
 use crate::arrays::dict::DictArraySlotsExt;
-use crate::arrays::filter::FilterReduce;
 
 impl TakeExecute for Dict {
     fn take(
@@ -42,19 +41,6 @@ impl TakeExecute for Dict {
     }
 }
 
-impl FilterReduce for Dict {
-    fn filter(array: ArrayView<'_, Dict>, mask: &Mask) -> VortexResult<Option<ArrayRef>> {
-        let codes = array.codes().filter(mask.clone())?;
-
-        // SAFETY: Selection preserves the integer code type and non-null code bounds, so every code
-        // indexes the unchanged values. `new_unchecked` resets `all_values_referenced` because
-        // selection can remove a value's last reference.
-        let filtered = unsafe { DictArray::new_unchecked(codes, array.values().clone()) };
-
-        Ok(Some(filtered.into_array()))
-    }
-}
-
 #[cfg(test)]
 mod test {
     use std::sync::LazyLock;

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

@@ -31,7 +31,7 @@ impl SliceReduce for Dict {
             let sliced_code = <Primitive as SliceReduce>::slice(codes, range)?
                 .vortex_expect("Primitive SliceReduce should always return Some");
             // Because we specialize the primitive branch here, we have to make sure to handle the stat inheritance
-            inherit_slice_stats(array.codes(), &sliced_code);
+            inherit_selection_stats(array.codes(), &sliced_code);
             sliced_code
         } else {
             array.codes().slice(range)?
@@ -49,15 +49,17 @@ impl SliceReduce for Dict {
     }
 }
 
-fn inherit_slice_stats(source: &ArrayRef, sliced: &ArrayRef) {
+/// Copies the stats of `source` that hold for any order-preserving selection of its rows, such as a
+/// slice or a filter.
+pub(super) fn inherit_selection_stats(source: &ArrayRef, selected: &ArrayRef) {
     source.statistics().with_iter(|iter| {
-        sliced
+        selected
             .statistics()
-            .inherit(iter.filter(|(stat, value)| is_inheritable_true_slice_stat(*stat, value)));
+            .inherit(iter.filter(|(stat, value)| is_inheritable_selection_stat(*stat, value)));
     });
 }
 
-fn is_inheritable_true_slice_stat(stat: Stat, value: &Precision<ScalarValue>) -> bool {
+fn is_inheritable_selection_stat(stat: Stat, value: &Precision<ScalarValue>) -> bool {
     matches!(
         stat,
         Stat::IsConstant | Stat::IsSorted | Stat::IsStrictSorted
@@ -67,7 +69,7 @@ fn is_inheritable_true_slice_stat(stat: Stat, value: &Precision<ScalarValue>) ->
         .is_some_and(|value| matches!(value, ScalarValue::Bool(true)))
 }
 
-fn slice_constant_code(
+pub(super) fn slice_constant_code(
     array: ArrayView<'_, Dict>,
     code: &Scalar,
     len: usize,

vortex-array/src/arrays/dict/take.rs

@@ -11,6 +11,7 @@ use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::array::ArrayView;
 use crate::array::VTable;
+use crate::arrays::Constant;
 use crate::arrays::ConstantArray;
 use crate::arrays::dict::DictArraySlotsExt;
 use crate::expr::stats::Precision;
@@ -170,3 +171,35 @@ pub(crate) fn propagate_take_stats(
         )
     })
 }
+
+/// Executes a dictionary whose codes are constant by expanding the one value they select, rather
+/// than materializing the codes and gathering through them.
+///
+/// The executor offers the codes child to its parent's kernels before the values child, so this
+/// runs ahead of the values' take kernel, whatever their encoding.
+#[derive(Debug)]
+pub(crate) struct DictConstantCodesKernel;
+
+impl ExecuteParentKernel<Constant> for DictConstantCodesKernel {
+    type Parent = Dict;
+
+    fn execute_parent(
+        &self,
+        array: ArrayView<'_, Constant>,
+        parent: <Self::Parent as Matcher>::Match<'_>,
+        child_idx: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        if child_idx != 0 || parent.is_empty() {
+            return Ok(None);
+        }
+        let value = match array.scalar().as_primitive().as_::<usize>() {
+            Some(code) => parent
+                .values()
+                .execute_scalar(code, ctx)?
+                .cast(parent.dtype())?,
+            None => Scalar::null(parent.dtype().as_nullable()),
+        };
+        Ok(Some(ConstantArray::new(value, parent.len()).into_array()))
+    }
+}

vortex-array/src/arrays/dict/tests.rs

@@ -2,18 +2,24 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_buffer::buffer;
+use vortex_error::VortexResult;
 
 use super::DictArray;
 use crate::IntoArray;
 use crate::VortexSessionExecute;
 use crate::array_session;
 use crate::arrays::BoolArray;
+use crate::arrays::ConstantArray;
 use crate::arrays::FixedSizeListArray;
 use crate::arrays::ListArray;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::VarBinArray;
 use crate::assert_arrays_eq;
 use crate::builders::builder_with_capacity_in;
+use crate::dtype::DType;
+use crate::dtype::Nullability;
+use crate::dtype::PType;
+use crate::scalar::Scalar;
 use crate::validity::Validity;
 
 #[test]
@@ -30,7 +36,7 @@ fn test_scalar_at_null_code() {
 }
 
 #[test]
-fn test_canonicalize_dict_with_all_null_array_validity_codes() -> vortex_error::VortexResult<()> {
+fn test_canonicalize_dict_with_all_null_array_validity_codes() -> VortexResult<()> {
     let mut ctx = array_session().create_execution_ctx();
     let codes = PrimitiveArray::new(
         buffer![0u32, 1],
@@ -133,3 +139,42 @@ fn test_dict_list_dict_display() {
         "[[\"Bonjour\", \"Bonjour\", null], [\"你好\", \"Hello\"], null, [\"Hello\", \"Hello\"], null, []]"
     )
 }
+
+#[test]
+fn canonicalize_constant_nullable_code_over_non_nullable_values() -> VortexResult<()> {
+    let mut ctx = array_session().create_execution_ctx();
+    let code = Scalar::primitive(1u8, Nullability::Nullable);
+    let dict = DictArray::try_new(
+        ConstantArray::new(code, 3).into_array(),
+        VarBinArray::from_iter(
+            ["a", "long enough to live out of line"].map(Some),
+            DType::Utf8(Nullability::NonNullable),
+        )
+        .into_array(),
+    )?;
+    let expected = VarBinArray::from_iter(
+        [Some("long enough to live out of line"); 3],
+        DType::Utf8(Nullability::Nullable),
+    );
+    assert_arrays_eq!(dict, expected, &mut ctx);
+    Ok(())
+}
+
+#[test]
+fn canonicalize_constant_null_code() -> VortexResult<()> {
+    let mut ctx = array_session().create_execution_ctx();
+    let dict = DictArray::try_new(
+        ConstantArray::new(
+            Scalar::null(DType::Primitive(PType::U8, Nullability::Nullable)),
+            3,
+        )
+        .into_array(),
+        buffer![10i32, 20, 30].into_array(),
+    )?;
+    assert_arrays_eq!(
+        dict,
+        PrimitiveArray::from_option_iter([None::<i32>, None, None]),
+        &mut ctx
+    );
+    Ok(())
+}

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

@@ -5,8 +5,10 @@ use vortex_session::VortexSession;
 
 use crate::ArrayVTable;
 use crate::arrays::Chunked;
+use crate::arrays::Constant;
 use crate::arrays::Dict;
 use crate::arrays::dict::TakeExecuteAdaptor;
+use crate::arrays::dict::take::DictConstantCodesKernel;
 use crate::optimizer::kernels::ArrayKernelsExt;
 use crate::scalar_fn::ScalarFnVTable;
 use crate::scalar_fn::fns::binary::Binary;
@@ -18,6 +20,7 @@ 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(), Constant, DictConstantCodesKernel);
     kernels.register_execute_parent_kernel(Dict.id(), Dict, TakeExecuteAdaptor(Dict));
     kernels.register_execute_parent_kernel(FillNull.id(), Dict, FillNullExecuteAdaptor(Dict));
 }