Use Range<u64> as argument for the RangeQueryHandle::get()

rerun-io · Sep 10, 2024 · c92b0bf · c92b0bf
1 parent 1923ba5
commit c92b0bf
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 26 deletions.
diff --git a/crates/store/re_dataframe/examples/range_paginated.rs b/crates/store/re_dataframe/examples/range_paginated.rs
@@ -61,21 +61,21 @@ fn main() -> anyhow::Result<()> {
             query_handle.num_rows()
         );
 
-        let (offset, len) = (0, 4);
-        println!("offset:{offset} len:{len}");
-        concat_and_print(query_handle.get(offset, len));
+        let row_range = 0..4;
+        println!("range: {row_range:?}");
+        concat_and_print(query_handle.get(row_range));
 
-        let (offset, len) = (2, 4);
-        println!("offset:{offset} len:{len}");
-        concat_and_print(query_handle.get(offset, len));
+        let row_range = 2..6;
+        println!("range: {row_range:?}");
+        concat_and_print(query_handle.get(row_range));
 
-        let (offset, len) = (10, 5);
-        println!("offset:{offset} len:{len}");
-        concat_and_print(query_handle.get(offset, len));
+        let row_range = 10..15;
+        println!("range: {row_range:?}");
+        concat_and_print(query_handle.get(row_range));
 
-        let (offset, len) = (0, 15);
-        println!("offset:{offset} len:{len}");
-        concat_and_print(query_handle.get(offset, len));
+        let row_range = 0..15;
+        println!("range: {row_range:?}");
+        concat_and_print(query_handle.get(row_range));
     }
 
     Ok(())

diff --git a/crates/store/re_dataframe/src/range.rs b/crates/store/re_dataframe/src/range.rs
@@ -1,3 +1,4 @@
+use std::ops::Range;
 use std::sync::{atomic::AtomicU64, OnceLock};
 
 use ahash::HashMap;
@@ -219,7 +220,9 @@ impl RangeQueryHandle<'_> {
     /// This is the most performant way to iterate over the dataset.
     //
     // TODO(cmc): This could be turned into an actual lazy iterator at some point.
-    pub fn get(&self, offset: u64, mut len: u64) -> Vec<RecordBatch> {
+    pub fn get(&self, row_range: Range<u64>) -> Vec<RecordBatch> {
+        let offset = row_range.start;
+        let mut len = row_range.end.saturating_sub(row_range.start);
         re_tracing::profile_function!(format!("get({offset}, {len}, {})", self.query));
 
         let state = self.init();
@@ -533,7 +536,7 @@ mod tests {
 
         // Paginated API
         {
-            let batch = handle.get(0, 0).pop().unwrap();
+            let batch = handle.get(0..0).pop().unwrap();
             // The output should be an empty recordbatch with the right schema and empty arrays.
             assert_eq!(0, batch.num_rows());
             assert!(
@@ -543,9 +546,9 @@ mod tests {
             assert!(itertools::izip!(handle.schema(), batch.data.iter())
                 .all(|(descr, array)| descr.datatype() == array.data_type()));
 
-            let _batch = handle.get(0, 1).pop().unwrap();
+            let _batch = handle.get(0..1).pop().unwrap();
 
-            let batch = handle.get(1, 1).pop();
+            let batch = handle.get(1..2).pop();
             assert!(batch.is_none());
         }
     }
@@ -633,7 +636,7 @@ mod tests {
 
         // Paginated API
         {
-            let batch = handle.get(0, 1).pop().unwrap();
+            let batch = handle.get(0..1).pop().unwrap();
             // The output should be an empty recordbatch with the right schema and empty arrays.
             assert_eq!(1, batch.num_rows());
             assert_eq!(
@@ -654,9 +657,9 @@ mod tests {
                     .all(|(descr, field)| descr.to_arrow_field() == *field)
             );
 
-            let _batch = handle.get(1, 1).pop().unwrap();
+            let _batch = handle.get(1..2).pop().unwrap();
 
-            let batch = handle.get(2, 1).pop();
+            let batch = handle.get(2..3).pop();
             assert!(batch.is_none());
         }
     }

diff --git a/crates/viewer/re_space_view_dataframe/src/dataframe_ui.rs b/crates/viewer/re_space_view_dataframe/src/dataframe_ui.rs
@@ -44,15 +44,15 @@ impl QueryHandle<'_> {
         }
     }
 
-    fn get(&self, start: u64, num_rows: u64) -> Vec<RecordBatch> {
+    fn get(&self, row_range: Range<u64>) -> Vec<RecordBatch> {
         match self {
             QueryHandle::LatestAt(query_handle) => {
                 // latest-at queries only have one row
-                debug_assert_eq!((start, num_rows), (0, 1));
+                debug_assert_eq!(row_range, 0..1);
 
                 vec![query_handle.get()]
             }
-            QueryHandle::Range(query_handle) => query_handle.get(start, num_rows),
+            QueryHandle::Range(query_handle) => query_handle.get(row_range),
         }
     }
 
@@ -178,10 +178,7 @@ impl<'a> egui_table::TableDelegate for DataframeTableDelegate<'a> {
 
         let data = RowsDisplayData::try_new(
             &info.visible_rows,
-            self.query_handle.get(
-                info.visible_rows.start,
-                info.visible_rows.end - info.visible_rows.start,
-            ),
+            self.query_handle.get(info.visible_rows.clone()),
             self.schema,
             &self.query_handle.timeline(),
         );