From 8edab0949cec2749217c83454f75e9928bb41f7a Mon Sep 17 00:00:00 2001
From: Clement Rey <cr.rey.clement@gmail.com>
Date: Thu, 30 Nov 2023 13:11:15 +0100
Subject: [PATCH] retain() rather than remove()s -- doesnt change anything but
 nicer

---
 crates/re_arrow_store/benches/gc.rs   |  2 +-
 crates/re_arrow_store/src/store_gc.rs | 10 ++++------
 2 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/crates/re_arrow_store/benches/gc.rs b/crates/re_arrow_store/benches/gc.rs
index c0ab4b6a4bdf..81951f1776b2 100644
--- a/crates/re_arrow_store/benches/gc.rs
+++ b/crates/re_arrow_store/benches/gc.rs
@@ -77,6 +77,7 @@ fn plotting_dashboard(c: &mut Criterion) {
         enable_batching: false,
     };
 
+    // NOTE: insert in multiple timelines to more closely match real world scenarios.
     let mut timegen = |i| {
         [
             build_log_time(Time::from_seconds_since_epoch(i as _)),
@@ -275,7 +276,6 @@ where
         (0..NUM_ROWS_PER_ENTITY_PATH).map(move |i| {
             DataRow::from_component_batches(
                 RowId::random(),
-                // NOTE: insert in multiple timelines to more closely match real world scenarios.
                 timegen(i),
                 entity_path.clone(),
                 datagen(i)
diff --git a/crates/re_arrow_store/src/store_gc.rs b/crates/re_arrow_store/src/store_gc.rs
index 155d0a50f578..148479d976ac 100644
--- a/crates/re_arrow_store/src/store_gc.rs
+++ b/crates/re_arrow_store/src/store_gc.rs
@@ -662,7 +662,7 @@ impl IndexedTable {
         let mut dropped_num_bytes = 0u64;
         let mut dropped_num_rows = 0u64;
 
-        let mut dropped_bucket_times = Vec::new();
+        let mut dropped_bucket_times = HashSet::default();
 
         // TODO(cmc): scaling linearly with the number of buckets could be improved, although this
         // is quite fast in practice because of the early check.
@@ -681,7 +681,7 @@ impl IndexedTable {
                 ..
             } = std::mem::take(inner);
 
-            dropped_bucket_times.push(*bucket_time);
+            dropped_bucket_times.insert(*bucket_time);
 
             while let Some(row_id) = col_row_id.pop_front() {
                 let mut diff = StoreDiff::deletion(row_id, ent_path.clone());
@@ -713,10 +713,8 @@ impl IndexedTable {
             dropped_num_rows += col_time.len() as u64;
         }
 
-        for bucket_time in dropped_bucket_times {
-            let previous = self.buckets.remove(&bucket_time);
-            debug_assert!(previous.is_some());
-        }
+        self.buckets
+            .retain(|bucket_time, _| !dropped_bucket_times.contains(bucket_time));
 
         if self.buckets.is_empty() {
             let Self {