fix imports, move fsst+dict bench into fsst

Signed-off-by: Onur Satici <[email protected]>

Author: onursatici

Cargo.lock

@@ -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 }

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_array_builder"
+harness = false
+required-features = ["test-harness"]

encodings/fsst/Cargo.toml

@@ -4,11 +4,12 @@
 use divan::Bencher;
 use rand::distr::{Distribution, StandardUniform};
 use vortex_array::arrays::ChunkedArray;
+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_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();

…ay/benches/chunked_dict_array_builder.rs …st/benches/chunked_dict_array_builder.rsvortex-array/benches/chunked_dict_array_builder.rs renamed to encodings/fsst/benches/chunked_dict_array_builder.rs vortex-array/benches/chunked_dict_array_builder.rs renamed to encodings/fsst/benches/chunked_dict_array_builder.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/lib.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()
+}

encodings/fsst/src/test_utils.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();

vortex-array/benches/dict_compare.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() {

vortex-array/benches/dict_compress.rs

@@ -7,9 +7,8 @@ use divan::Bencher;
 use rand::rngs::StdRng;
 use rand::{Rng, SeedableRng};
 use vortex_array::IntoArray;
-use vortex_array::arrays::PrimitiveArray;
+use vortex_array::arrays::{DictArray, PrimitiveArray};
 use vortex_array::compute::{mask, warm_up_vtables};
-use vortex_dict::DictArray;
 use vortex_mask::Mask;
 
 fn main() {

vortex-array/benches/dict_mask.rs

@@ -4,16 +4,17 @@
 use std::fmt::Debug;
 use std::hash::Hash;
 
-use vortex_array::stats::{ArrayStats, StatsSetRef};
-use vortex_array::vtable::{ArrayVTable, NotSupported, VTable, ValidityVTable};
-use vortex_array::{
-    Array, ArrayEq, ArrayHash, ArrayRef, EncodingId, EncodingRef, Precision, ToCanonical, vtable,
-};
 use vortex_buffer::BitBuffer;
 use vortex_dtype::{DType, match_each_integer_ptype};
 use vortex_error::{VortexExpect as _, VortexResult, vortex_bail};
 use vortex_mask::{AllOr, Mask};
 
+use crate::stats::{ArrayStats, StatsSetRef};
+use crate::vtable::{ArrayVTable, NotSupported, VTable, ValidityVTable};
+use crate::{
+    Array, ArrayEq, ArrayHash, ArrayRef, EncodingId, EncodingRef, Precision, ToCanonical, vtable,
+};
+
 vtable!(Dict);
 
 impl VTable for DictVTable {
@@ -194,17 +195,17 @@ mod test {
     use rand::distr::{Distribution, StandardUniform};
     use rand::prelude::StdRng;
     use rand::{Rng, SeedableRng};
-    use vortex_array::arrays::{ChunkedArray, PrimitiveArray};
-    use vortex_array::builders::builder_with_capacity;
-    use vortex_array::validity::Validity;
-    use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical, assert_arrays_eq};
     use vortex_buffer::{BitBuffer, buffer};
     use vortex_dtype::Nullability::NonNullable;
     use vortex_dtype::{DType, NativePType, PType, UnsignedPType};
     use vortex_error::{VortexExpect, VortexUnwrap, vortex_panic};
     use vortex_mask::AllOr;
 
-    use crate::DictArray;
+    use crate::arrays::dict::DictArray;
+    use crate::arrays::{ChunkedArray, PrimitiveArray};
+    use crate::builders::builder_with_capacity;
+    use crate::validity::Validity;
+    use crate::{Array, ArrayRef, IntoArray, ToCanonical, assert_arrays_eq};
 
     #[test]
     fn nullable_codes_validity() {