@@ -3,14 +3,19 @@ use std::{
33 sync:: atomic:: { AtomicU64 , Ordering } ,
44} ;
55
6- use arrow2:: array:: {
7- Array as ArrowArray , ListArray as ArrowListArray , PrimitiveArray as ArrowPrimitiveArray ,
8- StructArray as ArrowStructArray ,
6+ use arrow2:: {
7+ array:: {
8+ Array as ArrowArray , ListArray as ArrowListArray , PrimitiveArray as ArrowPrimitiveArray ,
9+ StructArray as ArrowStructArray ,
10+ } ,
11+ Either ,
912} ;
1013
1114use itertools:: { izip, Itertools } ;
1215use re_log_types:: { EntityPath , ResolvedTimeRange , Time , TimeInt , TimePoint , Timeline } ;
13- use re_types_core:: { ComponentName , Loggable , LoggableBatch , SerializationError , SizeBytes } ;
16+ use re_types_core:: {
17+ ComponentName , DeserializationError , Loggable , LoggableBatch , SerializationError , SizeBytes ,
18+ } ;
1419
1520use crate :: { ChunkId , RowId } ;
1621
@@ -26,8 +31,18 @@ pub enum ChunkError {
2631 #[ error( transparent) ]
2732 Arrow ( #[ from] arrow2:: error:: Error ) ,
2833
34+ #[ error( "{kind} index out of bounds: {index} (len={len})" ) ]
35+ IndexOutOfBounds {
36+ kind : String ,
37+ len : usize ,
38+ index : usize ,
39+ } ,
40+
2941 #[ error( transparent) ]
3042 Serialization ( #[ from] SerializationError ) ,
43+
44+ #[ error( transparent) ]
45+ Deserialization ( #[ from] DeserializationError ) ,
3146}
3247
3348pub type ChunkResult < T > = Result < T , ChunkError > ;
@@ -822,6 +837,32 @@ impl Chunk {
822837 . map ( |( & time, & counter) | RowId :: from_u128 ( ( time as u128 ) << 64 | ( counter as u128 ) ) )
823838 }
824839
840+ /// Returns an iterator over the [`RowId`]s of a [`Chunk`], for a given component.
841+ ///
842+ /// This is different than [`Self::row_ids`]: it will only yield `RowId`s for rows at which
843+ /// there is data for the specified `component_name`.
844+ #[ inline]
845+ pub fn component_row_ids (
846+ & self ,
847+ component_name : & ComponentName ,
848+ ) -> impl Iterator < Item = RowId > + ' _ {
849+ let Some ( list_array) = self . components . get ( component_name) else {
850+ return Either :: Left ( std:: iter:: empty ( ) ) ;
851+ } ;
852+
853+ let row_ids = self . row_ids ( ) ;
854+
855+ if let Some ( validity) = list_array. validity ( ) {
856+ Either :: Right ( Either :: Left (
857+ row_ids
858+ . enumerate ( )
859+ . filter_map ( |( i, o) | validity. get_bit ( i) . then_some ( o) ) ,
860+ ) )
861+ } else {
862+ Either :: Right ( Either :: Right ( row_ids) )
863+ }
864+ }
865+
825866 /// Returns the [`RowId`]-range covered by this [`Chunk`].
826867 ///
827868 /// `None` if the chunk `is_empty`.
0 commit comments