Dict: gather dictionary values straight into a VarBinBuilder

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

@@ -5,11 +5,12 @@
 
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
 
+use crate::ArrayView;
 use crate::Canonical;
 use crate::CanonicalView;
 use crate::ExecutionCtx;
-use crate::IntoArray;
 use crate::arrays::Bool;
 use crate::arrays::BoolArray;
 use crate::arrays::Decimal;
@@ -20,6 +21,7 @@ use crate::arrays::FixedSizeList;
 use crate::arrays::FixedSizeListArray;
 use crate::arrays::ListView;
 use crate::arrays::ListViewArray;
+use crate::arrays::Null;
 use crate::arrays::NullArray;
 use crate::arrays::Primitive;
 use crate::arrays::PrimitiveArray;
@@ -38,148 +40,138 @@ use crate::arrays::variant::VariantArraySlotsExt;
 /// by looking up each code in the values array.
 pub(crate) fn take_canonical(
     values: CanonicalView,
-    codes: &PrimitiveArray,
+    codes: ArrayView<'_, Primitive>,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<Canonical> {
-    let values = Canonical::from(values);
     Ok(match values {
-        Canonical::Null(a) => Canonical::Null(take_null(&a, codes)),
-        Canonical::Bool(a) => Canonical::Bool(take_bool(&a, codes, ctx)?),
-        Canonical::Primitive(a) => Canonical::Primitive(take_primitive(&a, codes, ctx)),
-        Canonical::Decimal(a) => Canonical::Decimal(take_decimal(&a, codes, ctx)),
-        Canonical::VarBinView(a) => Canonical::VarBinView(take_varbinview(&a, codes, ctx)),
-        Canonical::List(a) => Canonical::List(take_listview(&a, codes, ctx)),
-        Canonical::Map(_) => vortex_error::vortex_bail!("Map arrays don't support take"),
-        Canonical::FixedSizeList(a) => {
-            Canonical::FixedSizeList(take_fixed_size_list(&a, codes, ctx))
+        CanonicalView::Null(a) => Canonical::Null(take_null(a, codes)),
+        CanonicalView::Bool(a) => Canonical::Bool(take_bool(a, codes, ctx)?),
+        CanonicalView::Primitive(a) => Canonical::Primitive(take_primitive(a, codes, ctx)),
+        CanonicalView::Decimal(a) => Canonical::Decimal(take_decimal(a, codes, ctx)),
+        CanonicalView::VarBinView(a) => Canonical::VarBinView(take_varbinview(a, codes, ctx)),
+        CanonicalView::List(a) => Canonical::List(take_listview(a, codes, ctx)),
+        CanonicalView::Map(_) => vortex_bail!("Map arrays don't support take"),
+        CanonicalView::FixedSizeList(a) => {
+            Canonical::FixedSizeList(take_fixed_size_list(a, codes, ctx))
         }
-        Canonical::Struct(a) => Canonical::Struct(take_struct(&a, codes)),
-        Canonical::Union(_) => {
+        CanonicalView::Struct(a) => Canonical::Struct(take_struct(a, codes)),
+        CanonicalView::Union(_) => {
             todo!(
                 "TODO(connor)[Union]: implement dictionary execution after Union take supports \
                  nullable indices and outer null propagation"
             )
         }
-        Canonical::Extension(a) => Canonical::Extension(take_extension(&a, codes, ctx)),
-        Canonical::Variant(a) => {
-            let indices = codes.clone().into_array();
+        CanonicalView::Extension(a) => Canonical::Extension(take_extension(a, codes, ctx)),
+        CanonicalView::Variant(a) => {
+            let indices = codes.array().clone();
             let taken_core_storage = a.core_storage().take(indices.clone())?;
             let taken_shredded = a
                 .shredded()
-                .map(|shredded| shredded.take(indices.clone()))
+                .map(|shredded| shredded.take(indices))
                 .transpose()?;
             Canonical::Variant(VariantArray::try_new(taken_core_storage, taken_shredded)?)
         }
     })
 }
 
 /// Take for NullArray is trivial - just create a new NullArray with the new length.
-fn take_null(_array: &NullArray, codes: &PrimitiveArray) -> NullArray {
+fn take_null(_array: ArrayView<'_, Null>, codes: ArrayView<'_, Primitive>) -> NullArray {
     NullArray::new(codes.len())
 }
 
-// TODO(joe): use dict_bool_take
 fn take_bool(
-    array: &BoolArray,
-    codes: &PrimitiveArray,
+    array: ArrayView<'_, Bool>,
+    codes: ArrayView<'_, Primitive>,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<BoolArray> {
-    let codes_ref = codes.clone().into_array();
-    let array = array.as_view();
-    Ok(<Bool as TakeExecute>::take(array, &codes_ref, ctx)?
+    let codes_ref = codes.array();
+    Ok(<Bool as TakeExecute>::take(array, codes_ref, ctx)?
         .vortex_expect("take bool should not return None")
         .as_::<Bool>()
         .into_owned())
 }
 
 fn take_primitive(
-    array: &PrimitiveArray,
-    codes: &PrimitiveArray,
+    array: ArrayView<'_, Primitive>,
+    codes: ArrayView<'_, Primitive>,
     ctx: &mut ExecutionCtx,
 ) -> PrimitiveArray {
-    let codes_ref = codes.clone().into_array();
-    let array = array.as_view();
-    <Primitive as TakeExecute>::take(array, &codes_ref, ctx)
+    let codes_ref = codes.array();
+    <Primitive as TakeExecute>::take(array, codes_ref, ctx)
         .vortex_expect("take primitive array")
         .vortex_expect("take primitive should not return None")
         .as_::<Primitive>()
         .into_owned()
 }
 
 fn take_decimal(
-    array: &DecimalArray,
-    codes: &PrimitiveArray,
+    array: ArrayView<'_, Decimal>,
+    codes: ArrayView<'_, Primitive>,
     ctx: &mut ExecutionCtx,
 ) -> DecimalArray {
-    let codes_ref = codes.clone().into_array();
-    let array = array.as_view();
-    <Decimal as TakeExecute>::take(array, &codes_ref, ctx)
+    let codes_ref = codes.array();
+    <Decimal as TakeExecute>::take(array, codes_ref, ctx)
         .vortex_expect("take decimal array")
         .vortex_expect("take decimal should not return None")
         .as_::<Decimal>()
         .into_owned()
 }
 
 fn take_varbinview(
-    array: &VarBinViewArray,
-    codes: &PrimitiveArray,
+    array: ArrayView<'_, VarBinView>,
+    codes: ArrayView<'_, Primitive>,
     ctx: &mut ExecutionCtx,
 ) -> VarBinViewArray {
-    let codes_ref = codes.clone().into_array();
-    let array = array.as_view();
-    <VarBinView as TakeExecute>::take(array, &codes_ref, ctx)
+    let codes_ref = codes.array();
+    <VarBinView as TakeExecute>::take(array, codes_ref, ctx)
         .vortex_expect("take varbinview array")
         .vortex_expect("take varbinview should not return None")
         .as_::<VarBinView>()
         .into_owned()
 }
 
 fn take_listview(
-    array: &ListViewArray,
-    codes: &PrimitiveArray,
+    array: ArrayView<'_, ListView>,
+    codes: ArrayView<'_, Primitive>,
     ctx: &mut ExecutionCtx,
 ) -> ListViewArray {
-    let codes_ref = codes.clone().into_array();
-    let array = array.as_view();
-    <ListView as TakeExecute>::take(array, &codes_ref, ctx)
+    let codes_ref = codes.array();
+    <ListView as TakeExecute>::take(array, codes_ref, ctx)
         .vortex_expect("take listview execute")
         .vortex_expect("ListView TakeExecute should not return None")
         .as_::<ListView>()
         .into_owned()
 }
 
 fn take_fixed_size_list(
-    array: &FixedSizeListArray,
-    codes: &PrimitiveArray,
+    array: ArrayView<'_, FixedSizeList>,
+    codes: ArrayView<'_, Primitive>,
     ctx: &mut ExecutionCtx,
 ) -> FixedSizeListArray {
-    let codes_ref = codes.clone().into_array();
-    let array = array.as_view();
-    <FixedSizeList as TakeExecute>::take(array, &codes_ref, ctx)
+    let codes_ref = codes.array();
+    <FixedSizeList as TakeExecute>::take(array, codes_ref, ctx)
         .vortex_expect("take fixed size list array")
         .vortex_expect("take fixed size list should not return None")
         .as_::<FixedSizeList>()
         .into_owned()
 }
 
-fn take_struct(array: &StructArray, codes: &PrimitiveArray) -> StructArray {
-    let codes_ref = codes.clone().into_array();
-    let array = array.as_view();
-    <Struct as TakeReduce>::take(array, &codes_ref)
+fn take_struct(array: ArrayView<'_, Struct>, codes: ArrayView<'_, Primitive>) -> StructArray {
+    let codes_ref = codes.array();
+    <Struct as TakeReduce>::take(array, codes_ref)
         .vortex_expect("take struct array")
         .vortex_expect("take struct should not return None")
         .as_::<Struct>()
         .into_owned()
 }
 
 fn take_extension(
-    array: &ExtensionArray,
-    codes: &PrimitiveArray,
+    array: ArrayView<'_, Extension>,
+    codes: ArrayView<'_, Primitive>,
     ctx: &mut ExecutionCtx,
 ) -> ExtensionArray {
-    let codes_ref = codes.clone().into_array();
-    let array = array.as_view();
-    <Extension as TakeExecute>::take(array, &codes_ref, ctx)
+    let codes_ref = codes.array();
+    <Extension as TakeExecute>::take(array, codes_ref, ctx)
         .vortex_expect("take extension storage")
         .vortex_expect("take extension should not return None")
         .as_::<Extension>()

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

@@ -3,13 +3,16 @@
 
 use std::hash::Hasher;
 
+use num_traits::AsPrimitive;
 use prost::Message;
 use smallvec::smallvec;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
 use vortex_error::vortex_err;
 use vortex_error::vortex_panic;
+use vortex_mask::AllOr;
+use vortex_mask::Mask;
 use vortex_session::VortexSession;
 use vortex_session::registry::CachedId;
 
@@ -24,6 +27,7 @@ use crate::ArrayEq;
 use crate::ArrayHash;
 use crate::ArrayRef;
 use crate::Canonical;
+use crate::CanonicalView;
 use crate::EqMode;
 use crate::IntoArray;
 use crate::array::Array;
@@ -34,17 +38,22 @@ use crate::array::VTable;
 use crate::array::with_empty_buffers;
 use crate::arrays::ConstantArray;
 use crate::arrays::Primitive;
+use crate::arrays::VarBinView;
 use crate::arrays::dict::DictArrayExt;
 use crate::arrays::dict::DictArraySlotsExt;
 use crate::arrays::dict::compute::rules::PARENT_RULES;
 use crate::arrays::dict::execute::take_canonical;
 use crate::buffer::BufferHandle;
 use crate::builders::ArrayBuilder;
+use crate::builders::VarBinBuilder;
 use crate::dtype::DType;
 use crate::dtype::Nullability;
+use crate::dtype::OffsetBuilderPType;
 use crate::dtype::PType;
 use crate::executor::ExecutionCtx;
 use crate::executor::ExecutionResult;
+use crate::match_each_integer_ptype;
+use crate::match_each_varbin_builder;
 use crate::require_child;
 use crate::scalar::Scalar;
 use crate::serde::ArrayChildren;
@@ -205,7 +214,7 @@ impl VTable for Dict {
 
         Ok(ExecutionResult::done(take_canonical(
             values.as_::<AnyCanonical>(),
-            &codes.downcast::<Primitive>(),
+            codes.as_::<Primitive>(),
             ctx,
         )?))
     }
@@ -222,8 +231,15 @@ impl VTable for Dict {
             )
             && !codes.validity()?.definitely_all_null()
         {
-            let codes = codes.into_owned();
-            let canonical = take_canonical(values, &codes, ctx)?.into_array();
+            if let CanonicalView::VarBinView(values) = values
+                && let Some(result) = match_each_varbin_builder!(builder, |builder| {
+                    let validity = array.validity()?.execute_mask(array.len(), ctx)?;
+                    append_dict_to_varbin(codes, values, validity, builder)
+                })
+            {
+                return result;
+            }
+            let canonical = take_canonical(values, codes, ctx)?.into_array();
             canonical.append_to_builder(builder, ctx)?;
             return Ok(());
         }
@@ -245,3 +261,86 @@ impl VTable for Dict {
         PARENT_RULES.evaluate(array, parent, child_idx)
     }
 }
+
+/// Gathers the dictionary values straight into `builder`.
+///
+/// The canonical route first takes the values to full logical length, which allocates and then
+/// re-reads a views buffer proportional to the row count. The dictionary is usually far smaller
+/// than the column, so resolving each code against it in place skips that intermediate entirely
+/// and leaves one `memcpy` per row as the only work.
+fn append_dict_to_varbin<O: OffsetBuilderPType>(
+    codes: ArrayView<'_, Primitive>,
+    values: ArrayView<'_, VarBinView>,
+    validity: Mask,
+    builder: &mut VarBinBuilder<O>,
+) -> VortexResult<()>
+where
+    usize: AsPrimitive<O>,
+{
+    let len = codes.as_ref().len();
+
+    // Resolve the dictionary's storage once so that looking up a code is an O(1) read.
+    let views = values.views();
+    let buffers = values
+        .data_buffers()
+        .iter()
+        .map(|buffer| buffer.as_host().as_slice())
+        .collect::<Vec<_>>();
+
+    match_each_integer_ptype!(codes.ptype(), |C| {
+        let codes = codes.as_slice::<C>();
+        let view = |row: usize| &views[AsPrimitive::<usize>::as_(codes[row])];
+
+        // Both passes below resolve a row through its code, so the byte total comes from the same
+        // walk over the valid rows that the copy will make.
+        let num_bytes = match validity.bit_buffer() {
+            AllOr::All => (0..len).map(|row| view(row).len() as usize).sum(),
+            AllOr::None => {
+                builder.push_nulls(len);
+                return Ok(());
+            }
+            AllOr::Some(bits) => {
+                let mut total = 0;
+                bits.for_each_set_index(|row| total += view(row).len() as usize);
+                total
+            }
+        };
+
+        builder.append_valid_slices(num_bytes, &validity, |row| view(row).bytes(&buffers))
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::VortexSessionExecute;
+    use crate::array_session;
+    use crate::arrays::PrimitiveArray;
+    use crate::arrays::VarBinViewArray;
+    use crate::arrays::dict::DictArray;
+    use crate::assert_arrays_eq;
+    use crate::dtype::Nullability::Nullable;
+
+    const LONG: &str = "a string that is far too long to be inlined in a view";
+
+    #[test]
+    fn append_to_builder_gathers_through_the_dictionary() -> VortexResult<()> {
+        let mut ctx = array_session().create_execution_ctx();
+        let dict = DictArray::try_new(
+            PrimitiveArray::from_option_iter([Some(0u32), Some(2), None, Some(1), Some(0)])
+                .into_array(),
+            VarBinViewArray::from_iter([Some(LONG), None, Some("short")], DType::Utf8(Nullable))
+                .into_array(),
+        )?;
+
+        let mut builder = VarBinBuilder::<i32>::new(DType::Utf8(Nullable));
+        dict.append_to_builder(&mut builder, &mut ctx)?;
+
+        let expected = VarBinViewArray::from_iter(
+            [Some(LONG), Some("short"), None, None, Some(LONG)],
+            DType::Utf8(Nullable),
+        );
+        assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx);
+        Ok(())
+    }
+}