Skip to content

Commit

Permalink
Port parts of viewer to arrow-rs (#8534)
Browse files Browse the repository at this point in the history
### Related
* Part if #3741

### TODO
* [x] `@rerun-bot full-check`
  • Loading branch information
emilk authored Dec 19, 2024
1 parent 2c72d87 commit 662e404
Show file tree
Hide file tree
Showing 42 changed files with 694 additions and 618 deletions.
6 changes: 6 additions & 0 deletions Cargo.lock
Original file line number Diff line number Diff line change
Expand Up @@ -5701,6 +5701,7 @@ dependencies = [
name = "re_component_ui"
version = "0.22.0-alpha.1+dev"
dependencies = [
"arrow",
"egui",
"egui_extras",
"egui_plot",
Expand Down Expand Up @@ -6297,6 +6298,7 @@ dependencies = [
name = "re_selection_panel"
version = "0.22.0-alpha.1+dev"
dependencies = [
"arrow",
"egui",
"egui_tiles",
"itertools 0.13.0",
Expand Down Expand Up @@ -6559,6 +6561,7 @@ dependencies = [
"glam",
"itertools 0.13.0",
"nohash-hasher",
"re_arrow2",
"re_chunk_store",
"re_entity_db",
"re_log",
Expand Down Expand Up @@ -6596,6 +6599,7 @@ name = "re_view_dataframe"
version = "0.22.0-alpha.1+dev"
dependencies = [
"anyhow",
"arrow",
"egui",
"egui_table",
"itertools 0.13.0",
Expand Down Expand Up @@ -6872,6 +6876,7 @@ dependencies = [
"ahash",
"anyhow",
"arboard",
"arrow",
"bit-vec",
"bitflags 2.6.0",
"bytemuck",
Expand Down Expand Up @@ -6950,6 +6955,7 @@ name = "re_viewport_blueprint"
version = "0.22.0-alpha.1+dev"
dependencies = [
"ahash",
"arrow",
"egui",
"egui_tiles",
"itertools 0.13.0",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -137,7 +137,7 @@ fn generate_component_reflection(
|| contents.contains(&format!("impl Default for super::{}", &obj.name))
});
let custom_placeholder = if auto_derive_default || has_custom_default_impl {
quote! { Some(#type_name::default().to_arrow2()?) }
quote! { Some(#type_name::default().to_arrow()?) }
} else {
quote! { None }
};
Expand Down
22 changes: 20 additions & 2 deletions crates/store/re_chunk/src/builder.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use arrow::array::ArrayRef;
use arrow2::{
array::{Array as Arrow2Array, PrimitiveArray as Arrow2PrimitiveArray},
datatypes::DataType as Arrow2Datatype,
Expand Down Expand Up @@ -102,6 +103,23 @@ impl ChunkBuilder {
/// Add a row's worth of data using the given component data.
#[inline]
pub fn with_row(
self,
row_id: RowId,
timepoint: impl Into<TimePoint>,
components: impl IntoIterator<Item = (ComponentDescriptor, ArrayRef)>,
) -> Self {
self.with_sparse_row(
row_id,
timepoint,
components
.into_iter()
.map(|(component_descr, array)| (component_descr, Some(array.into()))),
)
}

/// Add a row's worth of data using the given component data.
#[inline]
pub fn with_row_arrow2(
self,
row_id: RowId,
timepoint: impl Into<TimePoint>,
Expand Down Expand Up @@ -142,7 +160,7 @@ impl ChunkBuilder {
timepoint: impl Into<TimePoint>,
component_batch: &dyn ComponentBatch,
) -> Self {
self.with_row(
self.with_row_arrow2(
row_id,
timepoint,
component_batch
Expand All @@ -160,7 +178,7 @@ impl ChunkBuilder {
timepoint: impl Into<TimePoint>,
component_batches: impl IntoIterator<Item = &'a dyn ComponentBatch>,
) -> Self {
self.with_row(
self.with_row_arrow2(
row_id,
timepoint,
component_batches.into_iter().filter_map(|component_batch| {
Expand Down
16 changes: 12 additions & 4 deletions crates/store/re_chunk/src/helpers.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
use std::sync::Arc;

use arrow::array::ArrayRef;
use arrow2::array::Array as Arrow2Array;

use re_log_types::{TimeInt, Timeline};
Expand Down Expand Up @@ -262,7 +263,14 @@ impl UnitChunkShared {

/// Returns the raw data for the specified component.
#[inline]
pub fn component_batch_raw(
pub fn component_batch_raw(&self, component_name: &ComponentName) -> Option<ArrayRef> {
self.component_batch_raw_arrow2(component_name)
.map(|array| array.into())
}

/// Returns the raw data for the specified component.
#[inline]
pub fn component_batch_raw_arrow2(
&self,
component_name: &ComponentName,
) -> Option<Box<dyn Arrow2Array>> {
Expand All @@ -276,7 +284,7 @@ impl UnitChunkShared {
/// Returns an error if the data cannot be deserialized.
#[inline]
pub fn component_batch<C: Component>(&self) -> Option<ChunkResult<Vec<C>>> {
let data = C::from_arrow2(&*self.component_batch_raw(&C::name())?);
let data = C::from_arrow2(&*self.component_batch_raw_arrow2(&C::name())?);
Some(data.map_err(Into::into))
}

Expand All @@ -291,7 +299,7 @@ impl UnitChunkShared {
component_name: &ComponentName,
instance_index: usize,
) -> Option<ChunkResult<Box<dyn Arrow2Array>>> {
let array = self.component_batch_raw(component_name)?;
let array = self.component_batch_raw_arrow2(component_name)?;
if array.len() > instance_index {
Some(Ok(array.sliced(instance_index, 1)))
} else {
Expand Down Expand Up @@ -334,7 +342,7 @@ impl UnitChunkShared {
&self,
component_name: &ComponentName,
) -> Option<ChunkResult<Box<dyn Arrow2Array>>> {
let array = self.component_batch_raw(component_name)?;
let array = self.component_batch_raw_arrow2(component_name)?;
if array.len() == 1 {
Some(Ok(array.sliced(0, 1)))
} else {
Expand Down
1 change: 1 addition & 0 deletions crates/store/re_chunk/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ pub use re_log_types::{EntityPath, TimeInt, TimePoint, Timeline, TimelineName};
pub use re_types_core::{ArchetypeFieldName, ArchetypeName, ComponentName};

pub mod external {
pub use arrow;
pub use arrow2;
pub use nohash_hasher;

Expand Down
4 changes: 2 additions & 2 deletions crates/store/re_chunk/src/merge.rs
Original file line number Diff line number Diff line change
Expand Up @@ -804,7 +804,7 @@ mod tests {
<MyPoint64 as re_types_core::LoggableBatch>::to_arrow2(&MyPoint64::new(1.0, 1.0))?;

let chunk1 = Chunk::builder(entity_path.into())
.with_row(
.with_row_arrow2(
row_id1,
timepoint1,
[
Expand All @@ -814,7 +814,7 @@ mod tests {
.build()?;

let chunk2 = Chunk::builder(entity_path.into())
.with_row(
.with_row_arrow2(
row_id2,
timepoint2,
[
Expand Down
2 changes: 1 addition & 1 deletion crates/store/re_chunk_store/tests/gc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ fn query_latest_array(
})
.max_by_key(|(index, _chunk)| *index)?;

unit.component_batch_raw(&component_name)
unit.component_batch_raw_arrow2(&component_name)
.map(|array| (data_time, row_id, array))
}

Expand Down
2 changes: 1 addition & 1 deletion crates/store/re_chunk_store/tests/reads.rs
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ fn query_latest_array(
})
.max_by_key(|(index, _chunk)| *index)?;

unit.component_batch_raw(&component_desc.component_name)
unit.component_batch_raw_arrow2(&component_desc.component_name)
.map(|array| (data_time, row_id, array))
}

Expand Down
2 changes: 1 addition & 1 deletion crates/store/re_query/examples/latest_at.rs
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ fn main() -> anyhow::Result<()> {
// data directly:
let colors = colors
.context("missing")?
.component_batch_raw(&MyColor::name())
.component_batch_raw_arrow2(&MyColor::name())
.context("invalid")?;
let colors = colors
.as_any()
Expand Down
15 changes: 12 additions & 3 deletions crates/store/re_query/src/latest_at.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,8 @@ use re_chunk::{Chunk, RowId, UnitChunkShared};
use re_chunk_store::{ChunkStore, LatestAtQuery, TimeInt};
use re_log_types::EntityPath;
use re_types_core::{
components::ClearIsRecursive, Component, ComponentDescriptor, ComponentName, SizeBytes,
components::ClearIsRecursive, external::arrow::array::ArrayRef, Component, ComponentDescriptor,
ComponentName, SizeBytes,
};

use crate::{QueryCache, QueryCacheKey, QueryError};
Expand Down Expand Up @@ -313,13 +314,21 @@ impl LatestAtResults {

/// Returns the raw data for the specified component.
#[inline]
pub fn component_batch_raw(
pub fn component_batch_raw(&self, component_name: &ComponentName) -> Option<ArrayRef> {
self.components
.get(component_name)?
.component_batch_raw(component_name)
}

/// Returns the raw data for the specified component.
#[inline]
pub fn component_batch_raw_arrow2(
&self,
component_name: &ComponentName,
) -> Option<Box<dyn Arrow2Array>> {
self.components
.get(component_name)
.and_then(|unit| unit.component_batch_raw(component_name))
.and_then(|unit| unit.component_batch_raw_arrow2(component_name))
}

/// Returns the deserialized data for the specified component.
Expand Down
1 change: 1 addition & 0 deletions crates/store/re_types/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -277,6 +277,7 @@ pub mod external {
pub use re_types_core;

pub use anyhow;
pub use arrow;
pub use arrow2;
pub use ndarray;
pub use uuid;
Expand Down
Loading

0 comments on commit 662e404

Please sign in to comment.