From d78b7446555ff4ce6bdf7c531547064469276d8f Mon Sep 17 00:00:00 2001 From: Clement Rey Date: Tue, 30 Jul 2024 11:04:54 +0200 Subject: [PATCH 1/8] introduce new Chunk zero-copy zero-deser iterators --- crates/store/re_chunk/Cargo.toml | 1 + crates/store/re_chunk/src/iter.rs | 125 +++++++++++++++++++++++++++++- 2 files changed, 122 insertions(+), 4 deletions(-) diff --git a/crates/store/re_chunk/Cargo.toml b/crates/store/re_chunk/Cargo.toml index e6007cb1dbc6..8a4cd67eaa95 100644 --- a/crates/store/re_chunk/Cargo.toml +++ b/crates/store/re_chunk/Cargo.toml @@ -55,6 +55,7 @@ arrow2 = { workspace = true, features = [ "compute_filter", ] } backtrace.workspace = true +bytemuck.workspace = true document-features.workspace = true itertools.workspace = true nohash-hasher.workspace = true diff --git a/crates/store/re_chunk/src/iter.rs b/crates/store/re_chunk/src/iter.rs index 4dcc2406161e..7fe34094d970 100644 --- a/crates/store/re_chunk/src/iter.rs +++ b/crates/store/re_chunk/src/iter.rs @@ -1,13 +1,16 @@ use std::sync::Arc; use arrow2::{ - array::{Array as ArrowArray, PrimitiveArray}, + array::{ + Array as ArrowArray, FixedSizeListArray as ArrowFixedSizeListArray, + PrimitiveArray as ArrowPrimitiveArray, Utf8Array as ArrowUtf8Array, + }, Either, }; -use itertools::izip; +use itertools::{izip, Itertools}; use re_log_types::{TimeInt, Timeline}; -use re_types_core::{Component, ComponentName}; +use re_types_core::{ArrowString, Component, ComponentName}; use crate::{Chunk, ChunkTimeline, RowId}; @@ -125,6 +128,7 @@ impl Chunk { /// See also: /// * [`Self::iter_component`]. /// * [`Self::iter_primitive`]. + #[inline] pub fn iter_component_arrays( &self, component_name: &ComponentName, @@ -143,6 +147,9 @@ impl Chunk { /// Use this when working with simple arrow datatypes and performance matters (e.g. scalars, /// points, etc). /// + /// See also: + /// * [`Self::iter_primitive_array`] + /// * [`Self::iter_string`] /// * [`Self::iter_component_arrays`]. /// * [`Self::iter_component`]. #[inline] @@ -157,7 +164,7 @@ impl Chunk { let Some(values) = list_array .values() .as_any() - .downcast_ref::>() + .downcast_ref::>() else { if cfg!(debug_assertions) { panic!("downcast failed for {component_name}, data discarded"); @@ -174,6 +181,116 @@ impl Chunk { .map(move |(idx, len)| &values[idx..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 + /// component batch will be a slice reference into that global slice. + /// Use this when working with simple arrow datatypes and performance matters (e.g. scalars, + /// points, etc). + /// + /// See also: + /// * [`Self::iter_primitive`] + /// * [`Self::iter_string`] + /// * [`Self::iter_component_arrays`]. + /// * [`Self::iter_component`]. + pub fn iter_primitive_array( + &self, + component_name: &ComponentName, + ) -> impl Iterator + '_ + where + [T; N]: bytemuck::Pod, + { + let Some(list_array) = self.components.get(component_name) else { + return Either::Left(std::iter::empty()); + }; + + let Some(fixed_size_list_array) = 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 Some(values) = fixed_size_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 size = fixed_size_list_array.size(); + let values = values.values().as_slice(); + + // 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)| { + bytemuck::cast_slice(&values[idx * size..idx * size + len * size]) + }), + ) + } + + /// Returns an iterator over the raw primitive strings 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 (e.g. labels, etc). + /// + /// See also: + /// * [`Self::iter_primitive`] + /// * [`Self::iter_primitive_array`] + /// * [`Self::iter_component_arrays`]. + /// * [`Self::iter_component`]. + pub fn iter_string( + &self, + component_name: &ComponentName, + ) -> impl Iterator> + '_ { + let Some(list_array) = self.components.get(component_name) else { + return Either::Left(std::iter::empty()); + }; + + let Some(utf8_array) = 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 = utf8_array.values(); + let offsets = utf8_array.offsets(); + let lengths = utf8_array.offsets().lengths().collect_vec(); + + // 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)| { + let offsets = &offsets.as_slice()[idx..idx + len]; + let lengths = &lengths.as_slice()[idx..idx + len]; + izip!(offsets, lengths) + .map(|(&idx, &len)| ArrowString(values.clone().sliced(idx as _, len))) + .collect_vec() + }), + ) + } } // --- From 0852b4b662b5749b5076b132dc3ab5fc4bcf048e Mon Sep 17 00:00:00 2001 From: Clement Rey Date: Tue, 30 Jul 2024 11:06:12 +0200 Subject: [PATCH 2/8] introduce new HybridResults zero-deser helpers --- crates/viewer/re_space_view/Cargo.toml | 2 + .../viewer/re_space_view/src/results_ext2.rs | 80 +++++++++++++++++++ 2 files changed, 82 insertions(+) diff --git a/crates/viewer/re_space_view/Cargo.toml b/crates/viewer/re_space_view/Cargo.toml index 99453aa42b0d..488ca78f0f4d 100644 --- a/crates/viewer/re_space_view/Cargo.toml +++ b/crates/viewer/re_space_view/Cargo.toml @@ -35,5 +35,7 @@ re_viewer_context.workspace = true re_viewport_blueprint.workspace = true ahash.workspace = true +bytemuck.workspace = true egui.workspace = true +itertools.workspace = true nohash-hasher.workspace = true diff --git a/crates/viewer/re_space_view/src/results_ext2.rs b/crates/viewer/re_space_view/src/results_ext2.rs index 1c716d3b282a..cd6f842170d8 100644 --- a/crates/viewer/re_space_view/src/results_ext2.rs +++ b/crates/viewer/re_space_view/src/results_ext2.rs @@ -361,3 +361,83 @@ impl<'a> RangeResultsExt for HybridResults<'a> { } } } + +// --- + +use re_chunk::{RowId, TimeInt, Timeline}; +use re_chunk_store::external::{re_chunk, re_chunk::external::arrow2}; + +/// The iterator type backing [`HybridResults::iter_as`]. +pub struct HybridResultsChunkIter<'a> { + chunks: Cow<'a, [Chunk]>, + timeline: Timeline, + component_name: ComponentName, +} + +impl<'a> HybridResultsChunkIter<'a> { + /// Iterate as indexed primitives. + /// + /// See [`Chunk::iter_primitive`] for more information. + pub fn primitive( + &'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_primitive::(&self.component_name) + ) + }) + } + + /// Iterate as indexed primitive arrays. + /// + /// See [`Chunk::iter_primitive_array`] for more information. + pub fn primitive_array( + &'a self, + ) -> impl Iterator + 'a + where + [T; N]: bytemuck::Pod, + { + self.chunks.iter().flat_map(move |chunk| { + itertools::izip!( + chunk.iter_component_indices(&self.timeline, &self.component_name), + chunk.iter_primitive_array::(&self.component_name) + ) + }) + } + + /// Iterate as indexed UTF-8 strings. + /// + /// See [`Chunk::iter_string`] for more information. + pub fn string( + &'a self, + ) -> impl Iterator)> + 'a { + self.chunks.iter().flat_map(|chunk| { + itertools::izip!( + chunk.iter_component_indices(&self.timeline, &self.component_name), + chunk.iter_string(&self.component_name) + ) + }) + } +} + +impl<'a> HybridResults<'a> { + /// Returns a zero-copy iterator over all the results for the given `(timeline, component)` pair. + /// + /// Call one of the following methods on the returned [`HybridResultsChunkIter`]: + /// * [`HybridResultsChunkIter::primitive`] + /// * [`HybridResultsChunkIter::primitive_array`] + /// * [`HybridResultsChunkIter::string`] + pub fn iter_as( + &'a self, + timeline: Timeline, + component_name: ComponentName, + ) -> HybridResultsChunkIter<'a> { + let chunks = self.get_optional_chunks(&component_name); + HybridResultsChunkIter { + chunks, + timeline, + component_name, + } + } +} From a8404bbe796951700672935c9041a768777bf18d Mon Sep 17 00:00:00 2001 From: Clement Rey Date: Tue, 30 Jul 2024 11:06:54 +0200 Subject: [PATCH 3/8] introduce new chunk-based entity_iterator --- Cargo.lock | 4 + .../viewer/re_space_view_spatial/Cargo.toml | 1 + .../src/visualizers/mod.rs | 4 +- .../visualizers/utilities/entity_iterator.rs | 124 +++++++++++++++++- .../src/visualizers/utilities/labels.rs | 74 +++++++++++ .../src/visualizers/utilities/mod.rs | 3 +- 6 files changed, 206 insertions(+), 4 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index faa5aed94c38..43c5eaf72990 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4302,6 +4302,7 @@ dependencies = [ "ahash", "anyhow", "backtrace", + "bytemuck", "criterion", "crossbeam", "document-features", @@ -4919,7 +4920,9 @@ name = "re_space_view" version = "0.18.0-alpha.1+dev" dependencies = [ "ahash", + "bytemuck", "egui", + "itertools 0.13.0", "nohash-hasher", "re_chunk_store", "re_entity_db", @@ -4999,6 +5002,7 @@ dependencies = [ "re_log_types", "re_math", "re_query", + "re_query2", "re_renderer", "re_space_view", "re_tracing", diff --git a/crates/viewer/re_space_view_spatial/Cargo.toml b/crates/viewer/re_space_view_spatial/Cargo.toml index 10d53674e49d..d538505fde6b 100644 --- a/crates/viewer/re_space_view_spatial/Cargo.toml +++ b/crates/viewer/re_space_view_spatial/Cargo.toml @@ -28,6 +28,7 @@ re_log_types.workspace = true re_log.workspace = true re_math = { workspace = true, features = ["serde"] } re_query.workspace = true +re_query2.workspace = true re_renderer = { workspace = true, features = [ "import-gltf", "import-obj", diff --git a/crates/viewer/re_space_view_spatial/src/visualizers/mod.rs b/crates/viewer/re_space_view_spatial/src/visualizers/mod.rs index 5f4c19ce1021..f805c5f6aa65 100644 --- a/crates/viewer/re_space_view_spatial/src/visualizers/mod.rs +++ b/crates/viewer/re_space_view_spatial/src/visualizers/mod.rs @@ -27,8 +27,8 @@ pub use segmentation_images::SegmentationImageVisualizer; pub use transform3d_arrows::{add_axis_arrows, AxisLengthDetector, Transform3DArrowsVisualizer}; pub use utilities::{ bounding_box_for_textured_rect, entity_iterator, process_labels_2d, process_labels_3d, - textured_rect_from_image, textured_rect_from_tensor, SpatialViewVisualizerData, UiLabel, - UiLabelTarget, MAX_NUM_LABELS_PER_ENTITY, + process_labels_3d_2, textured_rect_from_image, textured_rect_from_tensor, + SpatialViewVisualizerData, UiLabel, UiLabelTarget, MAX_NUM_LABELS_PER_ENTITY, }; // --- diff --git a/crates/viewer/re_space_view_spatial/src/visualizers/utilities/entity_iterator.rs b/crates/viewer/re_space_view_spatial/src/visualizers/utilities/entity_iterator.rs index facb42d77ebf..fa599bbb5e46 100644 --- a/crates/viewer/re_space_view_spatial/src/visualizers/utilities/entity_iterator.rs +++ b/crates/viewer/re_space_view_spatial/src/visualizers/utilities/entity_iterator.rs @@ -2,7 +2,9 @@ use itertools::Either; use re_chunk_store::{LatestAtQuery, RangeQuery}; use re_log_types::{TimeInt, Timeline}; use re_space_view::{ - latest_at_with_blueprint_resolved_data, range_with_blueprint_resolved_data, HybridResults, + latest_at_with_blueprint_resolved_data, latest_at_with_blueprint_resolved_data2, + range_with_blueprint_resolved_data, range_with_blueprint_resolved_data2, HybridResults, + HybridResults2, }; use re_types::Archetype; use re_viewer_context::{ @@ -209,3 +211,123 @@ where Ok(()) } + +// --- Chunk-based APIs --- + +pub fn query_archetype_with_history2<'a, A: Archetype>( + ctx: &'a ViewContext<'a>, + timeline: &Timeline, + timeline_cursor: TimeInt, + query_range: &QueryRange, + data_result: &'a re_viewer_context::DataResult, +) -> HybridResults2<'a> { + match query_range { + QueryRange::TimeRange(time_range) => { + let range_query = RangeQuery::new( + *timeline, + re_log_types::ResolvedTimeRange::from_relative_time_range( + time_range, + timeline_cursor, + ), + ); + let results = range_with_blueprint_resolved_data2( + ctx, + None, + &range_query, + data_result, + A::all_components().iter().copied(), + ); + (range_query, results).into() + } + QueryRange::LatestAt => { + let latest_query = LatestAtQuery::new(*timeline, timeline_cursor); + let query_shadowed_defaults = false; + let results = latest_at_with_blueprint_resolved_data2( + ctx, + None, + &latest_query, + data_result, + A::all_components().iter().copied(), + query_shadowed_defaults, + ); + (latest_query, results).into() + } + } +} + +/// Iterates through all entity views for a given archetype. +/// +/// The callback passed in gets passed along an [`SpatialSceneEntityContext`] which contains +/// various useful information about an entity in the context of the current scene. +pub fn process_archetype2( + ctx: &ViewContext<'_>, + query: &ViewQuery<'_>, + view_ctx: &ViewContextCollection, + mut fun: F, +) -> Result<(), SpaceViewSystemExecutionError> +where + A: Archetype, + F: FnMut( + &QueryContext<'_>, + &SpatialSceneEntityContext<'_>, + &HybridResults2<'_>, + ) -> Result<(), SpaceViewSystemExecutionError>, +{ + let transforms = view_ctx.get::()?; + let depth_offsets = view_ctx.get::()?; + let annotations = view_ctx.get::()?; + + let latest_at = query.latest_at_query(); + + let system_identifier = System::identifier(); + + for data_result in query.iter_visible_data_results(ctx, system_identifier) { + // The transform that considers pinholes only makes sense if this is a 3D space-view + let world_from_entity = + if view_ctx.space_view_class_identifier() == SpatialSpaceView3D::identifier() { + transforms.reference_from_entity(&data_result.entity_path) + } else { + transforms.reference_from_entity_ignoring_pinhole( + &data_result.entity_path, + ctx.recording(), + &latest_at, + ) + }; + + let Some(world_from_entity) = world_from_entity else { + continue; + }; + let depth_offset_key = (system_identifier, data_result.entity_path.hash()); + let entity_context = SpatialSceneEntityContext { + world_from_entity, + depth_offset: depth_offsets + .per_entity_and_visualizer + .get(&depth_offset_key) + .copied() + .unwrap_or_default(), + annotations: annotations.0.find(&data_result.entity_path), + highlight: query + .highlights + .entity_outline_mask(data_result.entity_path.hash()), + space_view_class_identifier: view_ctx.space_view_class_identifier(), + }; + + let results = query_archetype_with_history2::( + ctx, + &query.timeline, + query.latest_at, + data_result.query_range(), + data_result, + ); + + let mut query_ctx = ctx.query_context(data_result, &latest_at); + query_ctx.archetype_name = Some(A::name()); + + { + re_tracing::profile_scope!(format!("{}", data_result.entity_path)); + fun(&query_ctx, &entity_context, &results)?; + } + } + + Ok(()) +} diff --git a/crates/viewer/re_space_view_spatial/src/visualizers/utilities/labels.rs b/crates/viewer/re_space_view_spatial/src/visualizers/utilities/labels.rs index 031c7ed15dfa..f0d221556d1d 100644 --- a/crates/viewer/re_space_view_spatial/src/visualizers/utilities/labels.rs +++ b/crates/viewer/re_space_view_spatial/src/visualizers/utilities/labels.rs @@ -39,6 +39,8 @@ pub const MAX_NUM_LABELS_PER_ENTITY: usize = 30; /// /// Does nothing if there's no positions or no labels passed. /// Otherwise, produces one label per position passed. +// +// TODO(cmc): remove pub fn process_labels_3d<'a>( entity_path: &'a EntityPath, positions: impl Iterator + 'a, @@ -67,10 +69,44 @@ pub fn process_labels_3d<'a>( }) } +/// Produces 3D ui labels from component data. +/// +/// Does nothing if there's no positions or no labels passed. +/// Otherwise, produces one label per position passed. +pub fn process_labels_3d_2<'a>( + entity_path: &'a EntityPath, + positions: impl Iterator + 'a, + labels: &'a [re_types::ArrowString], + colors: &'a [egui::Color32], + annotation_infos: &'a ResolvedAnnotationInfos, + world_from_obj: glam::Affine3A, +) -> impl Iterator + 'a { + let labels = izip!( + annotation_infos.iter(), + labels.iter().map(Some).chain(std::iter::repeat(None)) + ) + .map(|(annotation_info, label)| annotation_info.label(label.map(|l| l.as_str()))); + + let colors = clamped_or(colors, &egui::Color32::WHITE); + + itertools::izip!(positions, labels, colors) + .enumerate() + .filter_map(move |(i, (point, label, color))| { + label.map(|label| UiLabel { + text: label, + color: *color, + target: UiLabelTarget::Position3D(world_from_obj.transform_point3(point)), + labeled_instance: InstancePathHash::instance(entity_path, Instance::from(i as u64)), + }) + }) +} + /// Produces 2D ui labels from component data. /// /// Does nothing if there's no positions or no labels passed. /// Otherwise, produces one label per position passed. +// +// TODO(cmc): remove pub fn process_labels_2d<'a>( entity_path: &'a EntityPath, positions: impl Iterator + 'a, @@ -104,3 +140,41 @@ pub fn process_labels_2d<'a>( }) }) } + +/// Produces 2D ui labels from component data. +/// +/// Does nothing if there's no positions or no labels passed. +/// Otherwise, produces one label per position passed. +pub fn process_labels_2d_2<'a>( + entity_path: &'a EntityPath, + positions: impl Iterator + 'a, + labels: &'a [re_types::ArrowString], + colors: &'a [egui::Color32], + annotation_infos: &'a ResolvedAnnotationInfos, + world_from_obj: glam::Affine3A, +) -> impl Iterator + 'a { + let labels = izip!( + annotation_infos.iter(), + labels.iter().map(Some).chain(std::iter::repeat(None)) + ) + .map(|(annotation_info, label)| annotation_info.label(label.map(|l| l.as_str()))); + + let colors = clamped_or(colors, &egui::Color32::WHITE); + + itertools::izip!(positions, labels, colors) + .enumerate() + .filter_map(move |(i, (point, label, color))| { + label.map(|label| { + let point = world_from_obj.transform_point3(glam::Vec3::new(point.x, point.y, 0.0)); + UiLabel { + text: label, + color: *color, + target: UiLabelTarget::Point2D(egui::pos2(point.x, point.y)), + labeled_instance: InstancePathHash::instance( + entity_path, + Instance::from(i as u64), + ), + } + }) + }) +} diff --git a/crates/viewer/re_space_view_spatial/src/visualizers/utilities/mod.rs b/crates/viewer/re_space_view_spatial/src/visualizers/utilities/mod.rs index 57bca95c769a..b450dc6eff64 100644 --- a/crates/viewer/re_space_view_spatial/src/visualizers/utilities/mod.rs +++ b/crates/viewer/re_space_view_spatial/src/visualizers/utilities/mod.rs @@ -4,7 +4,8 @@ mod spatial_view_visualizer; mod textured_rect; pub use labels::{ - process_labels_2d, process_labels_3d, UiLabel, UiLabelTarget, MAX_NUM_LABELS_PER_ENTITY, + process_labels_2d, process_labels_2d_2, process_labels_3d, process_labels_3d_2, UiLabel, + UiLabelTarget, MAX_NUM_LABELS_PER_ENTITY, }; pub use spatial_view_visualizer::SpatialViewVisualizerData; pub use textured_rect::{ From a5004e162f815d986ce83b1810b69aac2f504172 Mon Sep 17 00:00:00 2001 From: Clement Rey Date: Tue, 30 Jul 2024 11:07:05 +0200 Subject: [PATCH 4/8] zero-deser point2d visualizer --- .../src/visualizers/points2d.rs | 77 +++++++++++-------- 1 file changed, 43 insertions(+), 34 deletions(-) diff --git a/crates/viewer/re_space_view_spatial/src/visualizers/points2d.rs b/crates/viewer/re_space_view_spatial/src/visualizers/points2d.rs index 0734d69cdf05..08cec84d7ff0 100644 --- a/crates/viewer/re_space_view_spatial/src/visualizers/points2d.rs +++ b/crates/viewer/re_space_view_spatial/src/visualizers/points2d.rs @@ -1,10 +1,10 @@ use itertools::Itertools as _; -use re_query::range_zip_1x5; use re_renderer::{LineDrawableBuilder, PickingLayerInstanceId, PointCloudBuilder}; use re_types::{ archetypes::Points2D, components::{ClassId, Color, DrawOrder, KeypointId, Position2D, Radius, Text}, + ArrowString, Loggable as _, }; use re_viewer_context::{ auto_color_for_entity_path, ApplicableEntities, IdentifiedViewSystem, QueryContext, @@ -23,7 +23,7 @@ use crate::{ }; use super::{ - filter_visualizable_2d_entities, process_labels_2d, SpatialViewVisualizerData, + filter_visualizable_2d_entities, utilities::process_labels_2d_2, SpatialViewVisualizerData, SIZE_BOOST_IN_POINTS_FOR_POINT_OUTLINES, }; @@ -141,10 +141,10 @@ impl Points2DVisualizer { ) }; - self.data.ui_labels.extend(process_labels_2d( + self.data.ui_labels.extend(process_labels_2d_2( entity_path, label_positions, - data.labels, + &data.labels, &colors, &annotation_infos, ent_context.world_from_entity, @@ -166,7 +166,7 @@ pub struct Points2DComponentData<'a> { // Clamped to edge pub colors: &'a [Color], pub radii: &'a [Radius], - pub labels: &'a [Text], + pub labels: Vec, pub keypoint_ids: &'a [KeypointId], pub class_ids: &'a [ClassId], } @@ -211,53 +211,62 @@ impl VisualizerSystem for Points2DVisualizer { line_builder .radius_boost_in_ui_points_for_outlines(SIZE_BOOST_IN_POINTS_FOR_POINT_OUTLINES); - super::entity_iterator::process_archetype::( + super::entity_iterator::process_archetype2::( ctx, view_query, context_systems, |ctx, spatial_ctx, results| { - use re_space_view::RangeResultsExt as _; + use re_space_view::RangeResultsExt2 as _; - let resolver = ctx.recording().resolver(); - - let positions = match results.get_required_component_dense::(resolver) { - Some(positions) => positions?, - _ => return Ok(()), + let Some(all_position_chunks) = results.get_required_chunks(&Position2D::name()) + else { + return Ok(()); }; - let num_positions = positions - .range_indexed() - .map(|(_, positions)| positions.len()) - .sum::(); + let num_positions = all_position_chunks + .iter() + .flat_map(|chunk| chunk.iter_primitive_array::<2, f32>(&Position2D::name())) + .map(|points| points.len()) + .sum(); + if num_positions == 0 { return Ok(()); } point_builder.reserve(num_positions)?; - let colors = results.get_or_empty_dense(resolver)?; - let radii = results.get_or_empty_dense(resolver)?; - let labels = results.get_or_empty_dense(resolver)?; - let class_ids = results.get_or_empty_dense(resolver)?; - let keypoint_ids = results.get_or_empty_dense(resolver)?; - - let data = range_zip_1x5( - positions.range_indexed(), - colors.range_indexed(), - radii.range_indexed(), - labels.range_indexed(), - class_ids.range_indexed(), - keypoint_ids.range_indexed(), + let timeline = ctx.query.timeline(); + let all_positions_indexed = all_position_chunks.iter().flat_map(|chunk| { + itertools::izip!( + chunk.iter_component_indices(&timeline, &Position2D::name()), + chunk.iter_primitive_array::<2, f32>(&Position2D::name()) + ) + }); + let all_colors = results.iter_as(timeline, Color::name()); + let all_radii = results.iter_as(timeline, Radius::name()); + let all_labels = results.iter_as(timeline, Text::name()); + let all_class_ids = results.iter_as(timeline, ClassId::name()); + let all_keypoint_ids = results.iter_as(timeline, KeypointId::name()); + + let data = re_query2::range_zip_1x5( + all_positions_indexed, + all_colors.primitive::(), + all_radii.primitive::(), + all_labels.string(), + all_class_ids.primitive::(), + all_keypoint_ids.primitive::(), ) .map( |(_index, positions, colors, radii, labels, class_ids, keypoint_ids)| { Points2DComponentData { - positions, - colors: colors.unwrap_or_default(), - radii: radii.unwrap_or_default(), + positions: bytemuck::cast_slice(positions), + colors: colors.map_or(&[], |colors| bytemuck::cast_slice(colors)), + radii: radii.map_or(&[], |radii| bytemuck::cast_slice(radii)), labels: labels.unwrap_or_default(), - class_ids: class_ids.unwrap_or_default(), - keypoint_ids: keypoint_ids.unwrap_or_default(), + class_ids: class_ids + .map_or(&[], |class_ids| bytemuck::cast_slice(class_ids)), + keypoint_ids: keypoint_ids + .map_or(&[], |keypoint_ids| bytemuck::cast_slice(keypoint_ids)), } }, ); From d1ebc4a88cb3cc36326f70f592321a38d07ea7f4 Mon Sep 17 00:00:00 2001 From: Clement Rey Date: Tue, 30 Jul 2024 11:07:12 +0200 Subject: [PATCH 5/8] zero-deser point3d visualizer --- .../src/visualizers/points3d.rs | 80 +++++++++++-------- 1 file changed, 45 insertions(+), 35 deletions(-) diff --git a/crates/viewer/re_space_view_spatial/src/visualizers/points3d.rs b/crates/viewer/re_space_view_spatial/src/visualizers/points3d.rs index 50873b163273..4243b8c2414f 100644 --- a/crates/viewer/re_space_view_spatial/src/visualizers/points3d.rs +++ b/crates/viewer/re_space_view_spatial/src/visualizers/points3d.rs @@ -1,10 +1,10 @@ -use itertools::Itertools as _; +use itertools::Itertools; -use re_query::range_zip_1x5; use re_renderer::{LineDrawableBuilder, PickingLayerInstanceId, PointCloudBuilder}; use re_types::{ archetypes::Points3D, components::{ClassId, Color, KeypointId, Position3D, Radius, Text}, + ArrowString, Loggable, }; use re_viewer_context::{ auto_color_for_entity_path, ApplicableEntities, IdentifiedViewSystem, QueryContext, @@ -23,7 +23,7 @@ use crate::{ }; use super::{ - filter_visualizable_3d_entities, process_labels_3d, SpatialViewVisualizerData, + filter_visualizable_3d_entities, process_labels_3d_2, SpatialViewVisualizerData, SIZE_BOOST_IN_POINTS_FOR_POINT_OUTLINES, }; @@ -48,7 +48,7 @@ struct Points3DComponentData<'a> { // Clamped to edge colors: &'a [Color], radii: &'a [Radius], - labels: &'a [Text], + labels: Vec, keypoint_ids: &'a [KeypointId], class_ids: &'a [ClassId], } @@ -144,10 +144,10 @@ impl Points3DVisualizer { positions }; - self.data.ui_labels.extend(process_labels_3d( + self.data.ui_labels.extend(process_labels_3d_2( entity_path, label_positions.iter().copied(), - data.labels, + &data.labels, &colors, &annotation_infos, ent_context.world_from_entity, @@ -199,53 +199,62 @@ impl VisualizerSystem for Points3DVisualizer { line_builder .radius_boost_in_ui_points_for_outlines(SIZE_BOOST_IN_POINTS_FOR_POINT_OUTLINES); - super::entity_iterator::process_archetype::( + super::entity_iterator::process_archetype2::( ctx, view_query, context_systems, |ctx, spatial_ctx, results| { - use re_space_view::RangeResultsExt as _; + use re_space_view::RangeResultsExt2 as _; - let resolver = ctx.recording().resolver(); - - let positions = match results.get_required_component_dense::(resolver) { - Some(positions) => positions?, - _ => return Ok(()), + let Some(all_position_chunks) = results.get_required_chunks(&Position3D::name()) + else { + return Ok(()); }; - let num_positions = positions - .range_indexed() - .map(|(_, positions)| positions.len()) - .sum::(); + let num_positions = all_position_chunks + .iter() + .flat_map(|chunk| chunk.iter_primitive_array::<3, f32>(&Position3D::name())) + .map(|points| points.len()) + .sum(); + if num_positions == 0 { return Ok(()); } point_builder.reserve(num_positions)?; - let colors = results.get_or_empty_dense(resolver)?; - let radii = results.get_or_empty_dense(resolver)?; - let labels = results.get_or_empty_dense(resolver)?; - let class_ids = results.get_or_empty_dense(resolver)?; - let keypoint_ids = results.get_or_empty_dense(resolver)?; - - let data = range_zip_1x5( - positions.range_indexed(), - colors.range_indexed(), - radii.range_indexed(), - labels.range_indexed(), - class_ids.range_indexed(), - keypoint_ids.range_indexed(), + let timeline = ctx.query.timeline(); + let all_positions_indexed = all_position_chunks.iter().flat_map(|chunk| { + itertools::izip!( + chunk.iter_component_indices(&timeline, &Position3D::name()), + chunk.iter_primitive_array::<3, f32>(&Position3D::name()) + ) + }); + let all_colors = results.iter_as(timeline, Color::name()); + let all_radii = results.iter_as(timeline, Radius::name()); + let all_labels = results.iter_as(timeline, Text::name()); + let all_class_ids = results.iter_as(timeline, ClassId::name()); + let all_keypoint_ids = results.iter_as(timeline, KeypointId::name()); + + let data = re_query2::range_zip_1x5( + all_positions_indexed, + all_colors.primitive::(), + all_radii.primitive::(), + all_labels.string(), + all_class_ids.primitive::(), + all_keypoint_ids.primitive::(), ) .map( |(_index, positions, colors, radii, labels, class_ids, keypoint_ids)| { Points3DComponentData { - positions, - colors: colors.unwrap_or_default(), - radii: radii.unwrap_or_default(), + positions: bytemuck::cast_slice(positions), + colors: colors.map_or(&[], |colors| bytemuck::cast_slice(colors)), + radii: radii.map_or(&[], |radii| bytemuck::cast_slice(radii)), labels: labels.unwrap_or_default(), - class_ids: class_ids.unwrap_or_default(), - keypoint_ids: keypoint_ids.unwrap_or_default(), + class_ids: class_ids + .map_or(&[], |class_ids| bytemuck::cast_slice(class_ids)), + keypoint_ids: keypoint_ids + .map_or(&[], |keypoint_ids| bytemuck::cast_slice(keypoint_ids)), } }, ); @@ -281,6 +290,7 @@ impl VisualizerSystem for Points3DVisualizer { } impl TypedComponentFallbackProvider for Points3DVisualizer { + #[inline] fn fallback_for(&self, ctx: &QueryContext<'_>) -> Color { auto_color_for_entity_path(ctx.target_entity_path) } From 76244c7244d643141c215c60c189dfd4b74ee947 Mon Sep 17 00:00:00 2001 From: Clement Rey Date: Tue, 30 Jul 2024 11:53:19 +0200 Subject: [PATCH 6/8] lints --- crates/viewer/re_space_view/src/lib.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/viewer/re_space_view/src/lib.rs b/crates/viewer/re_space_view/src/lib.rs index ac3015ceeb0f..7b90363a1eb2 100644 --- a/crates/viewer/re_space_view/src/lib.rs +++ b/crates/viewer/re_space_view/src/lib.rs @@ -24,7 +24,7 @@ pub use query2::{ pub use results_ext::{HybridLatestAtResults, HybridResults, RangeResultsExt}; pub use results_ext2::{ HybridLatestAtResults as HybridLatestAtResults2, HybridResults as HybridResults2, - RangeResultsExt as RangeResultsExt2, + HybridResultsChunkIter, RangeResultsExt as RangeResultsExt2, }; pub use screenshot::ScreenshotMode; pub use view_property_ui::view_property_ui; From 0cced748c10efe6475adaca55638104de13bc169 Mon Sep 17 00:00:00 2001 From: Clement Rey Date: Tue, 30 Jul 2024 15:38:09 +0200 Subject: [PATCH 7/8] unrelated fix to chunkified override APIs --- crates/viewer/re_space_view/src/query2.rs | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/crates/viewer/re_space_view/src/query2.rs b/crates/viewer/re_space_view/src/query2.rs index 72a7a2de8b47..762084943257 100644 --- a/crates/viewer/re_space_view/src/query2.rs +++ b/crates/viewer/re_space_view/src/query2.rs @@ -137,6 +137,12 @@ fn query_overrides<'a>( .resolved_component_overrides .get(component_name) { + let current_query = match override_value.store_kind { + re_log_types::StoreKind::Recording => ctx.current_query(), + re_log_types::StoreKind::Blueprint => ctx.blueprint_query.clone(), + }; + + #[allow(clippy::match_same_arms)] // see @jleibs comment below let component_override_result = match override_value.store_kind { re_log_types::StoreKind::Recording => { // TODO(jleibs): This probably is not right, but this code path is not used @@ -144,7 +150,7 @@ fn query_overrides<'a>( // component override data-references are resolved. ctx.store_context.blueprint.query_caches2().latest_at( ctx.store_context.blueprint.store(), - &ctx.current_query(), + ¤t_query, &override_value.path, [*component_name], ) @@ -152,7 +158,7 @@ fn query_overrides<'a>( re_log_types::StoreKind::Blueprint => { ctx.store_context.blueprint.query_caches2().latest_at( ctx.store_context.blueprint.store(), - ctx.blueprint_query, + ¤t_query, &override_value.path, [*component_name], ) @@ -168,10 +174,10 @@ fn query_overrides<'a>( // This is extra tricky since the promise hasn't been resolved yet so we can't // actually look at the data. if let Some(value) = component_override_result.components.get(component_name) { - let index = value.index(&ctx.current_query().timeline()); + let index = value.index(¤t_query.timeline()); // NOTE: This can never happen, but I'd rather it happens than an unwrap. - debug_assert!(index.is_some()); + debug_assert!(index.is_some(), "{value:#?}"); let index = index.unwrap_or((TimeInt::STATIC, RowId::ZERO)); overrides.add(*component_name, index, value.clone()); From 0e29c96e2b4920f330814b5b9143b22009063077 Mon Sep 17 00:00:00 2001 From: Clement Rey Date: Tue, 30 Jul 2024 16:49:43 +0200 Subject: [PATCH 8/8] review --- .../src/visualizers/points2d.rs | 14 ++--- .../src/visualizers/points3d.rs | 14 ++--- .../visualizers/utilities/entity_iterator.rs | 58 +++++++++++++++++++ 3 files changed, 72 insertions(+), 14 deletions(-) diff --git a/crates/viewer/re_space_view_spatial/src/visualizers/points2d.rs b/crates/viewer/re_space_view_spatial/src/visualizers/points2d.rs index 08cec84d7ff0..ea6fd138026a 100644 --- a/crates/viewer/re_space_view_spatial/src/visualizers/points2d.rs +++ b/crates/viewer/re_space_view_spatial/src/visualizers/points2d.rs @@ -211,7 +211,8 @@ impl VisualizerSystem for Points2DVisualizer { line_builder .radius_boost_in_ui_points_for_outlines(SIZE_BOOST_IN_POINTS_FOR_POINT_OUTLINES); - super::entity_iterator::process_archetype2::( + use super::entity_iterator::{iter_primitive_array, process_archetype2}; + process_archetype2::( ctx, view_query, context_systems, @@ -236,12 +237,11 @@ impl VisualizerSystem for Points2DVisualizer { point_builder.reserve(num_positions)?; let timeline = ctx.query.timeline(); - let all_positions_indexed = all_position_chunks.iter().flat_map(|chunk| { - itertools::izip!( - chunk.iter_component_indices(&timeline, &Position2D::name()), - chunk.iter_primitive_array::<2, f32>(&Position2D::name()) - ) - }); + let all_positions_indexed = iter_primitive_array::<2, f32>( + &all_position_chunks, + timeline, + Position2D::name(), + ); let all_colors = results.iter_as(timeline, Color::name()); let all_radii = results.iter_as(timeline, Radius::name()); let all_labels = results.iter_as(timeline, Text::name()); diff --git a/crates/viewer/re_space_view_spatial/src/visualizers/points3d.rs b/crates/viewer/re_space_view_spatial/src/visualizers/points3d.rs index 4243b8c2414f..04db96274ec8 100644 --- a/crates/viewer/re_space_view_spatial/src/visualizers/points3d.rs +++ b/crates/viewer/re_space_view_spatial/src/visualizers/points3d.rs @@ -199,7 +199,8 @@ impl VisualizerSystem for Points3DVisualizer { line_builder .radius_boost_in_ui_points_for_outlines(SIZE_BOOST_IN_POINTS_FOR_POINT_OUTLINES); - super::entity_iterator::process_archetype2::( + use super::entity_iterator::{iter_primitive_array, process_archetype2}; + process_archetype2::( ctx, view_query, context_systems, @@ -224,12 +225,11 @@ impl VisualizerSystem for Points3DVisualizer { point_builder.reserve(num_positions)?; let timeline = ctx.query.timeline(); - let all_positions_indexed = all_position_chunks.iter().flat_map(|chunk| { - itertools::izip!( - chunk.iter_component_indices(&timeline, &Position3D::name()), - chunk.iter_primitive_array::<3, f32>(&Position3D::name()) - ) - }); + let all_positions_indexed = iter_primitive_array::<3, f32>( + &all_position_chunks, + timeline, + Position3D::name(), + ); let all_colors = results.iter_as(timeline, Color::name()); let all_radii = results.iter_as(timeline, Radius::name()); let all_labels = results.iter_as(timeline, Text::name()); diff --git a/crates/viewer/re_space_view_spatial/src/visualizers/utilities/entity_iterator.rs b/crates/viewer/re_space_view_spatial/src/visualizers/utilities/entity_iterator.rs index fa599bbb5e46..2adf4a0a40ef 100644 --- a/crates/viewer/re_space_view_spatial/src/visualizers/utilities/entity_iterator.rs +++ b/crates/viewer/re_space_view_spatial/src/visualizers/utilities/entity_iterator.rs @@ -331,3 +331,61 @@ where Ok(()) } + +// --- + +use re_chunk_store::external::{re_chunk, re_chunk::external::arrow2}; + +/// Iterate `chunks` as indexed primitives. +/// +/// See [`re_chunk::Chunk::iter_primitive`] for more information. +#[allow(unused)] +pub fn iter_primitive<'a, T: arrow2::types::NativeType>( + chunks: impl IntoIterator + 'a, + timeline: re_chunk::Timeline, + component_name: re_chunk::ComponentName, +) -> impl Iterator + 'a { + chunks.into_iter().flat_map(move |chunk| { + itertools::izip!( + chunk.iter_component_indices(&timeline, &component_name), + chunk.iter_primitive::(&component_name) + ) + }) +} + +/// Iterate `chunks` as indexed primitive arrays. +/// +/// See [`re_chunk::Chunk::iter_primitive_array`] for more information. +#[allow(unused)] +pub fn iter_primitive_array<'a, const N: usize, T: arrow2::types::NativeType>( + chunks: &'a std::borrow::Cow<'a, [re_chunk::Chunk]>, + timeline: re_chunk::Timeline, + component_name: re_chunk::ComponentName, +) -> impl Iterator + 'a +where + [T; N]: bytemuck::Pod, +{ + chunks.iter().flat_map(move |chunk| { + itertools::izip!( + chunk.iter_component_indices(&timeline, &component_name), + chunk.iter_primitive_array::(&component_name) + ) + }) +} + +/// Iterate `chunks` as indexed UTF-8 strings. +/// +/// See [`re_chunk::Chunk::iter_string`] for more information. +#[allow(unused)] +pub fn iter_string<'a>( + chunks: impl Iterator + 'a, + timeline: &'a re_chunk::Timeline, + component_name: &'a re_chunk::ComponentName, +) -> impl Iterator)> + 'a { + chunks.into_iter().flat_map(|chunk| { + itertools::izip!( + chunk.iter_component_indices(timeline, component_name), + chunk.iter_string(component_name) + ) + }) +}