Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Find and annihilate slow boolean-component iterators #8541

Merged
merged 2 commits into from
Dec 19, 2024
Merged

Find and annihilate slow boolean-component iterators #8541

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c477579
introduce iter_bool
teh-cmc Dec 19, 2024
445fc7f
annihilate a bunch of component_slow
teh-cmc Dec 19, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 47 additions & 3 deletions crates/store/re_chunk/src/iter.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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<Item = Arrow2Bitmap> + '_ {
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::<Arrow2BooleanArray>()
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
Expand Down
14 changes: 13 additions & 1 deletion crates/viewer/re_view/src/results_ext.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<Item = ((TimeInt, RowId), Arrow2Bitmap)> + '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.
Expand Down
5 changes: 2 additions & 3 deletions crates/viewer/re_view_spatial/src/visualizers/arrows2d.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -244,8 +244,7 @@ impl VisualizerSystem for Arrows2DVisualizer {
all_labels.string(),
all_class_ids.primitive::<u16>(),
all_keypoint_ids.primitive::<u16>(),
// TODO(cmc): provide a `iter_bool`.
all_show_labels.component_slow::<ShowLabels>(),
all_show_labels.bool(),
)
.map(
|(
Expand All @@ -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),
}
},
);
Expand Down
4 changes: 2 additions & 2 deletions crates/viewer/re_view_spatial/src/visualizers/arrows3d.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -241,7 +241,7 @@ impl VisualizerSystem for Arrows3DVisualizer {
all_radii.primitive::<f32>(),
all_labels.string(),
all_class_ids.primitive::<u16>(),
all_show_labels.component_slow::<ShowLabels>(),
all_show_labels.bool(),
)
.map(
|(_index, vectors, origins, colors, radii, labels, class_ids, show_labels)| {
Expand All @@ -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),
}
},
);
Expand Down
4 changes: 2 additions & 2 deletions crates/viewer/re_view_spatial/src/visualizers/boxes2d.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -245,7 +245,7 @@ impl VisualizerSystem for Boxes2DVisualizer {
all_radii.primitive::<f32>(),
all_labels.string(),
all_class_ids.primitive::<u16>(),
all_show_labels.component_slow::<ShowLabels>(),
all_show_labels.bool(),
)
.map(
|(
Expand All @@ -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),
}
},
);
Expand Down
4 changes: 2 additions & 2 deletions crates/viewer/re_view_spatial/src/visualizers/boxes3d.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -183,7 +183,7 @@ impl VisualizerSystem for Boxes3DVisualizer {
all_radii.primitive::<f32>(),
all_labels.string(),
all_class_ids.primitive::<u16>(),
all_show_labels.component_slow::<ShowLabels>(),
all_show_labels.bool(),
)
.map(
|(_index, half_sizes, colors, radii, labels, class_ids, show_labels)| {
Expand All @@ -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),
}
},
);
Expand Down
4 changes: 2 additions & 2 deletions crates/viewer/re_view_spatial/src/visualizers/capsules3d.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -195,7 +195,7 @@ impl VisualizerSystem for Capsules3DVisualizer {
all_radii_indexed,
all_colors.primitive::<u32>(),
all_labels.string(),
all_show_labels.component_slow::<ShowLabels>(),
all_show_labels.bool(),
all_class_ids.primitive::<u16>(),
)
.map(
Expand All @@ -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),
}
},
);
Expand Down
4 changes: 2 additions & 2 deletions crates/viewer/re_view_spatial/src/visualizers/ellipsoids.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -181,7 +181,7 @@ impl VisualizerSystem for Ellipsoids3DVisualizer {
all_fill_modes.primitive::<u8>(),
all_labels.string(),
all_class_ids.primitive::<u16>(),
all_show_labels.component_slow::<ShowLabels>(),
all_show_labels.bool(),
)
.map(
|(
Expand Down Expand Up @@ -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),
}
},
);
Expand Down
4 changes: 2 additions & 2 deletions crates/viewer/re_view_spatial/src/visualizers/lines2d.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -232,7 +232,7 @@ impl VisualizerSystem for Lines2DVisualizer {
all_radii.primitive::<f32>(),
all_labels.string(),
all_class_ids.primitive::<u16>(),
all_show_labels.component_slow::<ShowLabels>(),
all_show_labels.bool(),
)
.map(
|(_index, strips, colors, radii, labels, class_ids, show_labels)| {
Expand All @@ -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),
}
},
);
Expand Down
4 changes: 2 additions & 2 deletions crates/viewer/re_view_spatial/src/visualizers/lines3d.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -236,7 +236,7 @@ impl VisualizerSystem for Lines3DVisualizer {
all_radii.primitive::<f32>(),
all_labels.string(),
all_class_ids.primitive::<u16>(),
all_show_labels.component_slow::<ShowLabels>(),
all_show_labels.bool(),
)
.map(
|(_index, strips, colors, radii, labels, class_ids, show_labels)| {
Expand All @@ -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),
}
},
);
Expand Down
4 changes: 2 additions & 2 deletions crates/viewer/re_view_spatial/src/visualizers/points2d.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -251,7 +251,7 @@ impl VisualizerSystem for Points2DVisualizer {
all_labels.string(),
all_class_ids.primitive::<u16>(),
all_keypoint_ids.primitive::<u16>(),
all_show_labels.component_slow::<ShowLabels>(),
all_show_labels.bool(),
)
.map(
|(
Expand All @@ -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),
}
},
);
Expand Down
4 changes: 2 additions & 2 deletions crates/viewer/re_view_spatial/src/visualizers/points3d.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -241,7 +241,7 @@ impl VisualizerSystem for Points3DVisualizer {
all_labels.string(),
all_class_ids.primitive::<u16>(),
all_keypoint_ids.primitive::<u16>(),
all_show_labels.component_slow::<ShowLabels>(),
all_show_labels.bool(),
)
.map(
|(
Expand All @@ -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),
}
},
);
Expand Down
Loading