Dict: gather dictionary values straight into a VarBinBuilder

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

@@ -3,13 +3,17 @@
 
 use std::hash::Hasher;
 
+use num_traits::AsPrimitive;
 use prost::Message;
 use smallvec::smallvec;
+use vortex_buffer::BitBufferMut;
 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;
 
@@ -34,17 +38,24 @@ 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::VarBinBuilder;
 use crate::dtype::DType;
+use crate::dtype::IntegerPType;
 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;
@@ -223,6 +234,14 @@ impl VTable for Dict {
             && !codes.validity()?.definitely_all_null()
         {
             let codes = codes.into_owned();
+            if let Canonical::VarBinView(values) = Canonical::from(values)
+                && varbin_fast_path_is_host(&codes, &values)
+                && let Some(result) = match_each_varbin_builder!(builder, |builder| {
+                    append_dict_to_varbin(&codes, &values, builder, ctx)
+                })
+            {
+                return result;
+            }
             let canonical = take_canonical(values, &codes, ctx)?.into_array();
             canonical.append_to_builder(builder, ctx)?;
             return Ok(());
@@ -245,3 +264,225 @@ impl VTable for Dict {
         PARENT_RULES.evaluate(array, parent, child_idx)
     }
 }
+
+/// Whether every buffer that [`append_dict_to_varbin`] dereferences directly is host-resident.
+///
+/// The fast path reads the codes and the dictionary's views and data buffers as plain slices,
+/// which is only possible for host allocations. Device-resident dictionaries fall back to the
+/// canonical take, which goes through the compute stack and can therefore be served by a device
+/// kernel. This mirrors [`DictArrayExt::validate_all_values_referenced`], which likewise skips
+/// host-only work when the codes are not resident on the host.
+///
+/// This is a handful of enum discriminant checks over the array's buffer handles, so it stays
+/// outside the per-row loop and costs nothing measurable on the host-resident path.
+fn varbin_fast_path_is_host(codes: &PrimitiveArray, values: &VarBinViewArray) -> bool {
+    codes.buffer_handle().is_on_host()
+        && values.views_handle().is_on_host()
+        && values.data_buffers().iter().all(BufferHandle::is_on_host)
+}
+
+/// 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.
+///
+/// The caller must have checked [`varbin_fast_path_is_host`]; this function indexes the backing
+/// buffers as host slices and panics otherwise.
+fn append_dict_to_varbin<O: OffsetBuilderPType>(
+    codes: &PrimitiveArray,
+    values: &VarBinViewArray,
+    builder: &mut VarBinBuilder<O>,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<()>
+where
+    usize: AsPrimitive<O>,
+{
+    let len = codes.as_ref().len();
+    let codes_validity = codes.as_ref().validity()?.execute_mask(len, ctx)?;
+    let values_validity = values.validity()?.execute_mask(values.len(), ctx)?;
+
+    // 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 validity = dict_validity(&codes_validity, codes, &values_validity);
+        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))
+    })
+}
+
+/// Combines the codes' validity with the dictionary values' validity resolved through the codes.
+///
+/// A code pointing at a null dictionary entry produces a null row, matching what
+/// [`take_canonical`] derives from [`Validity::take`](crate::validity::Validity::take).
+fn dict_validity<C: IntegerPType>(
+    codes_validity: &Mask,
+    codes: &[C],
+    values_validity: &Mask,
+) -> Mask {
+    if values_validity.all_true() {
+        return codes_validity.clone();
+    }
+
+    let mut valid = BitBufferMut::new_unset(codes.len());
+    match codes_validity.bit_buffer() {
+        AllOr::All => {
+            for (row, &code) in codes.iter().enumerate() {
+                valid.set_to(row, values_validity.value(code.as_()));
+            }
+        }
+        AllOr::None => {}
+        AllOr::Some(bits) => bits.for_each_set_index(|row| {
+            valid.set_to(row, values_validity.value(codes[row].as_()));
+        }),
+    }
+    Mask::from_buffer(valid.freeze())
+}
+
+#[cfg(test)]
+mod tests {
+    use std::any::Any;
+    use std::ops::Range;
+    use std::sync::Arc;
+
+    use futures::FutureExt;
+    use futures::future::BoxFuture;
+    use vortex_buffer::Alignment;
+    use vortex_buffer::ByteBuffer;
+    use vortex_buffer::buffer;
+
+    use super::*;
+    use crate::VortexSessionExecute;
+    use crate::array_session;
+    use crate::arrays::VarBinViewArray;
+    use crate::arrays::dict::DictArray;
+    use crate::assert_arrays_eq;
+    use crate::buffer::DeviceBuffer;
+    use crate::dtype::Nullability::NonNullable;
+    use crate::dtype::Nullability::Nullable;
+    use crate::validity::Validity;
+
+    const LONG: &str = "a string that is far too long to be inlined in a view";
+
+    /// A stand-in for a real device allocation so residency handling can be exercised without a
+    /// GPU. Only the host/device discriminant matters here; the payload is never read back.
+    #[derive(Debug, PartialEq, Eq, Hash)]
+    struct FakeDeviceBuffer(Vec<u8>);
+
+    impl DeviceBuffer for FakeDeviceBuffer {
+        fn as_any(&self) -> &dyn Any {
+            self
+        }
+
+        fn len(&self) -> usize {
+            self.0.len()
+        }
+
+        fn alignment(&self) -> Alignment {
+            Alignment::of::<u128>()
+        }
+
+        fn copy_to_host_sync(&self, alignment: Alignment) -> VortexResult<ByteBuffer> {
+            Ok(ByteBuffer::copy_from_aligned(&self.0, alignment))
+        }
+
+        fn copy_to_host(
+            &self,
+            alignment: Alignment,
+        ) -> VortexResult<BoxFuture<'static, VortexResult<ByteBuffer>>> {
+            let copied = self.copy_to_host_sync(alignment)?;
+            Ok(async move { Ok(copied) }.boxed())
+        }
+
+        fn slice(&self, range: Range<usize>) -> Arc<dyn DeviceBuffer> {
+            Arc::new(Self(self.0[range].to_vec()))
+        }
+
+        fn aligned(self: Arc<Self>, _alignment: Alignment) -> VortexResult<Arc<dyn DeviceBuffer>> {
+            Ok(self)
+        }
+    }
+
+    fn to_device(handle: &BufferHandle) -> BufferHandle {
+        BufferHandle::new_device(Arc::new(FakeDeviceBuffer(handle.as_host().to_vec())))
+    }
+
+    #[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(())
+    }
+
+    /// The fast path indexes the codes and the dictionary's buffers as host slices, so it must be
+    /// skipped whenever any of them lives on a device. Regression test for the fast path panicking
+    /// on device-resident dictionaries instead of deferring to the canonical take.
+    #[test]
+    fn fast_path_is_skipped_for_device_resident_buffers() {
+        let codes = PrimitiveArray::new(buffer![0u32, 1, 0], Validity::NonNullable);
+        let values = VarBinViewArray::from_iter_str([LONG, "short"]);
+
+        assert!(varbin_fast_path_is_host(&codes, &values));
+
+        let device_codes = PrimitiveArray::from_buffer_handle(
+            to_device(codes.buffer_handle()),
+            PType::U32,
+            Validity::NonNullable,
+        );
+        assert!(!varbin_fast_path_is_host(&device_codes, &values));
+
+        let device_views = VarBinViewArray::new_handle(
+            to_device(values.views_handle()),
+            Arc::clone(values.data_buffers()),
+            DType::Utf8(NonNullable),
+            Validity::NonNullable,
+        );
+        assert!(!varbin_fast_path_is_host(&codes, &device_views));
+
+        let device_data = VarBinViewArray::new_handle(
+            values.views_handle().clone(),
+            values.data_buffers().iter().map(to_device).collect(),
+            DType::Utf8(NonNullable),
+            Validity::NonNullable,
+        );
+        assert!(!varbin_fast_path_is_host(&codes, &device_data));
+    }
+}