From c47757998f03a95cb1f72a7193edf17ef2b50fac Mon Sep 17 00:00:00 2001 From: Clement Rey Date: Thu, 19 Dec 2024 12:52:28 +0100 Subject: [PATCH 1/2] introduce iter_bool --- crates/store/re_chunk/src/iter.rs | 50 ++++++++++++++++++++++-- crates/viewer/re_view/src/results_ext.rs | 14 ++++++- 2 files changed, 60 insertions(+), 4 deletions(-) diff --git a/crates/store/re_chunk/src/iter.rs b/crates/store/re_chunk/src/iter.rs index fd212996822c..c417e814c6ad 100644 --- a/crates/store/re_chunk/src/iter.rs +++ b/crates/store/re_chunk/src/iter.rs @@ -2,10 +2,11 @@ use std::sync::Arc; use arrow2::{ array::{ - Array as Arrow2Array, FixedSizeListArray as Arrow2FixedSizeListArray, - ListArray as Arrow2ListArray, PrimitiveArray as Arrow2PrimitiveArray, - Utf8Array as Arrow2Utf8Array, + Array as Arrow2Array, BooleanArray as Arrow2BooleanArray, + FixedSizeListArray as Arrow2FixedSizeListArray, ListArray as Arrow2ListArray, + PrimitiveArray as Arrow2PrimitiveArray, Utf8Array as Arrow2Utf8Array, }, + bitmap::Bitmap as Arrow2Bitmap, Either, }; use itertools::{izip, Itertools}; @@ -264,6 +265,49 @@ impl Chunk { ) } + /// Returns an iterator over the raw boolean values of a [`Chunk`], for a given component. + /// + /// This is a very fast path: the entire column will be downcasted at once, and then every + /// component batch will be a slice reference into that global slice. + /// Use this when working with simple arrow datatypes and performance matters. + /// + /// See also: + /// * [`Self::iter_primitive_array`] + /// * [`Self::iter_primitive_array_list`] + /// * [`Self::iter_string`] + /// * [`Self::iter_buffer`]. + /// * [`Self::iter_component_arrays`]. + /// * [`Self::iter_component`]. + #[inline] + pub fn iter_bool( + &self, + component_name: &ComponentName, + ) -> impl Iterator + '_ { + let Some(list_array) = self.get_first_component(component_name) else { + return Either::Left(std::iter::empty()); + }; + + let Some(values) = list_array + .values() + .as_any() + .downcast_ref::() + else { + if cfg!(debug_assertions) { + panic!("downcast failed for {component_name}, data discarded"); + } else { + re_log::error_once!("downcast failed for {component_name}, data discarded"); + } + return Either::Left(std::iter::empty()); + }; + let values = values.values().clone(); + + // NOTE: No need for validity checks here, `iter_offsets` already takes care of that. + Either::Right( + self.iter_component_offsets(component_name) + .map(move |(idx, len)| values.clone().sliced(idx, len)), + ) + } + /// Returns an iterator over the raw primitive arrays of a [`Chunk`], for a given component. /// /// This is a very fast path: the entire column will be downcasted at once, and then every diff --git a/crates/viewer/re_view/src/results_ext.rs b/crates/viewer/re_view/src/results_ext.rs index ed5a05b659d5..8e4118a48519 100644 --- a/crates/viewer/re_view/src/results_ext.rs +++ b/crates/viewer/re_view/src/results_ext.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use itertools::Itertools as _; use re_chunk_store::{Chunk, LatestAtQuery, RangeQuery, UnitChunkShared}; -use re_log_types::external::arrow2::array::Array as Arrow2Array; +use re_log_types::external::arrow2::{array::Array as Arrow2Array, bitmap::Bitmap as Arrow2Bitmap}; use re_log_types::hash::Hash64; use re_query::{LatestAtResults, RangeResults}; use re_types_core::ComponentName; @@ -437,6 +437,18 @@ impl<'a> HybridResultsChunkIter<'a> { }) } + /// Iterate as indexed booleans. + /// + /// See [`Chunk::iter_bool`] for more information. + pub fn bool(&'a self) -> impl Iterator + 'a { + self.chunks.iter().flat_map(move |chunk| { + itertools::izip!( + chunk.iter_component_indices(&self.timeline, &self.component_name), + chunk.iter_bool(&self.component_name) + ) + }) + } + /// Iterate as indexed primitives. /// /// See [`Chunk::iter_primitive`] for more information. From 445fc7ff7fd3636e43cfd52ac4481864a877325e Mon Sep 17 00:00:00 2001 From: Clement Rey Date: Thu, 19 Dec 2024 12:52:35 +0100 Subject: [PATCH 2/2] annihilate a bunch of component_slow --- crates/viewer/re_view_spatial/src/visualizers/arrows2d.rs | 5 ++--- crates/viewer/re_view_spatial/src/visualizers/arrows3d.rs | 4 ++-- crates/viewer/re_view_spatial/src/visualizers/boxes2d.rs | 4 ++-- crates/viewer/re_view_spatial/src/visualizers/boxes3d.rs | 4 ++-- crates/viewer/re_view_spatial/src/visualizers/capsules3d.rs | 4 ++-- crates/viewer/re_view_spatial/src/visualizers/ellipsoids.rs | 4 ++-- crates/viewer/re_view_spatial/src/visualizers/lines2d.rs | 4 ++-- crates/viewer/re_view_spatial/src/visualizers/lines3d.rs | 4 ++-- crates/viewer/re_view_spatial/src/visualizers/points2d.rs | 4 ++-- crates/viewer/re_view_spatial/src/visualizers/points3d.rs | 4 ++-- 10 files changed, 20 insertions(+), 21 deletions(-) diff --git a/crates/viewer/re_view_spatial/src/visualizers/arrows2d.rs b/crates/viewer/re_view_spatial/src/visualizers/arrows2d.rs index c9d062306b2c..93b2e76457a6 100644 --- a/crates/viewer/re_view_spatial/src/visualizers/arrows2d.rs +++ b/crates/viewer/re_view_spatial/src/visualizers/arrows2d.rs @@ -244,8 +244,7 @@ impl VisualizerSystem for Arrows2DVisualizer { all_labels.string(), all_class_ids.primitive::(), all_keypoint_ids.primitive::(), - // TODO(cmc): provide a `iter_bool`. - all_show_labels.component_slow::(), + all_show_labels.bool(), ) .map( |( @@ -269,7 +268,7 @@ impl VisualizerSystem for Arrows2DVisualizer { .map_or(&[], |class_ids| bytemuck::cast_slice(class_ids)), keypoint_ids: keypoint_ids .map_or(&[], |keypoint_ids| bytemuck::cast_slice(keypoint_ids)), - show_labels: show_labels.unwrap_or_default().first().copied(), + show_labels: show_labels.unwrap_or_default().get(0).map(Into::into), } }, ); diff --git a/crates/viewer/re_view_spatial/src/visualizers/arrows3d.rs b/crates/viewer/re_view_spatial/src/visualizers/arrows3d.rs index 401afd9ec21f..2e02f969f53d 100644 --- a/crates/viewer/re_view_spatial/src/visualizers/arrows3d.rs +++ b/crates/viewer/re_view_spatial/src/visualizers/arrows3d.rs @@ -241,7 +241,7 @@ impl VisualizerSystem for Arrows3DVisualizer { all_radii.primitive::(), all_labels.string(), all_class_ids.primitive::(), - all_show_labels.component_slow::(), + all_show_labels.bool(), ) .map( |(_index, vectors, origins, colors, radii, labels, class_ids, show_labels)| { @@ -253,7 +253,7 @@ impl VisualizerSystem for Arrows3DVisualizer { labels: labels.unwrap_or_default(), class_ids: class_ids .map_or(&[], |class_ids| bytemuck::cast_slice(class_ids)), - show_labels: show_labels.unwrap_or_default().first().copied(), + show_labels: show_labels.unwrap_or_default().get(0).map(Into::into), } }, ); diff --git a/crates/viewer/re_view_spatial/src/visualizers/boxes2d.rs b/crates/viewer/re_view_spatial/src/visualizers/boxes2d.rs index bd451dd39f72..2a236f2cd3e6 100644 --- a/crates/viewer/re_view_spatial/src/visualizers/boxes2d.rs +++ b/crates/viewer/re_view_spatial/src/visualizers/boxes2d.rs @@ -245,7 +245,7 @@ impl VisualizerSystem for Boxes2DVisualizer { all_radii.primitive::(), all_labels.string(), all_class_ids.primitive::(), - all_show_labels.component_slow::(), + all_show_labels.bool(), ) .map( |( @@ -266,7 +266,7 @@ impl VisualizerSystem for Boxes2DVisualizer { labels: labels.unwrap_or_default(), class_ids: class_ids .map_or(&[], |class_ids| bytemuck::cast_slice(class_ids)), - show_labels: show_labels.unwrap_or_default().first().copied(), + show_labels: show_labels.unwrap_or_default().get(0).map(Into::into), } }, ); diff --git a/crates/viewer/re_view_spatial/src/visualizers/boxes3d.rs b/crates/viewer/re_view_spatial/src/visualizers/boxes3d.rs index 9853118eab2b..2d48e092ed32 100644 --- a/crates/viewer/re_view_spatial/src/visualizers/boxes3d.rs +++ b/crates/viewer/re_view_spatial/src/visualizers/boxes3d.rs @@ -183,7 +183,7 @@ impl VisualizerSystem for Boxes3DVisualizer { all_radii.primitive::(), all_labels.string(), all_class_ids.primitive::(), - all_show_labels.component_slow::(), + all_show_labels.bool(), ) .map( |(_index, half_sizes, colors, radii, labels, class_ids, show_labels)| { @@ -196,7 +196,7 @@ impl VisualizerSystem for Boxes3DVisualizer { labels: labels.unwrap_or_default(), class_ids: class_ids .map_or(&[], |class_ids| bytemuck::cast_slice(class_ids)), - show_labels: show_labels.unwrap_or_default().first().copied(), + show_labels: show_labels.unwrap_or_default().get(0).map(Into::into), } }, ); diff --git a/crates/viewer/re_view_spatial/src/visualizers/capsules3d.rs b/crates/viewer/re_view_spatial/src/visualizers/capsules3d.rs index 70705080e575..a2e2c5f1b41c 100644 --- a/crates/viewer/re_view_spatial/src/visualizers/capsules3d.rs +++ b/crates/viewer/re_view_spatial/src/visualizers/capsules3d.rs @@ -195,7 +195,7 @@ impl VisualizerSystem for Capsules3DVisualizer { all_radii_indexed, all_colors.primitive::(), all_labels.string(), - all_show_labels.component_slow::(), + all_show_labels.bool(), all_class_ids.primitive::(), ) .map( @@ -207,7 +207,7 @@ impl VisualizerSystem for Capsules3DVisualizer { labels: labels.unwrap_or_default(), class_ids: class_ids .map_or(&[], |class_ids| bytemuck::cast_slice(class_ids)), - show_labels: show_labels.unwrap_or_default().first().copied(), + show_labels: show_labels.unwrap_or_default().get(0).map(Into::into), } }, ); diff --git a/crates/viewer/re_view_spatial/src/visualizers/ellipsoids.rs b/crates/viewer/re_view_spatial/src/visualizers/ellipsoids.rs index fd52a0269b02..fa4f55c57f6c 100644 --- a/crates/viewer/re_view_spatial/src/visualizers/ellipsoids.rs +++ b/crates/viewer/re_view_spatial/src/visualizers/ellipsoids.rs @@ -181,7 +181,7 @@ impl VisualizerSystem for Ellipsoids3DVisualizer { all_fill_modes.primitive::(), all_labels.string(), all_class_ids.primitive::(), - all_show_labels.component_slow::(), + all_show_labels.bool(), ) .map( |( @@ -209,7 +209,7 @@ impl VisualizerSystem for Ellipsoids3DVisualizer { labels: labels.unwrap_or_default(), class_ids: class_ids .map_or(&[], |class_ids| bytemuck::cast_slice(class_ids)), - show_labels: show_labels.unwrap_or_default().first().copied(), + show_labels: show_labels.unwrap_or_default().get(0).map(Into::into), } }, ); diff --git a/crates/viewer/re_view_spatial/src/visualizers/lines2d.rs b/crates/viewer/re_view_spatial/src/visualizers/lines2d.rs index 009b33a3f5dc..d5b5db1f7cae 100644 --- a/crates/viewer/re_view_spatial/src/visualizers/lines2d.rs +++ b/crates/viewer/re_view_spatial/src/visualizers/lines2d.rs @@ -232,7 +232,7 @@ impl VisualizerSystem for Lines2DVisualizer { all_radii.primitive::(), all_labels.string(), all_class_ids.primitive::(), - all_show_labels.component_slow::(), + all_show_labels.bool(), ) .map( |(_index, strips, colors, radii, labels, class_ids, show_labels)| { @@ -243,7 +243,7 @@ impl VisualizerSystem for Lines2DVisualizer { labels: labels.unwrap_or_default(), class_ids: class_ids .map_or(&[], |class_ids| bytemuck::cast_slice(class_ids)), - show_labels: show_labels.unwrap_or_default().first().copied(), + show_labels: show_labels.unwrap_or_default().get(0).map(Into::into), } }, ); diff --git a/crates/viewer/re_view_spatial/src/visualizers/lines3d.rs b/crates/viewer/re_view_spatial/src/visualizers/lines3d.rs index 31d33029c371..9ef6d4511fb2 100644 --- a/crates/viewer/re_view_spatial/src/visualizers/lines3d.rs +++ b/crates/viewer/re_view_spatial/src/visualizers/lines3d.rs @@ -236,7 +236,7 @@ impl VisualizerSystem for Lines3DVisualizer { all_radii.primitive::(), all_labels.string(), all_class_ids.primitive::(), - all_show_labels.component_slow::(), + all_show_labels.bool(), ) .map( |(_index, strips, colors, radii, labels, class_ids, show_labels)| { @@ -247,7 +247,7 @@ impl VisualizerSystem for Lines3DVisualizer { labels: labels.unwrap_or_default(), class_ids: class_ids .map_or(&[], |class_ids| bytemuck::cast_slice(class_ids)), - show_labels: show_labels.unwrap_or_default().first().copied(), + show_labels: show_labels.unwrap_or_default().get(0).map(Into::into), } }, ); diff --git a/crates/viewer/re_view_spatial/src/visualizers/points2d.rs b/crates/viewer/re_view_spatial/src/visualizers/points2d.rs index 8929bec8da31..2b54e568c1d7 100644 --- a/crates/viewer/re_view_spatial/src/visualizers/points2d.rs +++ b/crates/viewer/re_view_spatial/src/visualizers/points2d.rs @@ -251,7 +251,7 @@ impl VisualizerSystem for Points2DVisualizer { all_labels.string(), all_class_ids.primitive::(), all_keypoint_ids.primitive::(), - all_show_labels.component_slow::(), + all_show_labels.bool(), ) .map( |( @@ -273,7 +273,7 @@ impl VisualizerSystem for Points2DVisualizer { .map_or(&[], |class_ids| bytemuck::cast_slice(class_ids)), keypoint_ids: keypoint_ids .map_or(&[], |keypoint_ids| bytemuck::cast_slice(keypoint_ids)), - show_labels: show_labels.unwrap_or_default().first().copied(), + show_labels: show_labels.unwrap_or_default().get(0).map(Into::into), } }, ); diff --git a/crates/viewer/re_view_spatial/src/visualizers/points3d.rs b/crates/viewer/re_view_spatial/src/visualizers/points3d.rs index 57abfef3309a..7928fea07477 100644 --- a/crates/viewer/re_view_spatial/src/visualizers/points3d.rs +++ b/crates/viewer/re_view_spatial/src/visualizers/points3d.rs @@ -241,7 +241,7 @@ impl VisualizerSystem for Points3DVisualizer { all_labels.string(), all_class_ids.primitive::(), all_keypoint_ids.primitive::(), - all_show_labels.component_slow::(), + all_show_labels.bool(), ) .map( |( @@ -263,7 +263,7 @@ impl VisualizerSystem for Points3DVisualizer { .map_or(&[], |class_ids| bytemuck::cast_slice(class_ids)), keypoint_ids: keypoint_ids .map_or(&[], |keypoint_ids| bytemuck::cast_slice(keypoint_ids)), - show_labels: show_labels.unwrap_or_default().first().copied(), + show_labels: show_labels.unwrap_or_default().get(0).map(Into::into), } }, );