Add ChunkStore::drop_time_range (#7602)

### What
Adds `ChunkStore::drop_time_range` to drop all events within some time
range. Will be used to implement undo.

* Part of #3135
* Extracted from #7546

---------

Co-authored-by: Clement Rey <cr.rey.clement@gmail.com>

crates/store/re_chunk/src/chunk.rs

@@ -116,6 +116,18 @@ impl PartialEq for Chunk {
 }
 
 impl Chunk {
+ /// Returns a version of us with a new [`ChunkId`].
+ ///
+ /// Reminder:
+ /// * The returned [`Chunk`] will re-use the exact same [`RowId`]s as `self`.
+ /// * Duplicated [`RowId`]s in the `ChunkStore` is undefined behavior.
+ #[must_use]
+ #[inline]
+ pub fn with_id(mut self, id: ChunkId) -> Self {
+ self.id = id;
+ self
+ }
+
  /// Returns `true` is two [`Chunk`]s are similar, although not byte-for-byte equal.
  ///
  /// In particular, this ignores chunks and row IDs, as well as temporal timestamps.

crates/store/re_chunk/src/shuffle.rs

@@ -111,6 +111,9 @@ impl Chunk {
  /// The underlying arrow data will be copied and shuffled in memory in order to make it contiguous.
  ///
  /// This is a no-op if the underlying timeline is already sorted appropriately (happy path).
+ ///
+ /// WARNING: the returned chunk has the same old [`crate::ChunkId`]! Change it with [`Self::with_id`].
+ #[must_use]
  pub fn sorted_by_timeline_if_unsorted(&self, timeline: &Timeline) -> Self {
  let mut chunk = self.clone();
 

crates/store/re_chunk/src/slice.rs

@@ -61,6 +61,9 @@ impl Chunk {
  /// This cannot fail nor panic: `index` and `len` will be capped so that they cannot
  /// run out of bounds.
  /// This can result in an empty [`Chunk`] being returned if the slice is completely OOB.
+ ///
+ /// WARNING: the returned chunk has the same old [`crate::ChunkId`]! Change it with [`Self::with_id`].
+ #[must_use]
  #[inline]
  pub fn row_sliced(&self, index: usize, len: usize) -> Self {
  re_tracing::profile_function!();
@@ -144,6 +147,9 @@ impl Chunk {
  ///
  /// If `timeline` is not found within the [`Chunk`], the end result will be the same as the
  /// current chunk but without any timeline column.
+ ///
+ /// WARNING: the returned chunk has the same old [`crate::ChunkId`]! Change it with [`Self::with_id`].
+ #[must_use]
  #[inline]
  pub fn timeline_sliced(&self, timeline: Timeline) -> Self {
  let Self {
@@ -184,6 +190,9 @@ impl Chunk {
  ///
  /// If `component_name` is not found within the [`Chunk`], the end result will be the same as the
  /// current chunk but without any component column.
+ ///
+ /// WARNING: the returned chunk has the same old [`crate::ChunkId`]! Change it with [`Self::with_id`].
+ #[must_use]
  #[inline]
  pub fn component_sliced(&self, component_name: ComponentName) -> Self {
  let Self {
@@ -224,6 +233,9 @@ impl Chunk {
  ///
  /// If none of the selected timelines exist in the [`Chunk`], the end result will be the same as the
  /// current chunk but without any timeline column.
+ ///
+ /// WARNING: the returned chunk has the same old [`crate::ChunkId`]! Change it with [`Self::with_id`].
+ #[must_use]
  #[inline]
  pub fn timelines_sliced(&self, timelines_to_keep: &IntSet<Timeline>) -> Self {
  let Self {
@@ -264,6 +276,9 @@ impl Chunk {
  ///
  /// If none of the `component_names` exist in the [`Chunk`], the end result will be the same as the
  /// current chunk but without any component column.
+ ///
+ /// WARNING: the returned chunk has the same old [`crate::ChunkId`]! Change it with [`Self::with_id`].
+ #[must_use]
  #[inline]
  pub fn components_sliced(&self, component_names: &IntSet<ComponentName>) -> Self {
  let Self {
@@ -306,6 +321,9 @@ impl Chunk {
  ///
  /// If `component_name` doesn't exist in this [`Chunk`], or if it is already dense, this method
  /// is a no-op.
+ ///
+ /// WARNING: the returned chunk has the same old [`crate::ChunkId`]! Change it with [`Self::with_id`].
+ #[must_use]
  #[inline]
  pub fn densified(&self, component_name_pov: ComponentName) -> Self {
  let Self {
@@ -400,6 +418,9 @@ impl Chunk {
  /// Empties the [`Chunk`] vertically.
  ///
  /// The result is a new [`Chunk`] with the same columns but zero rows.
+ ///
+ /// WARNING: the returned chunk has the same old [`crate::ChunkId`]! Change it with [`Self::with_id`].
+ #[must_use]
  #[inline]
  pub fn emptied(&self) -> Self {
  let Self {

crates/store/re_chunk_store/src/drop_time_range.rs

@@ -0,0 +1,114 @@
+use re_chunk::{ChunkId, Timeline};
+use re_log_types::ResolvedTimeRange;
+
+use crate::{ChunkStore, ChunkStoreEvent};
+
+impl ChunkStore {
+ /// Drop all events that are in the given range on the given timeline.
+ ///
+ /// Note that matching events will be dropped from all timelines they appear on.
+ ///
+ /// Static chunks are unaffected.
+ ///
+ /// Used to implement undo (erase the last event from the blueprint db).
+ pub fn drop_time_range(
+ &mut self,
+ timeline: &Timeline,
+ drop_range: ResolvedTimeRange,
+ ) -> Vec<ChunkStoreEvent> {
+ re_tracing::profile_function!();
+
+ if drop_range.max() < drop_range.min() {
+ return Default::default();
+ }
+
+ // Prepare the changes:
+
+ let mut chunk_ids_to_drop = vec![];
+ let mut new_chunks = vec![];
+
+ for (chunk_id, chunk) in &self.chunks_per_chunk_id {
+ let Some(time_column) = chunk.timelines().get(timeline) else {
+ // static chunk, or chunk that doesn't overlap this timeline
+ continue; // keep it
+ };
+
+ let chunk_range = time_column.time_range();
+
+ if drop_range.contains_range(chunk_range) {
+ // The whole chunk should be dropped!
+ chunk_ids_to_drop.push(*chunk_id);
+ } else if drop_range.intersects(chunk_range) {
+ let chunk = chunk.sorted_by_timeline_if_unsorted(timeline);
+
+ let num_rows = chunk.num_rows();
+
+ // Get the sorted times:
+ #[allow(clippy::unwrap_used)] // We already know the chunk has the timeline
+ let time_column = chunk.timelines().get(timeline).unwrap();
+ let times = time_column.times_raw();
+
+ let drop_range_min = drop_range.min().as_i64();
+ let drop_range_max = drop_range.max().as_i64();
+
+ let min_idx = times.partition_point(|&time| time < drop_range_min);
+ let max_idx = times.partition_point(|&time| time <= drop_range_max);
+
+ {
+ // Sanity check:
+ debug_assert!(min_idx <= max_idx);
+ debug_assert!(drop_range_min <= times[min_idx]);
+ if 0 < min_idx {
+ debug_assert!(times[min_idx - 1] < drop_range_min);
+ }
+ if max_idx < num_rows {
+ debug_assert!(drop_range_max < times[max_idx]);
+ if 0 < max_idx {
+ debug_assert!(times[max_idx - 1] <= drop_range_max);
+ }
+ }
+ }
+
+ if min_idx < max_idx {
+ chunk_ids_to_drop.push(*chunk_id);
+ if 0 < min_idx {
+ new_chunks.push(chunk.row_sliced(0, min_idx).with_id(ChunkId::new()));
+ }
+ if max_idx < num_rows {
+ new_chunks.push(
+ chunk
+ .row_sliced(max_idx, num_rows - max_idx)
+ .with_id(ChunkId::new()),
+ );
+ }
+ }
+ }
+ }
+
+ // ------------------
+ // Apply the changes:
+
+ let generation = self.generation();
+ let mut events: Vec<ChunkStoreEvent> = vec![];
+
+ for chunk_id in chunk_ids_to_drop {
+ for diff in self.remove_chunk(chunk_id) {
+ events.push(ChunkStoreEvent {
+ store_id: self.id.clone(),
+ store_generation: generation.clone(),
+ event_id: self
+ .event_id
+ .fetch_add(1, std::sync::atomic::Ordering::Relaxed),
+ diff,
+ });
+ }
+ }
+ for mut chunk in new_chunks {
+ chunk.sort_if_unsorted();
+ #[allow(clippy::unwrap_used)] // The chunk came from the store, so it should be fine
+ events.append(&mut self.insert_chunk(&chunk.into()).unwrap());
+ }
+
+ events
+ }
+}

crates/store/re_chunk_store/src/lib.rs

@@ -15,6 +15,7 @@
 //!
 
 mod dataframe;
+mod drop_time_range;
 mod events;
 mod gc;
 mod query;

crates/store/re_chunk_store/tests/drop_time_range.rs

@@ -0,0 +1,82 @@
+// https://github.com/rust-lang/rust-clippy/issues/10011
+#![cfg(test)]
+
+use std::sync::Arc;
+
+use re_chunk::{Chunk, RowId};
+use re_chunk_store::{ChunkStore, ChunkStoreConfig};
+use re_log_types::{example_components::MyColor, ResolvedTimeRange};
+use re_log_types::{EntityPath, TimePoint, Timeline};
+use re_types_core::Loggable as _;
+
+#[test]
+fn drop_time_range() -> anyhow::Result<()> {
+ re_log::setup_logging();
+
+ let entity_path = EntityPath::from("this/that");
+ let timeline = Timeline::new_sequence("timeline");
+ let data = MyColor::from_rgb(255, 0, 0);
+ let time_point_at = |time: i64| TimePoint::from([(timeline, time)]);
+
+ for config in [
+ ChunkStoreConfig::DEFAULT,
+ ChunkStoreConfig::COMPACTION_DISABLED,
+ ] {
+ let mut store = ChunkStore::new(
+ re_log_types::StoreId::random(re_log_types::StoreKind::Recording),
+ config,
+ );
+
+ let num_events = |store: &ChunkStore| {
+ store.num_temporal_events_for_component_on_timeline(
+ &timeline,
+ &entity_path,
+ MyColor::name(),
+ )
+ };
+
+ store.insert_chunk(&Arc::new(
+ Chunk::builder(entity_path.clone())
+ .with_component_batch(RowId::new(), time_point_at(0), &data)
+ .with_component_batch(RowId::new(), time_point_at(1), &data)
+ .with_component_batch(RowId::new(), time_point_at(2), &data)
+ .with_component_batch(RowId::new(), time_point_at(3), &data)
+ .build()?,
+ ))?;
+
+ store.insert_chunk(&Arc::new(
+ Chunk::builder(entity_path.clone())
+ .with_component_batch(RowId::new(), time_point_at(2), &data)
+ .with_component_batch(RowId::new(), time_point_at(3), &data)
+ .with_component_batch(RowId::new(), time_point_at(4), &data)
+ .with_component_batch(RowId::new(), time_point_at(5), &data)
+ .build()?,
+ ))?;
+
+ store.insert_chunk(&Arc::new(
+ Chunk::builder(entity_path.clone())
+ .with_component_batch(RowId::new(), time_point_at(4), &data)
+ .with_component_batch(RowId::new(), time_point_at(5), &data)
+ .with_component_batch(RowId::new(), time_point_at(6), &data)
+ .with_component_batch(RowId::new(), time_point_at(7), &data)
+ .build()?,
+ ))?;
+
+ assert_eq!(num_events(&store), 12);
+
+ // Drop nothing:
+ store.drop_time_range(&timeline, ResolvedTimeRange::new(10, 100));
+ store.drop_time_range(&timeline, ResolvedTimeRange::new(-100, -10));
+ assert_eq!(num_events(&store), 12);
+
+ // Drop stuff from the middle of the first chunk, and the start of the second:
+ store.drop_time_range(&timeline, ResolvedTimeRange::new(1, 2));
+ assert_eq!(num_events(&store), 9);
+
+ // Drop a bunch in the middle (including all of middle chunk):
+ store.drop_time_range(&timeline, ResolvedTimeRange::new(2, 5));
+ assert_eq!(num_events(&store), 3);
+ }
+
+ Ok(())
+}

crates/store/re_entity_db/src/entity_db.rs

@@ -446,6 +446,19 @@ impl EntityDb {
  store_events
  }
 
+ /// Drop all events in the given time range from the given timeline.
+ ///
+ /// Used to implement undo (erase the last event from the blueprint db).
+ pub fn drop_time_range(
+ &mut self,
+ timeline: &Timeline,
+ drop_range: ResolvedTimeRange,
+ ) -> Vec<ChunkStoreEvent> {
+ let store_events = self.data_store.drop_time_range(timeline, drop_range);
+ self.on_store_deletions(&store_events);
+ store_events
+ }
+
  /// Unconditionally drops all the data for a given [`EntityPath`] .
  ///
  /// This is _not_ recursive. Children of this entity will not be affected.

crates/store/re_log_types/src/resolved_time_range.rs

@@ -89,6 +89,12 @@ impl ResolvedTimeRange {
  self.min <= time && time <= self.max
  }
 
+ /// Does this range fully contain the other?
+ #[inline]
+ pub fn contains_range(&self, other: Self) -> bool {
+ self.min <= other.min && other.max <= self.max
+ }
+
  #[inline]
  pub fn intersects(&self, other: Self) -> bool {
  self.min <= other.max && self.max >= other.min

crates/store/re_log_types/src/time_point/time_int.rs

@@ -118,6 +118,7 @@ impl TimeInt {
 
  /// Always returns [`Self::STATIC`] for [`Self::STATIC`].
  #[inline]
+ #[must_use]
  pub fn inc(&self) -> Self {
  match self.0 {
  Some(t) => Self::new_temporal(t.get().saturating_add(1)),
@@ -127,6 +128,7 @@ impl TimeInt {
 
  /// Always returns [`Self::STATIC`] for [`Self::STATIC`].
  #[inline]
+ #[must_use]
  pub fn dec(&self) -> Self {
  match self.0 {
  Some(t) => Self::new_temporal(t.get().saturating_sub(1)),