From c82aa5856176153e3fb852a802979e551ea7b5d1 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 30 Jul 2026 22:16:46 +0100 Subject: [PATCH] Dict: gather dictionary values straight into a VarBinBuilder MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Appending a string dictionary to a builder went through `take_canonical`, which first materializes the values at full logical length — allocating and re-reading 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 and leaves one `memcpy` per row. The fast path indexes the codes and the dictionary's views and data buffers as host slices, so `varbin_fast_path_is_host` gates it on residency; a device-resident dictionary still takes the canonical route, where a device kernel can serve it. Signed-off-by: Robert Kruszewski Signed-off-by: Robert Kruszewski --- vortex-array/src/arrays/dict/vtable/mod.rs | 241 +++++++++++++++++++++ 1 file changed, 241 insertions(+) diff --git a/vortex-array/src/arrays/dict/vtable/mod.rs b/vortex-array/src/arrays/dict/vtable/mod.rs index d9ced9eeffa..892e74e2b94 100644 --- a/vortex-array/src/arrays/dict/vtable/mod.rs +++ b/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( + codes: &PrimitiveArray, + values: &VarBinViewArray, + builder: &mut VarBinBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> +where + usize: AsPrimitive, +{ + 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::>(); + + match_each_integer_ptype!(codes.ptype(), |C| { + let codes = codes.as_slice::(); + let validity = dict_validity(&codes_validity, codes, &values_validity); + let view = |row: usize| &views[AsPrimitive::::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( + 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); + + impl DeviceBuffer for FakeDeviceBuffer { + fn as_any(&self) -> &dyn Any { + self + } + + fn len(&self) -> usize { + self.0.len() + } + + fn alignment(&self) -> Alignment { + Alignment::of::() + } + + fn copy_to_host_sync(&self, alignment: Alignment) -> VortexResult { + Ok(ByteBuffer::copy_from_aligned(&self.0, alignment)) + } + + fn copy_to_host( + &self, + alignment: Alignment, + ) -> VortexResult>> { + let copied = self.copy_to_host_sync(alignment)?; + Ok(async move { Ok(copied) }.boxed()) + } + + fn slice(&self, range: Range) -> Arc { + Arc::new(Self(self.0[range].to_vec())) + } + + fn aligned(self: Arc, _alignment: Alignment) -> VortexResult> { + 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::::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)); + } +}