Dict compute to use codes if values are too large (#2804)

Fix #2796

Author: onursatici

encodings/dict/benches/dict_compare.rs

@@ -4,15 +4,15 @@ use std::str::from_utf8;
 
 use vortex_array::accessor::ArrayAccessor;
 use vortex_array::arrays::{ConstantArray, VarBinArray, VarBinViewArray};
-use vortex_array::compute::{Operator, compare};
+use vortex_array::compute::{Operator, compare, slice};
 use vortex_dict::builders::dict_encode;
 use vortex_dict::test::{gen_primitive_for_dict, gen_varbin_words};
 
 fn main() {
     divan::main();
 }
 
-const BENCH_ARGS: &[(usize, usize)] = &[
+const LENGTH_AND_UNIQUE_VALUES: &[(usize, usize)] = &[
     // length, unique_values
     (10_000, 2),
     (10_000, 4),
@@ -30,7 +30,7 @@ const BENCH_ARGS: &[(usize, usize)] = &[
     (100_000, 2048),
 ];
 
-#[divan::bench(args = BENCH_ARGS)]
+#[divan::bench(args = LENGTH_AND_UNIQUE_VALUES)]
 fn bench_compare_primitive(bencher: divan::Bencher, (len, uniqueness): (usize, usize)) {
     let primitive_arr = gen_primitive_for_dict::<i32>(len, uniqueness);
     let dict = dict_encode(&primitive_arr).unwrap();
@@ -41,7 +41,7 @@ fn bench_compare_primitive(bencher: divan::Bencher, (len, uniqueness): (usize, u
         .bench_refs(|dict| compare(dict, &ConstantArray::new(value, len), Operator::Eq).unwrap())
 }
 
-#[divan::bench(args = BENCH_ARGS)]
+#[divan::bench(args = LENGTH_AND_UNIQUE_VALUES)]
 fn bench_compare_varbin(bencher: divan::Bencher, (len, uniqueness): (usize, usize)) {
     let varbin_arr = VarBinArray::from(gen_varbin_words(len, uniqueness));
     let dict = dict_encode(&varbin_arr).unwrap();
@@ -55,7 +55,7 @@ fn bench_compare_varbin(bencher: divan::Bencher, (len, uniqueness): (usize, usiz
         .bench_refs(|dict| compare(dict, &ConstantArray::new(value, len), Operator::Eq).unwrap())
 }
 
-#[divan::bench(args = BENCH_ARGS)]
+#[divan::bench(args = LENGTH_AND_UNIQUE_VALUES)]
 fn bench_compare_varbinview(bencher: divan::Bencher, (len, uniqueness): (usize, usize)) {
     let varbinview_arr = VarBinViewArray::from_iter_str(gen_varbin_words(len, uniqueness));
     let dict = dict_encode(&varbinview_arr).unwrap();
@@ -67,3 +67,48 @@ fn bench_compare_varbinview(bencher: divan::Bencher, (len, uniqueness): (usize,
         .with_inputs(|| dict.clone())
         .bench_refs(|dict| compare(dict, &ConstantArray::new(value, len), Operator::Eq).unwrap())
 }
+
+const CODES_AND_VALUES_LENGTHS: &[(usize, usize)] = &[
+    (1_000, 10_000),
+    (2_000, 10_000),
+    (2_500, 10_000),
+    (3_333, 10_000),
+    (5_000, 10_000),
+    (7_500, 10_000),
+    (9_999, 10_000),
+    (10_000, 10_000),
+    (20_000, 10_000),
+];
+
+#[divan::bench(args = CODES_AND_VALUES_LENGTHS)]
+fn bench_compare_sliced_dict_primitive(
+    bencher: divan::Bencher,
+    (codes_len, values_len): (usize, usize),
+) {
+    let primitive_arr = gen_primitive_for_dict::<i32>(codes_len.max(values_len), values_len);
+    let dict = dict_encode(&primitive_arr).unwrap();
+    let dict = slice(&dict, 0, codes_len).unwrap();
+    let value = primitive_arr.as_slice::<i32>()[0];
+
+    bencher.with_inputs(|| dict.clone()).bench_refs(|dict| {
+        compare(dict, &ConstantArray::new(value, codes_len), Operator::Eq).unwrap()
+    })
+}
+
+#[divan::bench(args = CODES_AND_VALUES_LENGTHS)]
+fn bench_compare_sliced_dict_varbinview(
+    bencher: divan::Bencher,
+    (codes_len, values_len): (usize, usize),
+) {
+    let varbin_arr = VarBinArray::from(gen_varbin_words(codes_len.max(values_len), values_len));
+    let dict = dict_encode(&varbin_arr).unwrap();
+    let dict = slice(&dict, 0, codes_len).unwrap();
+    let bytes = varbin_arr
+        .with_iterator(|i| i.next().unwrap().unwrap().to_vec())
+        .unwrap();
+    let value = from_utf8(bytes.as_slice()).unwrap();
+
+    bencher.with_inputs(|| dict.clone()).bench_refs(|dict| {
+        compare(dict, &ConstantArray::new(value, codes_len), Operator::Eq).unwrap()
+    })
+}

encodings/dict/src/compute/binary_numeric.rs

@@ -13,6 +13,11 @@ impl BinaryNumericFn<&DictArray> for DictEncoding {
         rhs: &dyn Array,
         op: BinaryNumericOperator,
     ) -> VortexResult<Option<ArrayRef>> {
+        // if we have more values than codes, it is faster to canonicalise first.
+        if array.values().len() > array.codes().len() {
+            return Ok(None);
+        }
+
         let Some(rhs_scalar) = rhs.as_constant() else {
             return Ok(None);
         };

encodings/dict/src/compute/compare.rs

@@ -17,6 +17,10 @@ impl CompareFn<&DictArray> for DictEncoding {
         rhs: &dyn Array,
         operator: Operator,
     ) -> VortexResult<Option<ArrayRef>> {
+        // if we have more values than codes, it is faster to canonicalise first.
+        if lhs.values().len() > lhs.codes().len() {
+            return Ok(None);
+        }
         // If the RHS is constant, then we just need to compare against our encoded values.
         if let Some(rhs) = rhs.as_constant() {
             let compare_result = compare(

encodings/dict/src/compute/like.rs

@@ -12,6 +12,10 @@ impl LikeFn<&DictArray> for DictEncoding {
         pattern: &dyn Array,
         options: LikeOptions,
     ) -> VortexResult<Option<ArrayRef>> {
+        // if we have more values than codes, it is faster to canonicalise first.
+        if array.values().len() > array.codes().len() {
+            return Ok(None);
+        }
         if let Some(pattern) = pattern.as_constant() {
             let pattern = ConstantArray::new(pattern, array.values().len()).into_array();
             let values = like(array.values(), &pattern, options)?;