diff --git a/crates/store/re_chunk/src/util.rs b/crates/store/re_chunk/src/util.rs index 30ed831529bb..375d175e4d80 100644 --- a/crates/store/re_chunk/src/util.rs +++ b/crates/store/re_chunk/src/util.rs @@ -377,6 +377,28 @@ pub fn take_array( "index arrays with validity bits are technically valid, but generally a sign that something went wrong", ); + if indices.len() == array.len() { + let indices = indices.values().as_slice(); + + let starts_at_zero = || indices[0] == O::zero(); + let is_consecutive = || { + indices + .windows(2) + .all(|values| values[1] == values[0] + O::one()) + }; + + if starts_at_zero() && is_consecutive() { + #[allow(clippy::unwrap_used)] + return array + .clone() + .as_any() + .downcast_ref::() + // Unwrap: that's initial type that we got. + .unwrap() + .clone(); + } + } + #[allow(clippy::unwrap_used)] arrow2::compute::take::take(array, indices) // Unwrap: this literally cannot fail. diff --git a/crates/store/re_chunk/tests/memory_test.rs b/crates/store/re_chunk/tests/memory_test.rs index 9a502f675f3d..a23387914f89 100644 --- a/crates/store/re_chunk/tests/memory_test.rs +++ b/crates/store/re_chunk/tests/memory_test.rs @@ -61,7 +61,7 @@ use arrow2::{ }, offset::Offsets as ArrowOffsets, }; -use itertools::Itertools as _; +use itertools::Itertools; #[test] fn filter_does_allocate() { @@ -125,6 +125,72 @@ fn filter_does_allocate() { } } +#[test] +fn filter_empty_or_full_is_noop() { + re_log::setup_logging(); + + const NUM_SCALARS: i64 = 10_000_000; + + let (((unfiltered, unfiltered_size_bytes), (filtered, filtered_size_bytes)), total_size_bytes) = + memory_use(|| { + let unfiltered = memory_use(|| { + let scalars = ArrowPrimitiveArray::from_vec((0..NUM_SCALARS).collect_vec()); + ArrowListArray::::new( + ArrowListArray::::default_datatype(scalars.data_type().clone()), + ArrowOffsets::try_from_lengths( + std::iter::repeat(NUM_SCALARS as usize / 10).take(10), + ) + .unwrap() + .into(), + scalars.to_boxed(), + None, + ) + }); + + let filter = ArrowBooleanArray::from_slice( + std::iter::repeat(true) + .take(unfiltered.0.len()) + .collect_vec(), + ); + let filtered = memory_use(|| re_chunk::util::filter_array(&unfiltered.0, &filter)); + + (unfiltered, filtered) + }); + + eprintln!( + "unfiltered={} filtered={} total={}", + re_format::format_bytes(unfiltered_size_bytes as _), + re_format::format_bytes(filtered_size_bytes as _), + re_format::format_bytes(total_size_bytes as _), + ); + + assert!( + filtered_size_bytes < 1_000, + "filtered array should be the size of a few empty datastructures at most" + ); + + { + let unfiltered = unfiltered + .values() + .as_any() + .downcast_ref::>() + .unwrap(); + let filtered = filtered + .values() + .as_any() + .downcast_ref::>() + .unwrap(); + + assert!( + std::ptr::eq( + unfiltered.values().as_ptr_range().start, + filtered.values().as_ptr_range().start + ), + "whole thing should be a noop -- pointers should match" + ); + } +} + #[test] // TODO(cmc): That's the end goal, but it is simply impossible with `ListArray`'s encoding. // See `Chunk::take_array`'s doc-comment for more information. @@ -191,3 +257,65 @@ fn take_does_not_allocate() { ); } } + +#[test] +fn take_empty_or_full_is_noop() { + re_log::setup_logging(); + + const NUM_SCALARS: i64 = 10_000_000; + + let (((untaken, untaken_size_bytes), (taken, taken_size_bytes)), total_size_bytes) = + memory_use(|| { + let untaken = memory_use(|| { + let scalars = ArrowPrimitiveArray::from_vec((0..NUM_SCALARS).collect_vec()); + ArrowListArray::::new( + ArrowListArray::::default_datatype(scalars.data_type().clone()), + ArrowOffsets::try_from_lengths( + std::iter::repeat(NUM_SCALARS as usize / 10).take(10), + ) + .unwrap() + .into(), + scalars.to_boxed(), + None, + ) + }); + + let indices = ArrowPrimitiveArray::from_vec((0..untaken.0.len() as i32).collect_vec()); + let taken = memory_use(|| re_chunk::util::take_array(&untaken.0, &indices)); + + (untaken, taken) + }); + + eprintln!( + "untaken={} taken={} total={}", + re_format::format_bytes(untaken_size_bytes as _), + re_format::format_bytes(taken_size_bytes as _), + re_format::format_bytes(total_size_bytes as _), + ); + + assert!( + taken_size_bytes < 1_000, + "taken array should be the size of a few empty datastructures at most" + ); + + { + let untaken = untaken + .values() + .as_any() + .downcast_ref::>() + .unwrap(); + let taken = taken + .values() + .as_any() + .downcast_ref::>() + .unwrap(); + + assert!( + std::ptr::eq( + untaken.values().as_ptr_range().start, + taken.values().as_ptr_range().start + ), + "whole thing should be a noop -- pointers should match" + ); + } +}