diff --git a/crates/store/re_dataframe2/src/query.rs b/crates/store/re_dataframe2/src/query.rs index 397a46a324f5..ac1d21342ac2 100644 --- a/crates/store/re_dataframe2/src/query.rs +++ b/crates/store/re_dataframe2/src/query.rs @@ -3,6 +3,7 @@ use std::sync::{ OnceLock, }; +use ahash::HashSet; use arrow2::{ array::Array as ArrowArray, chunk::Chunk as ArrowChunk, datatypes::Schema as ArrowSchema, }; @@ -336,6 +337,31 @@ impl QueryHandle<'_> { &self.init().arrow_schema } + /// How many rows of data will be returned? + /// + /// The number of rows depends and only depends on the _view contents_. + /// The _selected contents_ has no influence on this value. + // + // TODO(cmc): implement this properly, cache the result, etc. + pub fn num_rows(&self) -> u64 { + let all_unique_timestamps: HashSet = self + .init() + .view_chunks + .iter() + .flat_map(|chunks| { + chunks.iter().filter_map(|(_cursor, chunk)| { + chunk + .timelines() + .get(&self.query.filtered_index) + .map(|time_column| time_column.times()) + }) + }) + .flatten() + .collect(); + + all_unique_timestamps.len() as _ + } + /// Returns the next row's worth of data. /// /// The returned vector of Arrow arrays strictly follows the schema specified by [`Self::schema`].