feat: use buffer for VarBinView views

encodings/dict/Cargo.toml

@@ -36,3 +36,7 @@ simplelog = { workspace = true }
 [[bench]]
 name = "dict_compress"
 harness = false
+
+[[bench]]
+name = "dict_canonical"
+harness = false

encodings/dict/benches/dict_canonical.rs

@@ -0,0 +1,34 @@
+#![allow(clippy::unwrap_used)]
+
+use criterion::{criterion_group, criterion_main, Criterion};
+use vortex::array::VarBinArray;
+use vortex::{IntoArray, IntoCanonical};
+use vortex_dict::{dict_encode_varbin, DictArray};
+use vortex_dtype::{DType, Nullability};
+
+fn fixture(len: usize) -> DictArray {
+    let values = [
+        Some("inlined"),
+        None,
+        Some("not inlined but repeated often"),
+    ];
+
+    let strings = VarBinArray::from_iter(
+        values.into_iter().cycle().take(len),
+        DType::Utf8(Nullability::Nullable),
+    );
+
+    let (codes, values) = dict_encode_varbin(&strings);
+    DictArray::try_new(codes.into_array(), values.into_array()).unwrap()
+}
+
+fn bench_canonical(c: &mut Criterion) {
+    let dict_array = fixture(1024).into_array();
+
+    c.bench_function("canonical", |b| {
+        b.iter(|| dict_array.clone().into_canonical().unwrap())
+    });
+}
+
+criterion_group!(bench_dict_canonical, bench_canonical);
+criterion_main!(bench_dict_canonical);

encodings/dict/src/array.rs

@@ -1,14 +1,14 @@
 use std::fmt::{Debug, Display};
 
-use arrow_buffer::BooleanBuffer;
+use arrow_buffer::{BooleanBuffer, ScalarBuffer};
 use serde::{Deserialize, Serialize};
 use vortex::array::visitor::{AcceptArrayVisitor, ArrayVisitor};
-use vortex::array::BoolArray;
-use vortex::compute::take;
-use vortex::compute::unary::scalar_at;
+use vortex::array::{BoolArray, ConstantArray, VarBinViewArray};
+use vortex::compute::unary::{scalar_at, try_cast};
+use vortex::compute::{compare, take, Operator};
 use vortex::encoding::ids;
 use vortex::stats::StatsSet;
-use vortex::validity::{ArrayValidity, LogicalValidity};
+use vortex::validity::{ArrayValidity, LogicalValidity, Validity};
 use vortex::{
     impl_encoding, Array, ArrayDType, ArrayTrait, Canonical, IntoArray, IntoArrayVariant,
     IntoCanonical,
@@ -41,6 +41,7 @@ impl DictArray {
             DictMetadata {
                 codes_ptype: PType::try_from(codes.dtype())
                     .vortex_expect("codes dtype must be uint"),
+
                 values_len: values.len(),
             },
             [values, codes].into(),
@@ -67,11 +68,52 @@ impl ArrayTrait for DictArray {}
 
 impl IntoCanonical for DictArray {
     fn into_canonical(self) -> VortexResult<Canonical> {
-        let canonical_values: Array = self.values().into_canonical()?.into();
-        take(canonical_values, self.codes())?.into_canonical()
+        match self.dtype() {
+            DType::Utf8(_) | DType::Binary(_) => canonicalize_string(self),
+            _ => canonicalize_primitive(self),
+        }
     }
 }
 
+/// Canonicalize a set of codes and values.
+fn canonicalize_string(array: DictArray) -> VortexResult<Canonical> {
+    let values = array.values().into_varbinview()?;
+    let codes = try_cast(array.codes(), PType::U64.into())?.into_primitive()?;
+
+    let value_views = ScalarBuffer::<u128>::from(values.views().clone().into_arrow());
+
+    // Gather the views from value_views into full_views using the dictionary codes.
+    let full_views: Vec<u128> = codes
+        .maybe_null_slice::<u64>()
+        .iter()
+        .map(|code| value_views[*code as usize])
+        .collect();
+
+    let validity = if array.dtype().is_nullable() {
+        // For nullable arrays, a code of 0 indicates null value.
+        Validity::Array(compare(
+            codes.as_ref(),
+            ConstantArray::new(0u64, codes.len()).as_ref(),
+            Operator::Eq,
+        )?)
+    } else {
+        Validity::NonNullable
+    };
+
+    VarBinViewArray::try_new(
+        full_views.into(),
+        values.buffers().collect(),
+        array.dtype().clone(),
+        validity,
+    )
+    .map(Canonical::VarBinView)
+}
+
+fn canonicalize_primitive(array: DictArray) -> VortexResult<Canonical> {
+    let canonical_values: Array = array.values().into_canonical()?.into();
+    take(canonical_values, array.codes())?.into_canonical()
+}
+
 impl ArrayValidity for DictArray {
     fn is_valid(&self, index: usize) -> bool {
         let values_index = scalar_at(self.codes(), index)

encodings/dict/src/compute.rs

@@ -42,9 +42,6 @@ impl ScalarAtFn for DictArray {
 
 impl TakeFn for DictArray {
     fn take(&self, indices: &Array) -> VortexResult<Array> {
-        // Dict
-        //   codes: 0 0 1
-        //   dict: a b c d e f g h
         let codes = take(self.codes(), indices)?;
         Self::try_new(codes, self.values()).map(|a| a.into_array())
     }

vortex-array/src/array/chunked/canonical.rs

@@ -1,4 +1,4 @@
-use arrow_buffer::{BooleanBufferBuilder, Buffer, MutableBuffer, ScalarBuffer};
+use arrow_buffer::{BooleanBufferBuilder, MutableBuffer, ScalarBuffer};
 use vortex_dtype::{DType, PType, StructDType};
 use vortex_error::{vortex_bail, vortex_err, ErrString, VortexResult};
 
@@ -179,11 +179,7 @@ fn pack_primitives(
         buffer.extend_from_slice(chunk.buffer());
     }
 
-    Ok(PrimitiveArray::new(
-        Buffer::from(buffer).into(),
-        ptype,
-        validity,
-    ))
+    Ok(PrimitiveArray::new(buffer.into(), ptype, validity))
 }
 
 /// Builds a new [VarBinViewArray] by repacking the values from the chunks into a single
@@ -231,8 +227,8 @@ fn pack_views(
         }
     }
 
-    let views_buffer: Buffer = ScalarBuffer::<u128>::from(views).into_inner();
-    VarBinViewArray::try_new(Array::from(views_buffer), buffers, dtype.clone(), validity)
+    let views_buffer = ScalarBuffer::<u128>::from(views).into_inner();
+    VarBinViewArray::try_new(views_buffer.into(), buffers, dtype.clone(), validity)
 }
 
 #[cfg(test)]

vortex-array/src/array/constant/canonical.rs

@@ -1,9 +1,9 @@
 use arrow_array::builder::make_view;
-use arrow_buffer::{BooleanBuffer, BufferBuilder};
+use arrow_buffer::{BooleanBuffer, BufferBuilder, MutableBuffer};
 use vortex_buffer::Buffer;
 use vortex_dtype::{match_each_native_ptype, DType, Nullability, PType};
 use vortex_error::{vortex_bail, VortexResult};
-use vortex_scalar::{BinaryScalar, BoolScalar, Scalar, Utf8Scalar};
+use vortex_scalar::{BinaryScalar, BoolScalar, Utf8Scalar};
 
 use crate::array::constant::ConstantArray;
 use crate::array::primitive::PrimitiveArray;
@@ -70,10 +70,10 @@ fn canonical_byte_view(
 ) -> VortexResult<VarBinViewArray> {
     match scalar_bytes {
         None => {
-            let views = ConstantArray::new(Scalar::null(dtype.clone()), len);
+            let views = MutableBuffer::from(Vec::<u128>::with_capacity(1));
 
             VarBinViewArray::try_new(
-                views.into_array(),
+                views.into(),
                 Vec::new(),
                 dtype.clone(),
                 Validity::AllInvalid,
@@ -100,9 +100,7 @@ fn canonical_byte_view(
             //   add u128 PType, see https://github.com/spiraldb/vortex/issues/1110
             let mut views = BufferBuilder::<u128>::new(len);
             views.append_n(len, view);
-            let views =
-                PrimitiveArray::new(views.finish().into(), PType::U8, Validity::NonNullable)
-                    .into_array();
+            let views = views.finish().into();
 
             let validity = if dtype.nullability() == Nullability::NonNullable {
                 Validity::NonNullable

vortex-array/src/array/varbin/mod.rs

@@ -62,7 +62,7 @@ impl VarBinArray {
             vortex_bail!(MismatchedTypes: "utf8 or binary", dtype);
         }
         if dtype.is_nullable() == (validity == Validity::NonNullable) {
-            vortex_bail!("incorrect validity {:?}", validity);
+            vortex_bail!("incorrect validity {:?} for {}", validity, dtype);
         }
 
         let length = offsets.len() - 1;

vortex-array/src/array/varbinview/compute.rs

@@ -3,19 +3,21 @@ use std::sync::Arc;
 use arrow_array::cast::AsArray;
 use arrow_array::types::ByteViewType;
 use arrow_array::{Datum, GenericByteViewArray};
+use arrow_buffer::ScalarBuffer;
 use arrow_ord::cmp;
 use arrow_schema::DataType;
 use vortex_buffer::Buffer;
+use vortex_dtype::PType;
 use vortex_error::{vortex_bail, VortexResult, VortexUnwrap};
 use vortex_scalar::Scalar;
 
 use crate::array::varbin::varbin_scalar;
 use crate::array::varbinview::{VarBinViewArray, VIEW_SIZE_BYTES};
-use crate::array::{varbinview_as_arrow, ConstantArray};
+use crate::array::ConstantArray;
 use crate::arrow::FromArrowArray;
-use crate::compute::unary::ScalarAtFn;
-use crate::compute::{slice, ArrayCompute, MaybeCompareFn, Operator, SliceFn, TakeFn};
-use crate::{Array, ArrayDType, IntoArray, IntoCanonical};
+use crate::compute::unary::{try_cast, ScalarAtFn};
+use crate::compute::{ArrayCompute, MaybeCompareFn, Operator, SliceFn, TakeFn};
+use crate::{Array, ArrayDType, IntoArray, IntoArrayVariant, IntoCanonical};
 
 impl ArrayCompute for VarBinViewArray {
     fn scalar_at(&self) -> Option<&dyn ScalarAtFn> {
@@ -45,11 +47,8 @@ impl ScalarAtFn for VarBinViewArray {
 impl SliceFn for VarBinViewArray {
     fn slice(&self, start: usize, stop: usize) -> VortexResult<Array> {
         Ok(Self::try_new(
-            slice(
-                self.views(),
-                start * VIEW_SIZE_BYTES,
-                stop * VIEW_SIZE_BYTES,
-            )?,
+            self.views()
+                .slice(start * VIEW_SIZE_BYTES..stop * VIEW_SIZE_BYTES),
             (0..self.metadata().buffer_lens.len())
                 .map(|i| self.buffer(i))
                 .collect::<Vec<_>>(),
@@ -63,12 +62,21 @@ impl SliceFn for VarBinViewArray {
 /// Take involves creating a new array that references the old array, just with the given set of views.
 impl TakeFn for VarBinViewArray {
     fn take(&self, indices: &Array) -> VortexResult<Array> {
-        let array_ref = varbinview_as_arrow(self);
-        let indices_arrow = indices.clone().into_canonical()?.into_arrow()?;
-
-        let take_arrow = arrow_select::take::take(&array_ref, &indices_arrow, None)?;
-        let nullable = take_arrow.is_nullable();
-        Ok(Array::from_arrow(take_arrow, nullable))
+        let views = ScalarBuffer::<u128>::from(self.views().clone().into_arrow());
+
+        let taken: Vec<u128> = try_cast(indices, PType::U64.into())?
+            .into_primitive()?
+            .maybe_null_slice::<u64>()
+            .iter()
+            .map(|idx| views[*idx as usize])
+            .collect();
+        VarBinViewArray::try_new(
+            ScalarBuffer::from(taken).into_inner().into(),
+            self.buffers().collect(),
+            self.dtype().clone(),
+            self.validity().take(indices)?,
+        )
+        .map(|a| a.into_array())
     }
 }