Skip to content

Commit

Permalink
Find and annihilate slow component iterators (#8540)
Browse files Browse the repository at this point in the history
* `HybridResultsChunkIter::component` is now
`HybridResultsChunkIter::component_slow`, in order to scare people off.
* I have removed *a lot* of those, and will remove a few more in follow
up PRs, but I need to add more iteration tools first.
  • Loading branch information
teh-cmc authored Dec 19, 2024
1 parent 5a54e08 commit 5acbbab
Show file tree
Hide file tree
Showing 23 changed files with 121 additions and 49 deletions.
4 changes: 4 additions & 0 deletions crates/store/re_chunk/src/iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -686,6 +686,10 @@ impl Chunk {
/// Use this when working with complex arrow datatypes and performance matters (e.g. ranging
/// through enum types across many timestamps).
///
/// TODO(#5305): Note that, while this is much faster than deserializing each row individually,
/// this still uses the old codegen'd deserialization path, which does some very unidiomatic Arrow
/// things, and is therefore very slow at the moment. Avoid this on performance critical paths.
///
/// See also:
/// * [`Self::iter_primitive`]
/// * [`Self::iter_primitive_array`]
Expand Down
31 changes: 31 additions & 0 deletions crates/store/re_types/src/components/colormap_ext.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
use super::Colormap;

impl Colormap {
/// Instantiate a new [`Colormap`] from a u8 value.
///
/// Returns `None` if the value doesn't match any of the enum's arms.
pub fn from_u8(value: u8) -> Option<Self> {
// NOTE: This code will be optimized out, it's only here to make sure this method fails to
// compile if the enum is modified.
match Self::default() {
Self::Grayscale
| Self::Inferno
| Self::Magma
| Self::Plasma
| Self::Turbo
| Self::Viridis
| Self::CyanToYellow => {}
}

match value {
v if v == Self::Grayscale as u8 => Some(Self::Grayscale),
v if v == Self::Inferno as u8 => Some(Self::Inferno),
v if v == Self::Magma as u8 => Some(Self::Magma),
v if v == Self::Plasma as u8 => Some(Self::Plasma),
v if v == Self::Turbo as u8 => Some(Self::Turbo),
v if v == Self::Viridis as u8 => Some(Self::Viridis),
v if v == Self::CyanToYellow as u8 => Some(Self::CyanToYellow),
_ => None,
}
}
}
21 changes: 21 additions & 0 deletions crates/store/re_types/src/components/fill_mode_ext.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
use super::FillMode;

impl FillMode {
/// Instantiate a new [`FillMode`] from a u8 value.
///
/// Returns `None` if the value doesn't match any of the enum's arms.
pub fn from_u8(value: u8) -> Option<Self> {
// NOTE: This code will be optimized out, it's only here to make sure this method fails to
// compile if the enum is modified.
match Self::default() {
Self::MajorWireframe | Self::DenseWireframe | Self::Solid => {}
}

match value {
v if v == Self::MajorWireframe as u8 => Some(Self::MajorWireframe),
v if v == Self::DenseWireframe as u8 => Some(Self::DenseWireframe),
v if v == Self::Solid as u8 => Some(Self::Solid),
_ => None,
}
}
}
2 changes: 2 additions & 0 deletions crates/store/re_types/src/components/mod.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

6 changes: 5 additions & 1 deletion crates/viewer/re_view/src/results_ext.rs
Original file line number Diff line number Diff line change
Expand Up @@ -421,8 +421,12 @@ pub struct HybridResultsChunkIter<'a> {
impl<'a> HybridResultsChunkIter<'a> {
/// Iterate as indexed deserialized batches.
///
/// TODO(#5305): Note that this uses the old codegen'd deserialization path, which does some
/// very unidiomatic Arrow things, and is therefore very slow at the moment. Avoid this on
/// performance critical paths.
///
/// See [`Chunk::iter_component`] for more information.
pub fn component<C: re_types_core::Component>(
pub fn component_slow<C: re_types_core::Component>(
&'a self,
) -> impl Iterator<Item = ((TimeInt, RowId), ChunkComponentIterItem<C>)> + 'a {
self.chunks.iter().flat_map(move |chunk| {
Expand Down
3 changes: 2 additions & 1 deletion crates/viewer/re_view_graph/src/visualizers/edges.rs
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,8 @@ impl VisualizerSystem for EdgesVisualizer {
let all_indexed_edges = results.iter_as(query.timeline, components::GraphEdge::name());
let graph_type = results.get_mono_with_fallback::<components::GraphType>();

for (_index, edges) in all_indexed_edges.component::<GraphEdge>() {
// TODO(cmc): Provide a `iter_struct`.
for (_index, edges) in all_indexed_edges.component_slow::<GraphEdge>() {
let edges = edges
.iter()
.enumerate()
Expand Down
9 changes: 5 additions & 4 deletions crates/viewer/re_view_graph/src/visualizers/nodes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -88,8 +88,9 @@ impl VisualizerSystem for NodeVisualizer {
.map_or(true, bool::from);

let data = range_zip_1x4(
all_indexed_nodes.component::<components::GraphNode>(),
all_colors.component::<components::Color>(),
// TODO(cmc): Provide a `iter_struct`.
all_indexed_nodes.component_slow::<components::GraphNode>(),
all_colors.primitive::<u32>(),
all_positions.primitive_array::<2, f32>(),
all_labels.string(),
all_radii.primitive::<f32>(),
Expand All @@ -100,7 +101,7 @@ impl VisualizerSystem for NodeVisualizer {
nodes.iter(),
(0..).map(Instance::from),
colors.unwrap_or_default().iter().map(Option::Some),
Option::<&Color>::default,
Option::<&u32>::default,
positions
.unwrap_or_default()
.iter()
Expand All @@ -113,7 +114,7 @@ impl VisualizerSystem for NodeVisualizer {
Option::<f32>::default,
)
.map(|(node, instance, color, position, label, radius)| {
let color = color.map(|&c| egui::Color32::from(c));
let color = color.map(|&c| egui::Color32::from(Color::new(c)));
let label = match (label, show_label) {
(Some(label), true) => Label::Text {
text: label.clone(),
Expand Down
23 changes: 11 additions & 12 deletions crates/viewer/re_view_map/src/visualizers/geo_line_strings.rs
Original file line number Diff line number Diff line change
Expand Up @@ -66,20 +66,20 @@ impl VisualizerSystem for GeoLineStringsVisualizer {

// iterate over each chunk and find all relevant component slices
for (_index, lines, colors, radii) in re_query::range_zip_1x2(
all_lines.component::<GeoLineString>(),
all_colors.component::<Color>(),
all_radii.component::<Radius>(),
all_lines.primitive_array_list::<2, f64>(),
all_colors.primitive::<u32>(),
all_radii.primitive::<f32>(),
) {
// required component
let lines = lines.as_slice();

// optional components
let colors = colors.as_ref().map(|c| c.as_slice()).unwrap_or(&[]);
let radii = radii.as_ref().map(|r| r.as_slice()).unwrap_or(&[]);
let colors = colors.unwrap_or(&[]);
let radii = radii.unwrap_or(&[]);

// optional components values to be used for instance clamping semantics
let last_color = colors.last().copied().unwrap_or(fallback_color);
let last_radii = radii.last().copied().unwrap_or(fallback_radius);
let last_color = colors.last().copied().unwrap_or(fallback_color.0 .0);
let last_radii = radii.last().copied().unwrap_or(fallback_radius.0 .0);

// iterate over all instances
for (instance_index, (line, color, radius)) in itertools::izip!(
Expand All @@ -90,13 +90,12 @@ impl VisualizerSystem for GeoLineStringsVisualizer {
.enumerate()
{
batch_data.lines.push(
line.0
.iter()
.map(|pos| walkers::Position::from_lat_lon(pos.x(), pos.y()))
line.iter()
.map(|pos| walkers::Position::from_lat_lon(pos[0], pos[1]))
.collect(),
);
batch_data.radii.push(*radius);
batch_data.colors.push(color.0.into());
batch_data.radii.push(Radius((*radius).into()));
batch_data.colors.push(Color::new(*color).into());
batch_data
.instance_id
.push(re_renderer::PickingLayerInstanceId(instance_index as _));
Expand Down
18 changes: 8 additions & 10 deletions crates/viewer/re_view_map/src/visualizers/geo_points.rs
Original file line number Diff line number Diff line change
Expand Up @@ -69,13 +69,12 @@ impl VisualizerSystem for GeoPointsVisualizer {

// iterate over each chunk and find all relevant component slices
for (_index, positions, colors, radii, class_ids) in re_query::range_zip_1x3(
all_positions.component::<LatLon>(),
all_positions.primitive_array::<2, f64>(),
all_colors.primitive::<u32>(),
all_radii.component::<Radius>(),
all_radii.primitive::<f32>(),
all_class_ids.primitive::<u16>(),
) {
// required component
let positions = positions.as_slice();
let num_instances = positions.len();

// Resolve annotation info (if needed).
Expand All @@ -94,10 +93,10 @@ impl VisualizerSystem for GeoPointsVisualizer {
&annotation_infos,
colors.map_or(&[], |colors| bytemuck::cast_slice(colors)),
);
let radii = radii.as_ref().map(|r| r.as_slice()).unwrap_or(&[]);
let radii = radii.unwrap_or(&[]);

// optional components values to be used for instance clamping semantics
let last_radii = radii.last().copied().unwrap_or(fallback_radius);
let last_radii = radii.last().copied().unwrap_or(fallback_radius.0 .0);

// iterate over all instances
for (instance_index, (position, color, radius)) in itertools::izip!(
Expand All @@ -107,11 +106,10 @@ impl VisualizerSystem for GeoPointsVisualizer {
)
.enumerate()
{
batch_data.positions.push(walkers::Position::from_lat_lon(
position.latitude(),
position.longitude(),
));
batch_data.radii.push(*radius);
batch_data
.positions
.push(walkers::Position::from_lat_lon(position[0], position[1]));
batch_data.radii.push(Radius((*radius).into()));
batch_data.colors.push(*color);
batch_data
.instance_id
Expand Down
3 changes: 2 additions & 1 deletion crates/viewer/re_view_spatial/src/visualizers/arrows2d.rs
Original file line number Diff line number Diff line change
Expand Up @@ -244,7 +244,8 @@ impl VisualizerSystem for Arrows2DVisualizer {
all_labels.string(),
all_class_ids.primitive::<u16>(),
all_keypoint_ids.primitive::<u16>(),
all_show_labels.component::<ShowLabels>(),
// TODO(cmc): provide a `iter_bool`.
all_show_labels.component_slow::<ShowLabels>(),
)
.map(
|(
Expand Down
2 changes: 1 addition & 1 deletion crates/viewer/re_view_spatial/src/visualizers/arrows3d.rs
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::<ShowLabels>(),
all_show_labels.component_slow::<ShowLabels>(),
)
.map(
|(_index, vectors, origins, colors, radii, labels, class_ids, show_labels)| {
Expand Down
2 changes: 1 addition & 1 deletion crates/viewer/re_view_spatial/src/visualizers/boxes2d.rs
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::<ShowLabels>(),
all_show_labels.component_slow::<ShowLabels>(),
)
.map(
|(
Expand Down
8 changes: 5 additions & 3 deletions crates/viewer/re_view_spatial/src/visualizers/boxes3d.rs
Original file line number Diff line number Diff line change
Expand Up @@ -159,9 +159,11 @@ impl VisualizerSystem for Boxes3DVisualizer {
let all_fill_modes = results.iter_as(timeline, FillMode::name());
// fill mode is currently a non-repeated component
let fill_mode: FillMode = all_fill_modes
.component::<FillMode>()
.primitive::<u8>()
.next()
.and_then(|(_, fill_modes)| fill_modes.as_slice().first().copied())
.and_then(|(_, fill_modes)| {
fill_modes.first().copied().and_then(FillMode::from_u8)
})
.unwrap_or_default();

match fill_mode {
Expand All @@ -181,7 +183,7 @@ impl VisualizerSystem for Boxes3DVisualizer {
all_radii.primitive::<f32>(),
all_labels.string(),
all_class_ids.primitive::<u16>(),
all_show_labels.component::<ShowLabels>(),
all_show_labels.component_slow::<ShowLabels>(),
)
.map(
|(_index, half_sizes, colors, radii, labels, class_ids, show_labels)| {
Expand Down
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::<ShowLabels>(),
all_show_labels.component_slow::<ShowLabels>(),
all_class_ids.primitive::<u16>(),
)
.map(
Expand Down
4 changes: 2 additions & 2 deletions crates/viewer/re_view_spatial/src/visualizers/depth_images.rs
Original file line number Diff line number Diff line change
Expand Up @@ -292,7 +292,7 @@ impl VisualizerSystem for DepthImageVisualizer {
let mut data = re_query::range_zip_1x5(
all_buffers_indexed,
all_formats_indexed,
all_colormaps.component::<components::Colormap>(),
all_colormaps.primitive::<u8>(),
all_value_ranges.primitive_array::<2, f64>(),
all_depth_meters.primitive::<f32>(),
all_fill_ratios.primitive::<f32>(),
Expand All @@ -310,7 +310,7 @@ impl VisualizerSystem for DepthImageVisualizer {
},
depth_meter: first_copied(depth_meter).map(Into::into),
fill_ratio: first_copied(fill_ratio).map(Into::into),
colormap: first_copied(colormap.as_deref()),
colormap: first_copied(colormap).and_then(Colormap::from_u8),
value_range: first_copied(value_range).map(Into::into),
})
},
Expand Down
5 changes: 3 additions & 2 deletions crates/viewer/re_view_spatial/src/visualizers/ellipsoids.rs
Original file line number Diff line number Diff line change
Expand Up @@ -178,10 +178,10 @@ impl VisualizerSystem for Ellipsoids3DVisualizer {
all_half_sizes_indexed,
all_colors.primitive::<u32>(),
all_line_radii.primitive::<f32>(),
all_fill_modes.component::<FillMode>(),
all_fill_modes.primitive::<u8>(),
all_labels.string(),
all_class_ids.primitive::<u16>(),
all_show_labels.component::<ShowLabels>(),
all_show_labels.component_slow::<ShowLabels>(),
)
.map(
|(
Expand All @@ -204,6 +204,7 @@ impl VisualizerSystem for Ellipsoids3DVisualizer {
.unwrap_or_default()
.first()
.copied()
.and_then(FillMode::from_u8)
.unwrap_or_default(),
labels: labels.unwrap_or_default(),
class_ids: class_ids
Expand Down
2 changes: 1 addition & 1 deletion crates/viewer/re_view_spatial/src/visualizers/lines2d.rs
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::<ShowLabels>(),
all_show_labels.component_slow::<ShowLabels>(),
)
.map(
|(_index, strips, colors, radii, labels, class_ids, show_labels)| {
Expand Down
2 changes: 1 addition & 1 deletion crates/viewer/re_view_spatial/src/visualizers/lines3d.rs
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::<ShowLabels>(),
all_show_labels.component_slow::<ShowLabels>(),
)
.map(
|(_index, strips, colors, radii, labels, class_ids, show_labels)| {
Expand Down
6 changes: 4 additions & 2 deletions crates/viewer/re_view_spatial/src/visualizers/meshes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -214,8 +214,10 @@ impl VisualizerSystem for Mesh3DVisualizer {
all_vertex_texcoords.primitive_array::<2, f32>(),
all_triangle_indices.primitive_array::<3, u32>(),
all_albedo_factors.primitive::<u32>(),
all_albedo_buffers.component::<ImageBuffer>(),
all_albedo_formats.component::<ImageFormat>(),
// TODO(cmc): Provide a `iter_blob`.
all_albedo_buffers.component_slow::<ImageBuffer>(),
// Legit call to `component_slow`, `ImageFormat` is real complicated.
all_albedo_formats.component_slow::<ImageFormat>(),
all_class_ids.primitive::<u16>(),
)
.map(
Expand Down
2 changes: 1 addition & 1 deletion crates/viewer/re_view_spatial/src/visualizers/points2d.rs
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::<ShowLabels>(),
all_show_labels.component_slow::<ShowLabels>(),
)
.map(
|(
Expand Down
2 changes: 1 addition & 1 deletion crates/viewer/re_view_spatial/src/visualizers/points3d.rs
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::<ShowLabels>(),
all_show_labels.component_slow::<ShowLabels>(),
)
.map(
|(
Expand Down
9 changes: 7 additions & 2 deletions crates/viewer/re_view_tensor/src/visualizer_system.rs
Original file line number Diff line number Diff line change
Expand Up @@ -69,13 +69,18 @@ impl VisualizerSystem for TensorSystem {
let all_ranges = results.iter_as(timeline, ValueRange::name());

for ((_, tensor_row_id), tensors, data_ranges) in
re_query::range_zip_1x1(all_tensors_indexed, all_ranges.component::<ValueRange>())
re_query::range_zip_1x1(all_tensors_indexed, all_ranges.primitive_array::<2, f64>())
{
let Some(tensor) = tensors.first() else {
continue;
};
let data_range = data_ranges
.and_then(|ranges| ranges.first().copied())
.and_then(|ranges| {
ranges
.first()
.copied()
.map(|range| ValueRange(range.into()))
})
.unwrap_or_else(|| {
let tensor_stats = ctx
.viewer_ctx
Expand Down
Loading

0 comments on commit 5acbbab

Please sign in to comment.