use std::sync::Arc;
use nohash_hasher::{IntMap, IntSet};
use re_chunk_store::{LatestAtQuery, RangeQuery, RowId};
use re_log_types::{
TimeInt,
external::arrow::{self, array::Array as _},
hash::Hash64,
};
use re_query::LatestAtResults;
use re_sdk_types::blueprint::datatypes::ComponentSourceKind;
use re_types_core::{Archetype, ComponentIdentifier};
use re_viewer_context::{DataResult, QueryRange, ViewContext, ViewQuery, ViewerContext};
use crate::blueprint_resolved_results::{
BlueprintResolvedLatestAtResults, BlueprintResolvedRangeResults, ComponentSourcesMap,
};
use crate::{BlueprintResolvedResults, ComponentMappingError};
pub type ComponentCastRule = fn(&arrow::datatypes::DataType) -> Option<arrow::datatypes::DataType>;
fn cast_list_array(
source: &arrow::array::ListArray,
target_list_datatype: &arrow::datatypes::DataType,
) -> Result<arrow::array::ListArray, arrow::error::ArrowError> {
if source.data_type() == target_list_datatype {
return Ok(source.clone());
}
let casted = arrow::compute::cast(source, target_list_datatype)?;
casted
.as_any()
.downcast_ref::<arrow::array::ListArray>()
.cloned()
.ok_or_else(|| {
arrow::error::ArrowError::CastError(format!(
"Expected ListArray after cast, got {:?}",
casted.data_type()
))
})
}
enum CastTarget {
Fixed(Option<arrow::datatypes::DataType>),
Polymorphic(ComponentCastRule),
}
fn transform_chunk(
target: &ComponentIdentifier,
source: &ComponentIdentifier,
selector: Option<&re_lenses_core::Selector>,
cast: &CastTarget,
chunk: &re_chunk_store::Chunk,
) -> Result<re_chunk_store::Chunk, ComponentMappingError> {
chunk.with_shadowed_component(*source, *target, |arr| {
let transformed = if let Some(sel) = selector {
sel.execute_per_row(&arr)
.map_err(ComponentMappingError::SelectorExecutionFailed)?
.unwrap_or_else(|| {
arrow::array::ListArray::new_null(
arrow::datatypes::Field::new_list_field(arr.value_type(), true).into(),
arr.len(),
)
})
} else {
arr
};
let target_datatype = match cast {
CastTarget::Polymorphic(rule) => rule(&transformed.value_type()),
CastTarget::Fixed(dt) => dt.clone(),
};
if let Some(dt) = target_datatype {
let target_list_datatype = arrow::datatypes::DataType::List(Arc::new(
arrow::datatypes::Field::new_list_field(dt.clone(), true),
));
cast_list_array(&transformed, &target_list_datatype).map_err(|err| {
ComponentMappingError::CastFailed {
source_datatype: transformed.data_type().clone(),
target_datatype: target_list_datatype,
err: Arc::new(err),
}
})
} else {
Ok(transformed)
}
})
}
#[derive(Debug, PartialEq, Eq, Hash)]
struct ActiveRemapping {
target: ComponentIdentifier,
source: ComponentIdentifier,
selector: Option<re_lenses_core::Selector>,
}
fn cast_target_for_remapping(
rule: Option<ComponentCastRule>,
target: &ComponentIdentifier,
reflection: &re_types_core::reflection::Reflection,
) -> CastTarget {
match rule {
Some(rule) => CastTarget::Polymorphic(rule),
None => CastTarget::Fixed(reflection.lookup_datatype(*target).cloned()),
}
}
fn component_not_found_error(
component: ComponentIdentifier,
entity_path: &re_log_types::EntityPath,
missing_virtual_chunks: &[re_chunk_store::ChunkId],
entity_db: &re_entity_db::EntityDb,
store_engine: &re_query::StorageEngineReadGuard<'_>,
timeline_name: Option<re_log_types::TimelineName>,
) -> ComponentMappingError {
if timeline_name.is_some_and(|timeline_name| {
entity_db.entity_has_temporal_data_on_timeline_for_component(
store_engine,
&timeline_name,
entity_path,
component,
)
}) {
ComponentMappingError::NoComponentDataForQuery(component)
} else {
if !missing_virtual_chunks.is_empty()
&& let Some(rrd_manifest) = entity_db.rrd_manifest_index().manifest()
{
let store = store_engine.store();
let timeline = timeline_name
.and_then(|timeline_name| store.schema().timelines().get(&timeline_name).copied());
for missing_root_chunk_id in missing_virtual_chunks
.iter()
.flat_map(|chunk_id| store.find_root_chunks(chunk_id))
{
if let Some(per_component) = rrd_manifest.static_map().get(entity_path)
&& per_component.get(&component) == Some(&missing_root_chunk_id)
{
return ComponentMappingError::NoComponentDataForQueryButIsFetchable(component);
}
if let Some(timeline) = &timeline
&& let Some(per_timeline) = rrd_manifest.temporal_map().get(entity_path)
&& let Some(per_component) = per_timeline.get(timeline)
&& let Some(per_chunk) = per_component.get(&component)
&& per_chunk.contains_key(&missing_root_chunk_id)
{
return ComponentMappingError::NoComponentDataForQueryButIsFetchable(component);
}
}
}
ComponentMappingError::ComponentNotPresentOnEntity(component)
}
}
pub fn range_with_blueprint_resolved_data<'a>(
ctx: &'a ViewContext<'a>,
annotations: Option<&re_viewer_context::Annotations>,
range_query: &RangeQuery,
data_result: &'a re_viewer_context::DataResult,
components: impl IntoIterator<Item = ComponentIdentifier>,
visualizer_instruction: &re_viewer_context::VisualizerInstruction,
) -> BlueprintResolvedRangeResults<'a> {
range_with_blueprint_resolved_data_polymorphic(
ctx,
annotations,
range_query,
data_result,
components,
visualizer_instruction,
&IntMap::default(),
)
}
pub fn range_with_blueprint_resolved_data_polymorphic<'a>(
ctx: &'a ViewContext<'a>,
_annotations: Option<&re_viewer_context::Annotations>,
range_query: &RangeQuery,
data_result: &'a re_viewer_context::DataResult,
components: impl IntoIterator<Item = ComponentIdentifier>,
visualizer_instruction: &re_viewer_context::VisualizerInstruction,
cast_rules: &IntMap<ComponentIdentifier, ComponentCastRule>,
) -> BlueprintResolvedRangeResults<'a> {
re_tracing::profile_function!(data_result.entity_path.to_string());
let mut queried_components = components.into_iter().collect::<IntSet<_>>();
let overrides = query_overrides(
ctx.viewer_ctx,
visualizer_instruction,
queried_components.iter().copied(),
);
let mut active_remappings = Vec::new();
let mut component_sources = IntMap::default();
let store_results = {
for (target_component, source) in &visualizer_instruction.component_mappings {
let source = if let re_viewer_context::VisualizerComponentSource::SourceComponent {
source_component,
selector,
} = source
&& queried_components.remove(target_component)
{
queried_components.insert(*source_component);
if selector.is_empty() {
active_remappings.push(ActiveRemapping {
target: *target_component,
source: *source_component,
selector: None,
});
Ok(source.source_kind())
} else {
match selector.parse::<re_lenses_core::Selector>() {
Ok(selector) => {
active_remappings.push(ActiveRemapping {
target: *target_component,
source: *source_component,
selector: Some(selector),
});
Ok(source.source_kind())
}
Err(err) => Err(ComponentMappingError::SelectorParseFailed(err)),
}
}
} else {
Ok(source.source_kind())
};
component_sources.insert(*target_component, source);
}
let engine = ctx.recording_engine();
let mut results = engine.cache().range(
re_chunk_store::ChunkTrackingMode::Report,
range_query,
&data_result.entity_path,
queried_components.iter().copied(),
);
let reflection = ctx.viewer_ctx.reflection();
for ActiveRemapping {
target,
source,
selector,
} in &active_remappings
{
let cast =
cast_target_for_remapping(cast_rules.get(target).copied(), target, reflection);
if let Some(mut chunks) = results.components.get(source).cloned() {
'ctx: {
for chunk in &mut chunks {
let result =
transform_chunk(target, source, selector.as_ref(), &cast, chunk);
match result {
Ok(modified_chunk) => *chunk = modified_chunk,
Err(err) => {
component_sources.insert(*target, Err(err));
break 'ctx;
}
}
}
results.components.insert(*target, chunks);
}
} else {
component_sources.insert(
*target,
Err(component_not_found_error(
*source,
&data_result.entity_path,
&results.missing_virtual,
ctx.recording(),
&engine,
Some(range_query.timeline),
)),
);
}
}
results
};
auto_determine_remaining_sources(
&mut component_sources,
queried_components,
|component| store_results.components.contains_key(&component),
&overrides,
);
let query_context = ctx.query_context(
data_result,
LatestAtQuery::new(range_query.timeline, range_query.range.min),
visualizer_instruction.id,
);
BlueprintResolvedRangeResults {
overrides,
store_results,
query_context,
view_defaults: &ctx.query_result.view_defaults,
component_sources,
component_mappings_hash: Hash64::hash(&active_remappings),
}
}
pub fn latest_at_with_blueprint_resolved_data<'a>(
ctx: &'a ViewContext<'a>,
annotations: Option<&'a re_viewer_context::Annotations>,
latest_at_query: &LatestAtQuery,
data_result: &'a re_viewer_context::DataResult,
components: impl IntoIterator<Item = ComponentIdentifier>,
visualizer_instruction: Option<&re_viewer_context::VisualizerInstruction>,
) -> BlueprintResolvedLatestAtResults<'a> {
latest_at_with_blueprint_resolved_data_polymorphic(
ctx,
annotations,
latest_at_query,
data_result,
components,
visualizer_instruction,
&IntMap::default(),
)
}
pub fn latest_at_with_blueprint_resolved_data_polymorphic<'a>(
ctx: &'a ViewContext<'a>,
_annotations: Option<&'a re_viewer_context::Annotations>,
latest_at_query: &LatestAtQuery,
data_result: &'a re_viewer_context::DataResult,
components: impl IntoIterator<Item = ComponentIdentifier>,
visualizer_instruction: Option<&re_viewer_context::VisualizerInstruction>,
cast_rules: &IntMap<ComponentIdentifier, ComponentCastRule>,
) -> BlueprintResolvedLatestAtResults<'a> {
let mut queried_components = components.into_iter().collect::<IntSet<_>>();
let overrides = if let Some(visualizer_instruction) = visualizer_instruction {
query_overrides(
ctx.viewer_ctx,
visualizer_instruction,
queried_components.iter().copied(),
)
} else {
query_overrides_at_path(
ctx.viewer_ctx,
data_result.override_base_path(),
queried_components.iter().copied(),
)
};
let mut active_remappings = Vec::new();
let mut component_sources = IntMap::default();
if let Some(visualizer_instruction) = visualizer_instruction {
for (target_component, source) in &visualizer_instruction.component_mappings {
let source_result =
if let re_viewer_context::VisualizerComponentSource::SourceComponent {
source_component,
selector,
} = source
&& queried_components.remove(target_component)
{
queried_components.insert(*source_component);
if selector.is_empty() {
active_remappings.push(ActiveRemapping {
target: *target_component,
source: *source_component,
selector: None,
});
Ok(source.source_kind())
} else {
match selector.parse::<re_lenses_core::Selector>() {
Ok(selector) => {
active_remappings.push(ActiveRemapping {
target: *target_component,
source: *source_component,
selector: Some(selector),
});
Ok(source.source_kind())
}
Err(err) => Err(ComponentMappingError::SelectorParseFailed(err)),
}
}
} else {
Ok(source.source_kind())
};
component_sources.insert(*target_component, source_result);
}
}
let engine = ctx.viewer_ctx.recording_engine();
let mut store_results = engine.cache().latest_at(
re_chunk_store::ChunkTrackingMode::Report,
latest_at_query,
&data_result.entity_path,
queried_components.iter().copied(),
);
let reflection = ctx.viewer_ctx.reflection();
for ActiveRemapping {
target,
source,
selector,
} in &active_remappings
{
let cast = cast_target_for_remapping(cast_rules.get(target).copied(), target, reflection);
if let Some(chunk) = store_results.components.get(source) {
let result = transform_chunk(target, source, selector.as_ref(), &cast, chunk);
match result {
Ok(modified_chunk) => {
let chunk = std::sync::Arc::new(modified_chunk)
.to_unit()
.expect("The source chunk was a unit chunk.");
store_results.components.insert(*target, chunk);
}
Err(err) => {
component_sources.insert(*target, Err(err));
}
}
} else {
component_sources.insert(
*target,
Err(component_not_found_error(
*source,
&data_result.entity_path,
&store_results.missing_virtual,
ctx.viewer_ctx.recording(),
&engine,
latest_at_query.timeline(),
)),
);
}
}
auto_determine_remaining_sources(
&mut component_sources,
queried_components,
|component| store_results.components.contains_key(&component),
&overrides,
);
let query_context = ctx.query_context(
data_result,
latest_at_query.clone(),
visualizer_instruction.map(|instruction| instruction.id),
);
BlueprintResolvedLatestAtResults {
overrides,
store_results,
view_defaults: &ctx.query_result.view_defaults,
query_context,
component_sources,
component_indices_hash: Hash64::hash(&active_remappings),
}
}
fn auto_determine_remaining_sources(
component_sources: &mut ComponentSourcesMap,
queried_components: IntSet<ComponentIdentifier>,
has_store_result: impl Fn(ComponentIdentifier) -> bool,
overrides: &LatestAtResults,
) {
#[expect(clippy::iter_over_hash_type)] for component in queried_components {
let std::collections::hash_map::Entry::Vacant(entry) = component_sources.entry(component)
else {
continue;
};
let source = if has_non_empty_override(overrides, component) {
ComponentSourceKind::Override
} else if has_store_result(component) {
ComponentSourceKind::SourceComponent
} else {
ComponentSourceKind::Default
};
entry.insert(Ok(source));
}
}
fn has_non_empty_override(overrides: &LatestAtResults, component: ComponentIdentifier) -> bool {
overrides
.get(component)
.and_then(|chunk| chunk.non_empty_component_batch_raw(component))
.is_some()
}
fn query_overrides(
ctx: &ViewerContext<'_>,
visualizer_instruction: &re_viewer_context::VisualizerInstruction,
components: impl IntoIterator<Item = ComponentIdentifier>,
) -> LatestAtResults {
if visualizer_instruction.component_overrides.is_empty() {
LatestAtResults::empty("<overrides>".into(), ctx.current_query())
} else {
query_overrides_at_path(
ctx,
&visualizer_instruction.override_path,
components
.into_iter()
.filter(|c| visualizer_instruction.component_overrides.contains(c)),
)
}
}
fn query_overrides_at_path(
ctx: &ViewerContext<'_>,
blueprint_path: &re_log_types::EntityPath,
components: impl IntoIterator<Item = ComponentIdentifier>,
) -> LatestAtResults {
let mut overrides = LatestAtResults::empty("<overrides>".into(), ctx.current_query());
let blueprint_engine = &ctx.store_context.blueprint.storage_engine();
for component in components {
let component_override_result = blueprint_engine.cache().latest_at(
re_chunk_store::ChunkTrackingMode::Report,
ctx.blueprint_query,
blueprint_path,
[component],
);
if let Some(value) = component_override_result.get(component) {
let index = value.index(ctx.blueprint_query.timeline().as_ref());
re_log::debug_assert!(index.is_some(), "{value:#?}");
let index = index.unwrap_or((TimeInt::STATIC, RowId::ZERO));
overrides.add(component, index, value.clone());
}
}
overrides
}
pub trait DataResultQuery {
fn latest_at_with_blueprint_resolved_data<'a, A: re_types_core::Archetype>(
&'a self,
ctx: &'a ViewContext<'a>,
latest_at_query: &'a LatestAtQuery,
visualizer_instruction: Option<&re_viewer_context::VisualizerInstruction>,
) -> BlueprintResolvedLatestAtResults<'a>;
fn latest_at_with_blueprint_resolved_data_for_component<'a>(
&'a self,
ctx: &'a ViewContext<'a>,
latest_at_query: &'a LatestAtQuery,
component: ComponentIdentifier,
visualizer_instruction: Option<&re_viewer_context::VisualizerInstruction>,
) -> BlueprintResolvedLatestAtResults<'a>;
fn query_components_with_history<'a>(
&'a self,
ctx: &'a ViewContext<'a>,
view_query: &ViewQuery<'_>,
component_descriptors: impl IntoIterator<Item = ComponentIdentifier>,
visualizer_instruction: &re_viewer_context::VisualizerInstruction,
) -> BlueprintResolvedResults<'a>;
fn query_archetype_with_history<'a, A: Archetype>(
&'a self,
ctx: &'a ViewContext<'a>,
view_query: &ViewQuery<'_>,
visualizer_instruction: &re_viewer_context::VisualizerInstruction,
) -> BlueprintResolvedResults<'a> {
self.query_components_with_history(
ctx,
view_query,
A::all_component_identifiers(),
visualizer_instruction,
)
}
}
impl DataResultQuery for DataResult {
fn latest_at_with_blueprint_resolved_data<'a, A: re_types_core::Archetype>(
&'a self,
ctx: &'a ViewContext<'a>,
latest_at_query: &'a LatestAtQuery,
visualizer_instruction: Option<&re_viewer_context::VisualizerInstruction>,
) -> BlueprintResolvedLatestAtResults<'a> {
latest_at_with_blueprint_resolved_data(
ctx,
None,
latest_at_query,
self,
A::all_component_identifiers(),
visualizer_instruction,
)
}
fn latest_at_with_blueprint_resolved_data_for_component<'a>(
&'a self,
ctx: &'a ViewContext<'a>,
latest_at_query: &'a LatestAtQuery,
component: ComponentIdentifier,
visualizer_instruction: Option<&re_viewer_context::VisualizerInstruction>,
) -> BlueprintResolvedLatestAtResults<'a> {
latest_at_with_blueprint_resolved_data(
ctx,
None,
latest_at_query,
self,
std::iter::once(component),
visualizer_instruction,
)
}
fn query_components_with_history<'a>(
&'a self,
ctx: &'a ViewContext<'a>,
view_query: &ViewQuery<'_>,
components: impl IntoIterator<Item = ComponentIdentifier>,
visualizer_instruction: &re_viewer_context::VisualizerInstruction,
) -> BlueprintResolvedResults<'a> {
match self.query_range() {
QueryRange::TimeRange(time_range) => {
let range_query = RangeQuery::new(
view_query.timeline,
re_log_types::AbsoluteTimeRange::from_relative_time_range(
time_range,
view_query.latest_at,
),
);
let results = range_with_blueprint_resolved_data(
ctx,
None,
&range_query,
self,
components,
visualizer_instruction,
);
(range_query, results).into()
}
QueryRange::LatestAt => {
let latest_query = LatestAtQuery::new(view_query.timeline, view_query.latest_at);
let results = latest_at_with_blueprint_resolved_data(
ctx,
None,
&latest_query,
self,
components,
Some(visualizer_instruction),
);
(latest_query, results).into()
}
}
}
}