feat: take_search constant time index lookups

encodings/fastlanes/src/bitpacking/serde.rs

@@ -71,14 +71,10 @@ impl SerdeVTable<BitPackedVTable> for BitPackedVTable {
             }
         };
 
-        let validity_idx = if let Some(patches_meta) = &metadata.patches {
-            if patches_meta.chunk_offsets_dtype().is_some() {
-                3
-            } else {
-                2
-            }
-        } else {
-            0
+        let validity_idx = match &metadata.patches {
+            None => 0,
+            Some(patches_meta) if patches_meta.chunk_offsets_dtype().is_some() => 3,
+            Some(_) => 2,
         };
 
         let validity = load_validity(validity_idx)?;

vortex-array/benches/take_patches.rs

@@ -38,7 +38,22 @@ const BENCH_ARGS: &[(f64, f64)] = &[
 #[divan::bench(args = BENCH_ARGS)]
 fn take_search(bencher: Bencher, (patches_sparsity, index_multiple): (f64, f64)) {
     let mut rng = StdRng::seed_from_u64(0);
-    let patches = fixture(16384, patches_sparsity, &mut rng);
+    let patches = fixture(65536, patches_sparsity, &mut rng);
+    let indices = indices(
+        patches.array_len(),
+        (patches.array_len() as f64 * index_multiple) as usize,
+        &mut rng,
+    );
+
+    bencher
+        .with_inputs(|| (&patches, indices.clone()))
+        .bench_values(|(patches, indices)| patches.take_search(indices.to_primitive(), false));
+}
+
+#[divan::bench(args = BENCH_ARGS)]
+fn take_search_chunked(bencher: Bencher, (patches_sparsity, index_multiple): (f64, f64)) {
+    let mut rng = StdRng::seed_from_u64(0);
+    let patches = fixture_with_chunk_offsets(65536, patches_sparsity, &mut rng);
     let indices = indices(
         patches.array_len(),
         (patches.array_len() as f64 * index_multiple) as usize,
@@ -53,7 +68,7 @@ fn take_search(bencher: Bencher, (patches_sparsity, index_multiple): (f64, f64))
 #[divan::bench(args = BENCH_ARGS)]
 fn take_map(bencher: Bencher, (patches_sparsity, index_multiple): (f64, f64)) {
     let mut rng = StdRng::seed_from_u64(0);
-    let patches = fixture(16384, patches_sparsity, &mut rng);
+    let patches = fixture(65536, patches_sparsity, &mut rng);
     let indices = indices(
         patches.array_len(),
         (patches.array_len() as f64 * index_multiple) as usize,
@@ -82,6 +97,32 @@ fn fixture(len: usize, sparsity: f64, rng: &mut StdRng) -> Patches {
     )
 }
 
+fn fixture_with_chunk_offsets(len: usize, sparsity: f64, rng: &mut StdRng) -> Patches {
+    let patch_indices = (0..len)
+        .filter(|_| rng.random_bool(sparsity))
+        .map(|x| x as u64)
+        .collect::<Vec<u64>>();
+
+    let sparse_len = patch_indices.len();
+    let values = Buffer::from_iter((0..sparse_len).map(|x| x as u64)).into_array();
+
+    const PATCH_CHUNK_SIZE: usize = 1024;
+    let chunk_offsets: Vec<u64> = (0..len)
+        .step_by(PATCH_CHUNK_SIZE)
+        .map(|chunk_start| {
+            patch_indices.partition_point(|&idx| (idx as usize) < chunk_start) as u64
+        })
+        .collect();
+
+    Patches::new(
+        len,
+        0,
+        Buffer::from(patch_indices).into_array(),
+        values,
+        Some(Buffer::from(chunk_offsets).into_array()),
+    )
+}
+
 fn indices(array_len: usize, n_indices: usize, rng: &mut StdRng) -> ArrayRef {
     Buffer::from_iter((0..n_indices).map(|_| rng.random_range(0..(array_len as u64)))).into_array()
 }