diff --git a/vortex-array/src/arrays/dict/execute.rs b/vortex-array/src/arrays/dict/execute.rs index 4a3ccda43c2..6462ed38d16 100644 --- a/vortex-array/src/arrays/dict/execute.rs +++ b/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,35 +40,34 @@ 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 { - 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)?) } @@ -74,32 +75,29 @@ pub(crate) fn take_canonical( } /// 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 { - let codes_ref = codes.clone().into_array(); - let array = array.as_view(); - Ok(::take(array, &codes_ref, ctx)? + let codes_ref = codes.array(); + Ok(::take(array, codes_ref, ctx)? .vortex_expect("take bool should not return None") .as_::() .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(); - ::take(array, &codes_ref, ctx) + let codes_ref = codes.array(); + ::take(array, codes_ref, ctx) .vortex_expect("take primitive array") .vortex_expect("take primitive should not return None") .as_::() @@ -107,13 +105,12 @@ fn take_primitive( } 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(); - ::take(array, &codes_ref, ctx) + let codes_ref = codes.array(); + ::take(array, codes_ref, ctx) .vortex_expect("take decimal array") .vortex_expect("take decimal should not return None") .as_::() @@ -121,13 +118,12 @@ fn take_decimal( } 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(); - ::take(array, &codes_ref, ctx) + let codes_ref = codes.array(); + ::take(array, codes_ref, ctx) .vortex_expect("take varbinview array") .vortex_expect("take varbinview should not return None") .as_::() @@ -135,13 +131,12 @@ fn take_varbinview( } 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(); - ::take(array, &codes_ref, ctx) + let codes_ref = codes.array(); + ::take(array, codes_ref, ctx) .vortex_expect("take listview execute") .vortex_expect("ListView TakeExecute should not return None") .as_::() @@ -149,23 +144,21 @@ fn take_listview( } 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(); - ::take(array, &codes_ref, ctx) + let codes_ref = codes.array(); + ::take(array, codes_ref, ctx) .vortex_expect("take fixed size list array") .vortex_expect("take fixed size list should not return None") .as_::() .into_owned() } -fn take_struct(array: &StructArray, codes: &PrimitiveArray) -> StructArray { - let codes_ref = codes.clone().into_array(); - let array = array.as_view(); - ::take(array, &codes_ref) +fn take_struct(array: ArrayView<'_, Struct>, codes: ArrayView<'_, Primitive>) -> StructArray { + let codes_ref = codes.array(); + ::take(array, codes_ref) .vortex_expect("take struct array") .vortex_expect("take struct should not return None") .as_::() @@ -173,13 +166,12 @@ fn take_struct(array: &StructArray, codes: &PrimitiveArray) -> StructArray { } 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(); - ::take(array, &codes_ref, ctx) + let codes_ref = codes.array(); + ::take(array, codes_ref, ctx) .vortex_expect("take extension storage") .vortex_expect("take extension should not return None") .as_::() diff --git a/vortex-array/src/arrays/dict/vtable/mod.rs b/vortex-array/src/arrays/dict/vtable/mod.rs index d9ced9eeffa..62e324f4ba5 100644 --- a/vortex-array/src/arrays/dict/vtable/mod.rs +++ b/vortex-array/src/arrays/dict/vtable/mod.rs @@ -3,6 +3,7 @@ use std::hash::Hasher; +use num_traits::AsPrimitive; use prost::Message; use smallvec::smallvec; use vortex_error::VortexResult; @@ -10,6 +11,8 @@ 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_::(), - &codes.downcast::(), + codes.as_::(), 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( + codes: ArrayView<'_, Primitive>, + values: ArrayView<'_, VarBinView>, + validity: Mask, + builder: &mut VarBinBuilder, +) -> VortexResult<()> +where + usize: AsPrimitive, +{ + 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::>(); + + match_each_integer_ptype!(codes.ptype(), |C| { + let codes = codes.as_slice::(); + 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)) + }) +} + +#[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::::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(()) + } +}