diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 70fba4acdcf..ec34c41422d 100644 --- a/vortex-array/Cargo.toml +++ b/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 diff --git a/vortex-array/benches/dict_append_varbin.rs b/vortex-array/benches/dict_append_varbin.rs new file mode 100644 index 00000000000..7235ab0e079 --- /dev/null +++ b/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 = 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)); +} diff --git a/vortex-array/src/arrays/dict/array.rs b/vortex-array/src/arrays/dict/array.rs index ba1201693b6..f23349cf0bf 100644 --- a/vortex-array/src/arrays/dict/array.rs +++ b/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>, + #[case] codes: Vec>, + #[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::(&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(); diff --git a/vortex-array/src/arrays/dict/vtable/mod.rs b/vortex-array/src/arrays/dict/vtable/mod.rs index d9ced9eeffa..efe72a078f1 100644 --- a/vortex-array/src/arrays/dict/vtable/mod.rs +++ b/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::() + && !codes.validity()?.definitely_all_null() + && builder.as_any().is::() + { + 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::(), @@ -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::(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::() + .vortex_expect("caller checked that the builder is a DynVarBinBuilder"); + + match_each_integer_ptype!(codes.ptype(), |P| { + let codes = codes.as_slice::

(); + 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(()) +}