perf(array): gather dict byte values straight into the offset builder

vortex-array/Cargo.toml

@@ -203,6 +203,10 @@ harness = false
 name = "validity_is_valid"
 harness = false
 
+[[bench]]
+name = "dict_append_varbin"
+harness = false
+
 [[bench]]
 name = "dict_unreferenced_mask"
 harness = false

vortex-array/benches/dict_append_varbin.rs

@@ -0,0 +1,71 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Benchmarks appending a UTF-8 `DictArray` into a `DynVarBinBuilder`, the path a `pa.string()`
+//! export of a low-cardinality string column takes.
+//!
+//! `cardinality` is the dictionary size; the win grows as it falls relative to `len`, because the
+//! skipped intermediate is proportional to `len`.
+
+#![allow(clippy::unwrap_used)]
+
+use std::sync::LazyLock;
+
+use divan::Bencher;
+use divan::black_box;
+use vortex_array::IntoArray;
+use vortex_array::VortexSessionExecute;
+use vortex_array::array_session;
+use vortex_array::arrays::DictArray;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::arrays::VarBinViewArray;
+use vortex_array::builders::DynVarBinBuilder;
+use vortex_array::dtype::DType;
+use vortex_array::dtype::Nullability;
+use vortex_session::VortexSession;
+
+fn main() {
+    divan::main();
+}
+
+const SHAPES: &[(usize, usize)] = &[(4096, 16), (4096, 256), (4096, 2048)];
+
+static SESSION: LazyLock<VortexSession> = LazyLock::new(array_session);
+
+fn dict(len: usize, cardinality: usize, nulls: bool) -> DictArray {
+    let values = VarBinViewArray::from_iter(
+        (0..cardinality)
+            .map(|i| Some(format!("https://example.com/some/path/segment/{i}").into_bytes())),
+        DType::Utf8(Nullability::Nullable),
+    );
+    let codes = PrimitiveArray::from_option_iter((0..len).map(|i| {
+        if nulls && i.is_multiple_of(4) {
+            None
+        } else {
+            Some(u32::try_from(i % cardinality).unwrap())
+        }
+    }));
+    DictArray::try_new(codes.into_array(), values.into_array()).unwrap()
+}
+
+fn bench_append(bencher: Bencher, array: DictArray) {
+    let mut ctx = SESSION.create_execution_ctx();
+    let len = array.len();
+    let dtype = array.dtype().clone();
+    let array = array.into_array();
+    bencher.bench_local(|| {
+        let mut builder = DynVarBinBuilder::with_capacity(dtype.clone(), false, len);
+        array.append_to_builder(&mut builder, &mut ctx).unwrap();
+        black_box(builder.finish_into_varbin())
+    });
+}
+
+#[divan::bench(args = SHAPES)]
+fn all_valid(bencher: Bencher, shape: (usize, usize)) {
+    bench_append(bencher, dict(shape.0, shape.1, false));
+}
+
+#[divan::bench(args = SHAPES)]
+fn with_null_codes(bencher: Bencher, shape: (usize, usize)) {
+    bench_append(bencher, dict(shape.0, shape.1, true));
+}

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

@@ -280,6 +280,7 @@ mod test {
     use rand::distr::Distribution;
     use rand::distr::StandardUniform;
     use rand::prelude::StdRng;
+    use rstest::rstest;
     use vortex_buffer::BitBuffer;
     use vortex_buffer::buffer;
     use vortex_error::VortexExpect;
@@ -288,6 +289,7 @@ mod test {
     use vortex_mask::AllOr;
 
     use crate::ArrayRef;
+    use crate::Canonical;
     use crate::IntoArray;
     use crate::VortexSessionExecute;
     use crate::array_session;
@@ -300,6 +302,7 @@ mod test {
     use crate::builders::builder_with_capacity;
     use crate::dtype::DType;
     use crate::dtype::NativePType;
+    use crate::dtype::Nullability;
     use crate::dtype::Nullability::NonNullable;
     use crate::dtype::PType;
     use crate::dtype::UnsignedPType;
@@ -457,6 +460,55 @@ mod test {
         Ok(())
     }
 
+    /// The gather path must agree with canonicalizing the dict and appending that.
+    #[rstest]
+    #[case::all_valid(vec![Some("aa"), Some("bb"), Some("cc")], vec![Some(2u8), Some(0), Some(1), Some(2)])]
+    #[case::null_codes(vec![Some("aa"), Some("bb")], vec![Some(1u8), None, Some(0), None])]
+    #[case::null_values(vec![Some("aa"), None], vec![Some(1u8), Some(0), Some(1)])]
+    #[case::null_codes_and_values(vec![Some("aa"), None], vec![Some(1u8), None, Some(0)])]
+    #[case::heap_values(
+        vec![
+            Some("a string comfortably longer than twelve bytes"),
+            Some("another string that also exceeds twelve bytes"),
+        ],
+        vec![Some(1u8), Some(0), Some(1), Some(0)],
+    )]
+    #[case::mixed_inlined_and_heap(
+        vec![Some("tiny"), Some("a string comfortably longer than twelve bytes"), None],
+        vec![Some(0u8), Some(1), Some(2), Some(1), Some(0)],
+    )]
+    #[case::leading_and_trailing_null_codes(
+        vec![Some("aa"), Some("bb")],
+        vec![None, Some(0u8), Some(1), None],
+    )]
+    #[case::single_value(vec![Some("only")], vec![Some(0u8), Some(0), Some(0)])]
+    fn dict_byte_gather_matches_canonical(
+        #[case] values: Vec<Option<&str>>,
+        #[case] codes: Vec<Option<u8>>,
+        #[values(false, true)] large_offsets: bool,
+        #[values(false, true)] sliced: bool,
+    ) -> VortexResult<()> {
+        let mut ctx = array_session().create_execution_ctx();
+        let values = VarBinViewArray::from_iter(values, DType::Utf8(Nullability::Nullable));
+        let codes = PrimitiveArray::from_option_iter(codes);
+        let dict = DictArray::try_new(codes.into_array(), values.into_array())?.into_array();
+
+        let dict = if sliced && dict.len() > 1 {
+            dict.slice(1..dict.len())?
+        } else {
+            dict
+        };
+
+        let expected = dict.clone().execute::<Canonical>(&mut ctx)?.into_array();
+
+        let mut builder =
+            DynVarBinBuilder::with_capacity(dict.dtype().clone(), large_offsets, dict.len());
+        dict.append_to_builder(&mut builder, &mut ctx)?;
+
+        assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx);
+        Ok(())
+    }
+
     #[test]
     fn test_dict_array_from_primitive_chunks() -> VortexResult<()> {
         let mut ctx = array_session().create_execution_ctx();

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::VortexExpect;
 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_session::VortexSession;
 use vortex_session::registry::CachedId;
 
@@ -34,17 +37,21 @@ use crate::array::VTable;
 use crate::array::with_empty_buffers;
 use crate::arrays::ConstantArray;
 use crate::arrays::Primitive;
+use crate::arrays::PrimitiveArray;
+use crate::arrays::VarBinViewArray;
 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::DynVarBinBuilder;
 use crate::dtype::DType;
 use crate::dtype::Nullability;
 use crate::dtype::PType;
 use crate::executor::ExecutionCtx;
 use crate::executor::ExecutionResult;
+use crate::match_each_integer_ptype;
 use crate::require_child;
 use crate::scalar::Scalar;
 use crate::serde::ArrayChildren;
@@ -215,6 +222,18 @@ impl VTable for Dict {
         builder: &mut dyn ArrayBuilder,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<()> {
+        // The generic path below takes the values to full logical length, allocating an
+        // intermediate that is then copied into the builder again. Gather by code instead.
+        if matches!(array.dtype(), DType::Utf8(_) | DType::Binary(_))
+            && !array.is_empty()
+            && let Some(codes) = array.codes().as_opt::<Primitive>()
+            && !codes.validity()?.definitely_all_null()
+            && builder.as_any().is::<DynVarBinBuilder>()
+        {
+            let codes = codes.into_owned();
+            return append_dict_bytes(array, &codes, builder, ctx);
+        }
+
         if !array.is_empty()
             && let (Some(codes), Some(values)) = (
                 array.codes().as_opt::<Primitive>(),
@@ -245,3 +264,71 @@ impl VTable for Dict {
         PARENT_RULES.evaluate(array, parent, child_idx)
     }
 }
+
+/// Gathers UTF-8 or binary dictionary values into a [`DynVarBinBuilder`] by code.
+///
+/// `builder` must be a [`DynVarBinBuilder`]; the caller checks this before dispatching here.
+fn append_dict_bytes(
+    array: ArrayView<'_, Dict>,
+    codes: &PrimitiveArray,
+    builder: &mut dyn ArrayBuilder,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<()> {
+    let len = array.len();
+    let values = array.values().clone().execute::<VarBinViewArray>(ctx)?;
+    let values_mask = values.validity()?.execute_mask(values.len(), ctx)?;
+    let codes_mask = codes.as_ref().validity()?.execute_mask(len, ctx)?;
+
+    let views = values.views();
+    let buffers: Vec<&[u8]> = (0..values.data_buffers().len())
+        .map(|idx| values.buffer(idx).as_slice())
+        .collect();
+
+    let view_bytes = |index: usize| -> &[u8] {
+        let view = &views[index];
+        if view.is_inlined() {
+            view.as_inlined().value()
+        } else {
+            let reference = view.as_view();
+            &buffers[reference.buffer_index as usize][reference.as_range()]
+        }
+    };
+
+    let builder = builder
+        .as_any_mut()
+        .downcast_mut::<DynVarBinBuilder>()
+        .vortex_expect("caller checked that the builder is a DynVarBinBuilder");
+
+    match_each_integer_ptype!(codes.ptype(), |P| {
+        let codes = codes.as_slice::<P>();
+        let append_code = |builder: &mut DynVarBinBuilder, row: usize| {
+            let code: usize = codes[row].as_();
+            if values_mask.value(code) {
+                builder.append_n_values(view_bytes(code), 1);
+            } else {
+                // The code may point at a null dictionary entry.
+                builder.append_nulls(1);
+            }
+        };
+
+        match codes_mask.bit_buffer() {
+            AllOr::All => {
+                for row in 0..len {
+                    append_code(&mut *builder, row);
+                }
+            }
+            AllOr::None => builder.append_nulls(len),
+            AllOr::Some(valid) => {
+                let mut row = 0;
+                valid.for_each_set_index(|index| {
+                    builder.append_nulls(index - row);
+                    append_code(&mut *builder, index);
+                    row = index + 1;
+                });
+                builder.append_nulls(len - row);
+            }
+        }
+    });
+
+    Ok(())
+}