move vortex-dict into vortex-array

Cargo.toml

@@ -226,7 +226,6 @@ vortex-compute = { version = "0.1.0", path = "./vortex-compute", default-feature
 vortex-datafusion = { version = "0.1.0", path = "./vortex-datafusion", default-features = false }
 vortex-datetime-parts = { version = "0.1.0", path = "./encodings/datetime-parts", default-features = false }
 vortex-decimal-byte-parts = { version = "0.1.0", path = "encodings/decimal-byte-parts", default-features = false }
-vortex-dict = { version = "0.1.0", path = "./encodings/dict", default-features = false }
 vortex-dtype = { version = "0.1.0", path = "./vortex-dtype", default-features = false }
 vortex-error = { version = "0.1.0", path = "./vortex-error", default-features = false }
 vortex-fastlanes = { version = "0.1.0", path = "./encodings/fastlanes", default-features = false }

encodings/fsst/Cargo.toml

@@ -20,6 +20,7 @@ workspace = true
 async-trait = { workspace = true }
 fsst-rs = { workspace = true }
 prost = { workspace = true }
+rand = { workspace = true, optional = true }
 vortex-array = { workspace = true }
 vortex-buffer = { workspace = true }
 vortex-dtype = { workspace = true }
@@ -28,6 +29,9 @@ vortex-mask = { workspace = true }
 vortex-scalar = { workspace = true }
 vortex-vector = { workspace = true }
 
+[features]
+test-harness = ["dep:rand", "vortex-array/test-harness"]
+
 [dev-dependencies]
 divan = { workspace = true }
 itertools = { workspace = true }
@@ -38,3 +42,8 @@ vortex-array = { workspace = true, features = ["test-harness"] }
 [[bench]]
 name = "fsst_compress"
 harness = false
+
+[[bench]]
+name = "chunked_dict_fsst_builder"
+harness = false
+required-features = ["test-harness"]

encodings/dict/benches/chunked_dict_array_builder.rs → encodings/fsst/benches/chunked_dict_fsst_builder.rs

@@ -2,13 +2,12 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use divan::Bencher;
-use rand::distr::{Distribution, StandardUniform};
 use vortex_array::arrays::ChunkedArray;
 use vortex_array::builders::builder_with_capacity;
 use vortex_array::compute::warm_up_vtables;
 use vortex_array::{Array, ArrayRef, IntoArray};
-use vortex_dict::test::{gen_dict_fsst_test_data, gen_dict_primitive_chunks};
 use vortex_dtype::NativePType;
+use vortex_fsst::test_utils::gen_dict_fsst_test_data;
 
 fn main() {
     warm_up_vtables();
@@ -24,36 +23,6 @@ const BENCH_ARGS: &[(usize, usize, usize)] = &[
     (1000, 1000, 100),
 ];
 
-#[divan::bench(types = [u32, u64, f32, f64], args = BENCH_ARGS)]
-fn chunked_dict_primitive_canonical_into<T: NativePType>(
-    bencher: Bencher,
-    (len, unique_values, chunk_count): (usize, usize, usize),
-) where
-    StandardUniform: Distribution<T>,
-{
-    let chunk = gen_dict_primitive_chunks::<T, u16>(len, unique_values, chunk_count);
-
-    bencher.with_inputs(|| chunk.clone()).bench_values(|chunk| {
-        let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count);
-        chunk.append_to_builder(builder.as_mut());
-        builder.finish()
-    })
-}
-
-#[divan::bench(types = [u32, u64, f32, f64], args = BENCH_ARGS)]
-fn chunked_dict_primitive_into_canonical<T: NativePType>(
-    bencher: Bencher,
-    (len, unique_values, chunk_count): (usize, usize, usize),
-) where
-    StandardUniform: Distribution<T>,
-{
-    let chunk = gen_dict_primitive_chunks::<T, u16>(len, unique_values, chunk_count);
-
-    bencher
-        .with_inputs(|| chunk.clone())
-        .bench_values(|chunk| chunk.to_canonical())
-}
-
 fn make_dict_fsst_chunks<T: NativePType>(
     len: usize,
     unique_values: usize,

encodings/fsst/src/lib.rs

@@ -17,6 +17,8 @@ mod compress;
 mod compute;
 mod ops;
 mod serde;
+#[cfg(feature = "test-harness")]
+pub mod test_utils;
 #[cfg(test)]
 mod tests;
 

encodings/fsst/src/test_utils.rs

@@ -0,0 +1,55 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#![allow(clippy::unwrap_used)]
+
+use rand::prelude::StdRng;
+use rand::{Rng, SeedableRng};
+use vortex_array::arrays::{DictArray, PrimitiveArray, VarBinArray};
+use vortex_array::{ArrayRef, IntoArray};
+use vortex_dtype::{DType, NativePType, Nullability};
+use vortex_error::VortexUnwrap;
+
+use crate::{fsst_compress, fsst_train_compressor};
+
+pub fn gen_fsst_test_data(len: usize, avg_str_len: usize, unique_chars: u8) -> ArrayRef {
+    let mut rng = StdRng::seed_from_u64(0);
+    let mut strings = Vec::with_capacity(len);
+
+    for _ in 0..len {
+        // Generate a random string with length around `avg_len`. The number of possible
+        // characters within the random string is defined by `unique_chars`.
+        let len = avg_str_len * rng.random_range(50..=150) / 100;
+        strings.push(Some(
+            (0..len)
+                .map(|_| rng.random_range(b'a'..(b'a' + unique_chars)))
+                .collect::<Vec<u8>>(),
+        ));
+    }
+
+    let varbin = VarBinArray::from_iter(
+        strings
+            .into_iter()
+            .map(|opt_s| opt_s.map(Vec::into_boxed_slice)),
+        DType::Binary(Nullability::NonNullable),
+    );
+    let compressor = fsst_train_compressor(varbin.as_ref()).vortex_unwrap();
+
+    fsst_compress(varbin.as_ref(), &compressor)
+        .vortex_unwrap()
+        .into_array()
+}
+
+pub fn gen_dict_fsst_test_data<T: NativePType>(
+    len: usize,
+    unique_values: usize,
+    str_len: usize,
+    unique_char_count: u8,
+) -> DictArray {
+    let values = gen_fsst_test_data(len, str_len, unique_char_count);
+    let mut rng = StdRng::seed_from_u64(0);
+    let codes = (0..len)
+        .map(|_| T::from(rng.random_range(0..unique_values)).unwrap())
+        .collect::<PrimitiveArray>();
+    DictArray::try_new(codes.into_array(), values).vortex_unwrap()
+}

vortex-array/Cargo.toml

@@ -140,3 +140,22 @@ harness = false
 name = "expr_large_struct_pack"
 path = "benches/expr/large_struct_pack.rs"
 harness = false
+
+[[bench]]
+name = "chunked_dict_builder"
+harness = false
+required-features = ["test-harness"]
+
+[[bench]]
+name = "dict_compress"
+harness = false
+required-features = ["test-harness"]
+
+[[bench]]
+name = "dict_compare"
+harness = false
+required-features = ["test-harness"]
+
+[[bench]]
+name = "dict_mask"
+harness = false

vortex-array/benches/chunked_dict_builder.rs

@@ -0,0 +1,54 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use divan::Bencher;
+use rand::distr::{Distribution, StandardUniform};
+use vortex_array::Array;
+use vortex_array::arrays::dict_test::gen_dict_primitive_chunks;
+use vortex_array::builders::builder_with_capacity;
+use vortex_array::compute::warm_up_vtables;
+use vortex_dtype::NativePType;
+
+fn main() {
+    warm_up_vtables();
+    divan::main();
+}
+
+const BENCH_ARGS: &[(usize, usize, usize)] = &[
+    (1000, 10, 10),
+    (1000, 100, 10),
+    (1000, 1000, 10),
+    (1000, 10, 100),
+    (1000, 100, 100),
+    (1000, 1000, 100),
+];
+
+#[divan::bench(types = [u32, u64, f32, f64], args = BENCH_ARGS)]
+fn chunked_dict_primitive_canonical_into<T: NativePType>(
+    bencher: Bencher,
+    (len, unique_values, chunk_count): (usize, usize, usize),
+) where
+    StandardUniform: Distribution<T>,
+{
+    let chunk = gen_dict_primitive_chunks::<T, u16>(len, unique_values, chunk_count);
+
+    bencher.with_inputs(|| chunk.clone()).bench_values(|chunk| {
+        let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count);
+        chunk.append_to_builder(builder.as_mut());
+        builder.finish()
+    })
+}
+
+#[divan::bench(types = [u32, u64, f32, f64], args = BENCH_ARGS)]
+fn chunked_dict_primitive_into_canonical<T: NativePType>(
+    bencher: Bencher,
+    (len, unique_values, chunk_count): (usize, usize, usize),
+) where
+    StandardUniform: Distribution<T>,
+{
+    let chunk = gen_dict_primitive_chunks::<T, u16>(len, unique_values, chunk_count);
+
+    bencher
+        .with_inputs(|| chunk.clone())
+        .bench_values(|chunk| chunk.to_canonical())
+}

encodings/dict/benches/dict_compare.rs → vortex-array/benches/dict_compare.rs

@@ -6,10 +6,10 @@
 use std::str::from_utf8;
 
 use vortex_array::accessor::ArrayAccessor;
+use vortex_array::arrays::dict_test::{gen_primitive_for_dict, gen_varbin_words};
 use vortex_array::arrays::{ConstantArray, VarBinArray, VarBinViewArray};
+use vortex_array::builders::dict::dict_encode;
 use vortex_array::compute::{Operator, compare, warm_up_vtables};
-use vortex_dict::builders::dict_encode;
-use vortex_dict::test::{gen_primitive_for_dict, gen_varbin_words};
 
 fn main() {
     warm_up_vtables();

encodings/dict/benches/dict_compress.rs → vortex-array/benches/dict_compress.rs

@@ -5,10 +5,10 @@
 
 use divan::Bencher;
 use rand::distr::{Distribution, StandardUniform};
+use vortex_array::arrays::dict_test::{gen_primitive_for_dict, gen_varbin_words};
 use vortex_array::arrays::{VarBinArray, VarBinViewArray};
+use vortex_array::builders::dict::dict_encode;
 use vortex_array::compute::warm_up_vtables;
-use vortex_dict::builders::dict_encode;
-use vortex_dict::test::{gen_primitive_for_dict, gen_varbin_words};
 use vortex_dtype::NativePType;
 
 fn main() {