diff --git a/crates/store/re_chunk/src/iter.rs b/crates/store/re_chunk/src/iter.rs index d93a95fe9ffd..ea4a548b5394 100644 --- a/crates/store/re_chunk/src/iter.rs +++ b/crates/store/re_chunk/src/iter.rs @@ -1,4 +1,7 @@ -use std::sync::Arc; +use std::{ + borrow::{Borrow as _, Cow}, + sync::Arc, +}; use arrow::{ array::{ @@ -743,6 +746,106 @@ impl ChunkComponentSlicer for bool { } } +/// Generic component slicer that casts to a primitive type. +/// +/// In the happy path (when the array is already the target type), this performs zero-copy slicing. +/// When casting is required, it allocates and owns the casted array. +pub struct CastToPrimitive +where + P: arrow::array::ArrowPrimitiveType, + T: arrow::datatypes::ArrowNativeType, +{ + _phantom: std::marker::PhantomData<(P, T)>, +} + +/// Iterator that owns the array values and component spans. +/// +/// This is necessary when we need to cast the array, as the casted array +/// must be owned by the iterator rather than borrowed from the caller. +struct OwnedSliceIterator<'a, T, I> +where + T: arrow::datatypes::ArrowNativeType, + I: Iterator>, +{ + values: Cow<'a, arrow::buffer::ScalarBuffer>, + component_spans: I, +} + +impl<'a, T, I> Iterator for OwnedSliceIterator<'a, T, I> +where + T: arrow::datatypes::ArrowNativeType + Clone, + I: Iterator>, +{ + type Item = Cow<'a, [T]>; + + fn next(&mut self) -> Option { + let span = self.component_spans.next()?; + match &self.values { + Cow::Borrowed(values) => Some(Cow::Borrowed(&values[span.range()])), + // TODO(grtlr): This `clone` here makes me sad, but I don't see a way around it. + Cow::Owned(values) => Some(Cow::Owned(values[span.range()].to_vec())), + } + } +} + +fn error_on_cast_failure( + component: ComponentIdentifier, + target: &arrow::datatypes::DataType, + actual: &arrow::datatypes::DataType, + error: &arrow::error::ArrowError, +) { + if cfg!(debug_assertions) { + panic!( + "[DEBUG ASSERT] cast from {actual:?} to {target:?} failed for {component}: {error}. Data discarded" + ); + } else { + re_log::error_once!( + "cast from {actual:?} to {target:?} failed for {component}: {error}. Data discarded" + ); + } +} + +impl ChunkComponentSlicer for CastToPrimitive +where + P: arrow::array::ArrowPrimitiveType, + T: arrow::datatypes::ArrowNativeType + Clone, +{ + type Item<'a> = Cow<'a, [T]>; + + fn slice<'a>( + component: ComponentIdentifier, + array: &'a dyn arrow::array::Array, + component_spans: impl Iterator> + 'a, + ) -> impl Iterator> + 'a { + // We first try to down cast (happy path - zero copy). + if let Some(values) = array.downcast_array_ref::>() { + return Either::Right(OwnedSliceIterator { + values: Cow::Borrowed(values.values()), + component_spans, + }); + } + + // Then we try to perform a primitive cast (requires ownership). + let casted = match arrow::compute::cast(array, &P::DATA_TYPE) { + Ok(casted) => casted, + Err(err) => { + error_on_cast_failure(component, &P::DATA_TYPE, array.data_type(), &err); + return Either::Left(std::iter::empty()); + } + }; + + let Some(values) = casted.downcast_array_ref::>() else { + error_on_downcast_failure(component, "ArrowPrimitiveArray", array.data_type()); + return Either::Left(std::iter::empty()); + }; + + Either::Right(OwnedSliceIterator { + values: Cow::Owned(values.values().clone()), + component_spans, + }) + } +} + // --- pub struct ChunkIndicesIter { @@ -945,11 +1048,11 @@ fn error_on_downcast_failure( ) { if cfg!(debug_assertions) { panic!( - "[DEBUG ASSERT] downcast failed to {target} failed for {component}. Array data type was {actual:?}. Data discarded" + "[DEBUG ASSERT] downcast to {target} failed for {component}. Array data type was {actual:?}. Data discarded" ); } else { re_log::error_once!( - "downcast failed to {target} for {component}. Array data type was {actual:?}. data discarded" + "downcast to {target} failed for {component}. Array data type was {actual:?}. Data discarded" ); } } diff --git a/crates/store/re_chunk/src/lib.rs b/crates/store/re_chunk/src/lib.rs index 2a4753eddacb..7ea8462fef94 100644 --- a/crates/store/re_chunk/src/lib.rs +++ b/crates/store/re_chunk/src/lib.rs @@ -24,7 +24,12 @@ pub use self::chunk::{ }; pub use self::helpers::{ChunkShared, UnitChunkShared}; pub use self::iter::{ - ChunkComponentIter, ChunkComponentIterItem, ChunkComponentSlicer, ChunkIndicesIter, + // TODO: Right place? Maybe `re_view`? + CastToPrimitive, + ChunkComponentIter, + ChunkComponentIterItem, + ChunkComponentSlicer, + ChunkIndicesIter, }; pub use self::latest_at::LatestAtQuery; pub use self::range::{RangeQuery, RangeQueryOptions}; diff --git a/crates/store/re_types/src/dynamic_archetype.rs b/crates/store/re_types/src/dynamic_archetype.rs index 25373d51470a..5e33caeaceb7 100644 --- a/crates/store/re_types/src/dynamic_archetype.rs +++ b/crates/store/re_types/src/dynamic_archetype.rs @@ -62,6 +62,36 @@ impl DynamicArchetype { self } + // TODO: hackity + /// Adds a field of arbitrary data to this archetype. + /// + /// In many cases, it might be more convenient to use [`Self::with_component`] to log an existing Rerun component instead. + #[inline] + pub fn with_component_from_data_with_type( + mut self, + field: impl AsRef, + component_type: ComponentType, + array: arrow::array::ArrayRef, + ) -> Self { + let field = field.as_ref(); + let component = field.into(); + + self.batches.insert( + component, + SerializedComponentBatch { + array, + descriptor: { + let mut desc = ComponentDescriptor::partial(component); + if let Some(archetype_name) = self.archetype_name { + desc = desc.with_builtin_archetype(archetype_name); + } + desc.with_component_type(component_type) + }, + }, + ); + self + } + /// Adds an existing Rerun [`Component`] to this archetype. #[inline] pub fn with_component( diff --git a/crates/viewer/re_view_time_series/src/series_query.rs b/crates/viewer/re_view_time_series/src/series_query.rs index 00d2c830f5e0..127929759f24 100644 --- a/crates/viewer/re_view_time_series/src/series_query.rs +++ b/crates/viewer/re_view_time_series/src/series_query.rs @@ -3,7 +3,9 @@ use itertools::Itertools as _; use re_chunk_store::RangeQuery; +use re_chunk_store::external::re_chunk::CastToPrimitive; use re_log_types::{EntityPath, TimeInt}; +use re_types::external::arrow; use re_types::external::arrow::datatypes::DataType as ArrowDatatype; use re_types::{ComponentDescriptor, ComponentIdentifier, Loggable as _, RowId, components}; use re_view::{ChunksWithComponent, HybridRangeResults, RangeResultsExt as _, clamped_or_nothing}; @@ -22,8 +24,10 @@ pub fn determine_num_series(all_scalar_chunks: &ChunksWithComponent<'_>) -> usiz .iter() .find_map(|chunk| { chunk - .iter_slices::() - .find_map(|slice| (!slice.is_empty()).then_some(slice.len())) + // TODO(grtlr): The comment above is even more important when we do casting (allocations) here. + // TODO(grtlr): Unless we're on the happy path this allocates and happens everytime the visualizer runs! + .iter_slices::>() + .find_map(|values| (!values.is_empty()).then_some(values.len())) }) .unwrap_or(1) } @@ -95,7 +99,10 @@ pub fn collect_scalars( let points = &mut *points_per_series[0]; all_scalar_chunks .iter() - .flat_map(|chunk| chunk.iter_slices::()) + .flat_map(|chunk| { + // TODO(grtlr): Unless we're on the happy path this allocates and happens everytime the visualizer runs! + chunk.iter_slices::>() + }) .enumerate() .for_each(|(i, values)| { if let Some(value) = values.first() { @@ -107,13 +114,17 @@ pub fn collect_scalars( } else { all_scalar_chunks .iter() - .flat_map(|chunk| chunk.iter_slices::()) + .flat_map(|chunk| { + // TODO(grtlr): Unless we're on the happy path this allocates and happens everytime the visualizer runs! + chunk.iter_slices::>() + }) .enumerate() .for_each(|(i, values)| { - for (points, value) in points_per_series.iter_mut().zip(values) { + let values_slice = values.as_ref(); + for (points, value) in points_per_series.iter_mut().zip(values_slice) { points[i].value = *value; } - for points in points_per_series.iter_mut().skip(values.len()) { + for points in points_per_series.iter_mut().skip(values_slice.len()) { points[i].attrs.kind = PlotSeriesKind::Clear; } }); diff --git a/crates/viewer/re_view_time_series/tests/blueprint.rs b/crates/viewer/re_view_time_series/tests/blueprint.rs index 3cf48c9fbee9..736b98f946a2 100644 --- a/crates/viewer/re_view_time_series/tests/blueprint.rs +++ b/crates/viewer/re_view_time_series/tests/blueprint.rs @@ -1,10 +1,14 @@ +use std::sync::Arc; + use re_chunk_store::RowId; use re_log_types::{EntityPath, TimePoint}; use re_test_context::TestContext; use re_test_viewport::TestContextExt as _; use re_types::{ + Archetype as _, DynamicArchetype, archetypes::{self, Scalars}, blueprint, components, + external::arrow::array::{Float32Array, Int64Array}, }; use re_view_time_series::TimeSeriesView; use re_viewer_context::{BlueprintContext as _, TimeControlCommand, ViewClass as _, ViewId}; @@ -70,3 +74,67 @@ fn setup_blueprint(test_context: &mut TestContext) -> ViewId { blueprint.add_view_at_root(view) }) } + +// TODO: Move this test to a better place. +#[test] +pub fn test_blueprint_f64_with_time_series() { + let mut test_context = TestContext::new_with_view_class::(); + + let timeline = re_log_types::Timeline::log_tick(); + + for i in 0..32 { + let timepoint = TimePoint::from([(timeline, i)]); + let t = i as f64 / 8.0; + test_context.log_entity("plots/sin", |builder| { + builder.with_archetype(RowId::new(), timepoint.clone(), &Scalars::single(t.sin())) + }); + test_context.log_entity("plots/cos", |builder| { + builder.with_archetype( + RowId::new(), + timepoint.clone(), + // an untagged component + &DynamicArchetype::new(Scalars::name()).with_component_from_data( + "scalars", + Arc::new(Float32Array::from(vec![t.cos() as f32])), + ), + ) + }); + test_context.log_entity("plots/line", |builder| { + builder.with_archetype( + RowId::new(), + timepoint, + // an untagged component + &DynamicArchetype::new(Scalars::name()).with_component_from_data( + "scalars", + // Something that stays in the same domain as a sine wave. + Arc::new(Int64Array::from(vec![(i % 2) * 2 - 1])), + ), + ) + }); + } + + // test_context + // .save_recording_to_file("/Users/goertler/Desktop/dyn_f64.rrd") + // .unwrap(); + + test_context.send_time_commands( + test_context.active_store_id(), + [TimeControlCommand::SetActiveTimeline(*timeline.name())], + ); + + let view_id = setup_descriptor_override_blueprint(&mut test_context); + test_context.run_view_ui_and_save_snapshot( + view_id, + "blueprint_f64_with_time_series", + egui::vec2(300.0, 300.0), + None, + ); +} + +fn setup_descriptor_override_blueprint(test_context: &mut TestContext) -> ViewId { + test_context.setup_viewport_blueprint(|_ctx, blueprint| { + let view = ViewBlueprint::new_with_root_wildcard(TimeSeriesView::identifier()); + + blueprint.add_view_at_root(view) + }) +} diff --git a/crates/viewer/re_view_time_series/tests/snapshots/blueprint_f64_with_time_series.png b/crates/viewer/re_view_time_series/tests/snapshots/blueprint_f64_with_time_series.png new file mode 100644 index 000000000000..bf1c75a0c401 --- /dev/null +++ b/crates/viewer/re_view_time_series/tests/snapshots/blueprint_f64_with_time_series.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:503b538603562fd0bd29610f8f9df20f20baa9fb8399869121a95c969cefb861 +size 63227