diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index f8676d76ef0..5489d22cfb9 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -124,6 +124,10 @@ harness = false name = "scalar_at_struct" harness = false +[[bench]] +name = "scalar_at_patches" +harness = false + [[bench]] name = "varbinview_compact" harness = false diff --git a/vortex-array/benches/scalar_at_patches.rs b/vortex-array/benches/scalar_at_patches.rs new file mode 100644 index 00000000000..42567fd63c8 --- /dev/null +++ b/vortex-array/benches/scalar_at_patches.rs @@ -0,0 +1,140 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::unwrap_used)] +#![expect(clippy::cast_possible_truncation)] + +use divan::Bencher; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex_array::IntoArray; +use vortex_array::patches::Patches; +use vortex_buffer::Buffer; + +fn main() { + divan::main(); +} + +const ARRAY_LEN: usize = 1_000_000; +const NUM_PATCHES: usize = 100; +const NUM_QUERIES: usize = 1_000; +const PATCH_LOW: usize = 100_000; +const PATCH_HIGH: usize = 110_000; + +fn narrow_band_patches() -> Patches { + let mut rng = StdRng::seed_from_u64(42); + let mut indices: Vec = (0..NUM_PATCHES) + .map(|_| rng.random_range((PATCH_LOW as u64)..(PATCH_HIGH as u64))) + .collect(); + indices.sort(); + indices.dedup(); + let values: Buffer = (0..indices.len() as i32).collect(); + Patches::new( + ARRAY_LEN, + 0, + Buffer::from(indices).into_array(), + values.into_array(), + None, + ) + .unwrap() +} + +fn full_range_patches() -> Patches { + let mut rng = StdRng::seed_from_u64(43); + let mut indices: Vec = (0..NUM_PATCHES) + .map(|_| rng.random_range(0..(ARRAY_LEN as u64))) + .collect(); + indices.sort(); + indices.dedup(); + let values: Buffer = (0..indices.len() as i32).collect(); + Patches::new( + ARRAY_LEN, + 0, + Buffer::from(indices).into_array(), + values.into_array(), + None, + ) + .unwrap() +} + +#[divan::bench] +fn search_index_below_min(bencher: Bencher) { + let patches = narrow_band_patches(); + let queries: Vec = (0..NUM_QUERIES).collect(); + + bencher.bench_local(|| { + for &q in &queries { + std::hint::black_box(patches.search_index(q).unwrap()); + } + }); +} + +#[divan::bench] +fn search_index_above_max(bencher: Bencher) { + let patches = narrow_band_patches(); + let queries: Vec = (PATCH_HIGH..(PATCH_HIGH + NUM_QUERIES)).collect(); + + bencher.bench_local(|| { + for &q in &queries { + std::hint::black_box(patches.search_index(q).unwrap()); + } + }); +} + +#[divan::bench] +fn search_index_mixed_out_of_range(bencher: Bencher) { + let patches = narrow_band_patches(); + let queries: Vec = (0..NUM_QUERIES / 2) + .map(|i| i * 100) + .chain((0..NUM_QUERIES / 2).map(|i| PATCH_HIGH + i * 50)) + .collect(); + + bencher.bench_local(|| { + for &q in &queries { + std::hint::black_box(patches.search_index(q).unwrap()); + } + }); +} + +#[divan::bench] +fn search_index_in_range(bencher: Bencher) { + let patches = narrow_band_patches(); + let mut rng = StdRng::seed_from_u64(7); + let queries: Vec = (0..NUM_QUERIES) + .map(|_| rng.random_range(PATCH_LOW..PATCH_HIGH)) + .collect(); + + bencher.bench_local(|| { + for &q in &queries { + std::hint::black_box(patches.search_index(q).unwrap()); + } + }); +} + +#[divan::bench] +fn search_index_full_range_random(bencher: Bencher) { + let patches = full_range_patches(); + let mut rng = StdRng::seed_from_u64(11); + let queries: Vec = (0..NUM_QUERIES) + .map(|_| rng.random_range(0..ARRAY_LEN)) + .collect(); + + bencher.bench_local(|| { + for &q in &queries { + std::hint::black_box(patches.search_index(q).unwrap()); + } + }); +} + +#[divan::bench] +fn get_patched_above_max(bencher: Bencher) { + let patches = narrow_band_patches(); + let queries: Vec = (PATCH_HIGH..(PATCH_HIGH + NUM_QUERIES)).collect(); + + bencher.bench_local(|| { + for &q in &queries { + std::hint::black_box(patches.get_patched(q).unwrap()); + } + }); +} diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index fa532546b50..43f267867f7 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -5,6 +5,7 @@ use std::cmp::Ordering; use std::fmt::Debug; use std::hash::Hash; use std::ops::Range; +use std::sync::OnceLock; use num_traits::NumCast; use vortex_buffer::BitBuffer; @@ -146,6 +147,10 @@ pub struct Patches { /// `offset_within_chunk` is necessary in order to keep track of how many /// elements were sliced off within the chunk. offset_within_chunk: Option, + /// Cached `indices[0] - offset`. + min_index: OnceLock, + /// Cached `indices[len - 1] - offset`. + max_index: OnceLock, } impl Patches { @@ -172,6 +177,9 @@ impl Patches { ); vortex_ensure!(!indices.is_empty(), "Patch indices must not be empty"); + let min_index = OnceLock::new(); + let max_index = OnceLock::new(); + // Perform validation of components when they are host-resident. // This is not possible to do eagerly when the data is on GPU memory. if indices.is_host() && values.is_host() { @@ -185,6 +193,14 @@ impl Patches { "Patch indices {max:?}, offset {offset} are longer than the array length {array_len}" ); + // Pre-populate bounds caches for the search_index fast path. + let min = usize::try_from( + &indices.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())?, + ) + .map_err(|_| vortex_err!("indices must be a number"))?; + let _ = max_index.set(max - offset); + let _ = min_index.set(min - offset); + #[cfg(debug_assertions)] { use crate::VortexSessionExecute; @@ -205,6 +221,8 @@ impl Patches { chunk_offsets: chunk_offsets.clone(), // Initialize with `Some(0)` only if `chunk_offsets` are set. offset_within_chunk: chunk_offsets.map(|_| 0), + min_index, + max_index, }) } @@ -232,6 +250,8 @@ impl Patches { values, chunk_offsets, offset_within_chunk, + min_index: OnceLock::new(), + max_index: OnceLock::new(), } } @@ -394,6 +414,10 @@ impl Patches { /// [`SearchResult::Found(patch_idx)`]: SearchResult::Found /// [`SearchResult::NotFound(insertion_point)`]: SearchResult::NotFound pub fn search_index(&self, index: usize) -> VortexResult { + if let Some(result) = self.search_index_out_of_range(index) { + return Ok(result); + } + if self.chunk_offsets.is_some() { return self.search_index_chunked(index); } @@ -401,6 +425,20 @@ impl Patches { Self::search_index_binary_search(&self.indices, index + self.offset) } + /// Returns the search result if `index` falls outside the cached bounds, + /// or `None` if the bounds are uncached or `index` is in range. + fn search_index_out_of_range(&self, index: usize) -> Option { + let min = *self.min_index.get()?; + let max = *self.max_index.get()?; + if index < min { + Some(SearchResult::NotFound(0)) + } else if index > max { + Some(SearchResult::NotFound(self.indices.len())) + } else { + None + } + } + /// Binary searches for `needle` in the indices array. /// /// # Returns @@ -551,19 +589,27 @@ impl Patches { }) } - /// Returns the minimum patch index + /// Returns the minimum patch index. pub fn min_index(&self) -> VortexResult { + if let Some(&v) = self.min_index.get() { + return Ok(v); + } let first = self .indices .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())? .as_primitive() .as_::() .ok_or_else(|| vortex_err!("index does not fit in usize"))?; - Ok(first - self.offset) + let result = first - self.offset; + let _ = self.min_index.set(result); + Ok(result) } - /// Returns the maximum patch index + /// Returns the maximum patch index. pub fn max_index(&self) -> VortexResult { + if let Some(&v) = self.max_index.get() { + return Ok(v); + } let last = self .indices .execute_scalar( @@ -573,7 +619,9 @@ impl Patches { .as_primitive() .as_::() .ok_or_else(|| vortex_err!("index does not fit in usize"))?; - Ok(last - self.offset) + let result = last - self.offset; + let _ = self.max_index.set(result); + Ok(result) } /// Filter the patches by a mask, resulting in new patches for the filtered array. @@ -648,6 +696,8 @@ impl Patches { // TODO(0ax1): Chunk offsets are invalid after a filter is applied. chunk_offsets: None, offset_within_chunk: self.offset_within_chunk, + min_index: OnceLock::new(), + max_index: OnceLock::new(), })) } @@ -695,6 +745,8 @@ impl Patches { values, chunk_offsets: new_chunk_offsets, offset_within_chunk, + min_index: OnceLock::new(), + max_index: OnceLock::new(), })) } @@ -830,6 +882,8 @@ impl Patches { .take(PrimitiveArray::new(values_indices, values_validity).into_array())?, chunk_offsets: None, offset_within_chunk: Some(0), // Reset when creating new Patches. + min_index: OnceLock::new(), + max_index: OnceLock::new(), })) } @@ -879,6 +933,8 @@ impl Patches { // TODO(0ax1): Chunk offsets are invalid after take is applied. chunk_offsets: None, offset_within_chunk: self.offset_within_chunk, + min_index: OnceLock::new(), + max_index: OnceLock::new(), })) } @@ -902,6 +958,9 @@ impl Patches { values, chunk_offsets: self.chunk_offsets, offset_within_chunk: self.offset_within_chunk, + // indices and offset are preserved, so the cached bounds remain valid. + min_index: self.min_index, + max_index: self.max_index, }) } } @@ -1750,6 +1809,66 @@ mod test { assert_eq!(patches.search_index(9).unwrap(), SearchResult::NotFound(3)); } + #[test] + fn test_search_index_out_of_range_fast_path() { + let patches = Patches::new( + 100, + 0, + buffer![10u64, 20, 30, 40].into_array(), + buffer![1i32, 2, 3, 4].into_array(), + None, + ) + .unwrap(); + + assert_eq!( + patches.search_index_out_of_range(0), + Some(SearchResult::NotFound(0)) + ); + assert_eq!( + patches.search_index_out_of_range(9), + Some(SearchResult::NotFound(0)) + ); + assert_eq!( + patches.search_index_out_of_range(41), + Some(SearchResult::NotFound(4)) + ); + assert_eq!( + patches.search_index_out_of_range(99), + Some(SearchResult::NotFound(4)) + ); + assert_eq!(patches.search_index_out_of_range(10), None); + assert_eq!(patches.search_index_out_of_range(25), None); + assert_eq!(patches.search_index_out_of_range(40), None); + + assert_eq!(patches.search_index(5).unwrap(), SearchResult::NotFound(0)); + assert_eq!(patches.search_index(50).unwrap(), SearchResult::NotFound(4)); + } + + #[test] + fn test_search_index_out_of_range_with_offset() { + let patches = Patches::new( + 100, + 0, + buffer![10u64, 50, 90].into_array(), + buffer![1i32, 2, 3].into_array(), + None, + ) + .unwrap(); + let sliced = patches.slice(40..95).unwrap().unwrap(); + + assert_eq!(sliced.min_index().unwrap(), 10); + assert_eq!(sliced.max_index().unwrap(), 50); + assert_eq!( + sliced.search_index_out_of_range(5), + Some(SearchResult::NotFound(0)) + ); + assert_eq!( + sliced.search_index_out_of_range(54), + Some(SearchResult::NotFound(2)) + ); + assert_eq!(sliced.search_index_out_of_range(30), None); + } + #[test] fn test_mask_boundary_patches() { // Test masking patches at array boundaries