Skip to content

Commit

Permalink
Update tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@emilk
emilk committed Jan 8, 2025
1 parent 288552d commit 55df031
Show file tree
Hide file tree
Showing 2 changed files with 61 additions and 36 deletions.
70 changes: 34 additions & 36 deletions crates/store/re_chunk/src/batcher.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -996,8 +996,6 @@ mod tests {
use re_log_types::example_components::{MyPoint, MyPoint64};
use re_types_core::{Component as _, Loggable as _};

use crate::arrow2_util;

use super::*;

/// A bunch of rows that don't fit any of the split conditions should end up together.
Expand All @@ -1011,9 +1009,9 @@ mod tests {
let timepoint2 = TimePoint::default().with(timeline1, 43);
let timepoint3 = TimePoint::default().with(timeline1, 44);

let points1 = MyPoint::to_arrow2([MyPoint::new(1.0, 2.0), MyPoint::new(3.0, 4.0)])?;
let points2 = MyPoint::to_arrow2([MyPoint::new(10.0, 20.0), MyPoint::new(30.0, 40.0)])?;
let points3 = MyPoint::to_arrow2([MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)])?;
let points1 = MyPoint::to_arrow([MyPoint::new(1.0, 2.0), MyPoint::new(3.0, 4.0)])?;
let points2 = MyPoint::to_arrow([MyPoint::new(10.0, 20.0), MyPoint::new(30.0, 40.0)])?;
let points3 = MyPoint::to_arrow([MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)])?;

let components1 = [(MyPoint::descriptor(), points1.clone())];
let components2 = [(MyPoint::descriptor(), points2.clone())];
Expand Down Expand Up @@ -1063,7 +1061,7 @@ mod tests {
)];
let expected_components = [(
MyPoint::descriptor(),
arrow2_util::arrays_to_list_array_opt(&[&*points1, &*points2, &*points3].map(Some))
arrow_util::arrays_to_list_array_opt(&[&*points1, &*points2, &*points3].map(Some))
.unwrap(),
)];
let expected_chunk = Chunk::from_native_row_ids(
Expand All @@ -1090,9 +1088,9 @@ mod tests {

let timeless = TimePoint::default();

let points1 = MyPoint::to_arrow2([MyPoint::new(1.0, 2.0), MyPoint::new(3.0, 4.0)])?;
let points2 = MyPoint::to_arrow2([MyPoint::new(10.0, 20.0), MyPoint::new(30.0, 40.0)])?;
let points3 = MyPoint::to_arrow2([MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)])?;
let points1 = MyPoint::to_arrow([MyPoint::new(1.0, 2.0), MyPoint::new(3.0, 4.0)])?;
let points2 = MyPoint::to_arrow([MyPoint::new(10.0, 20.0), MyPoint::new(30.0, 40.0)])?;
let points3 = MyPoint::to_arrow([MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)])?;

let components1 = [(MyPoint::descriptor(), points1.clone())];
let components2 = [(MyPoint::descriptor(), points2.clone())];
Expand Down Expand Up @@ -1135,7 +1133,7 @@ mod tests {
let expected_timelines = [];
let expected_components = [(
MyPoint::descriptor(),
arrow2_util::arrays_to_list_array_opt(&[&*points1, &*points2, &*points3].map(Some))
arrow_util::arrays_to_list_array_opt(&[&*points1, &*points2, &*points3].map(Some))
.unwrap(),
)];
let expected_chunk = Chunk::from_native_row_ids(
Expand Down Expand Up @@ -1166,9 +1164,9 @@ mod tests {
let timepoint2 = TimePoint::default().with(timeline1, 43);
let timepoint3 = TimePoint::default().with(timeline1, 44);

let points1 = MyPoint::to_arrow2([MyPoint::new(1.0, 2.0), MyPoint::new(3.0, 4.0)])?;
let points2 = MyPoint::to_arrow2([MyPoint::new(10.0, 20.0), MyPoint::new(30.0, 40.0)])?;
let points3 = MyPoint::to_arrow2([MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)])?;
let points1 = MyPoint::to_arrow([MyPoint::new(1.0, 2.0), MyPoint::new(3.0, 4.0)])?;
let points2 = MyPoint::to_arrow([MyPoint::new(10.0, 20.0), MyPoint::new(30.0, 40.0)])?;
let points3 = MyPoint::to_arrow([MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)])?;

let components1 = [(MyPoint::descriptor(), points1.clone())];
let components2 = [(MyPoint::descriptor(), points2.clone())];
Expand Down Expand Up @@ -1219,7 +1217,7 @@ mod tests {
)];
let expected_components = [(
MyPoint::descriptor(),
arrow2_util::arrays_to_list_array_opt(&[&*points1, &*points3].map(Some)).unwrap(),
arrow_util::arrays_to_list_array_opt(&[&*points1, &*points3].map(Some)).unwrap(),
)];
let expected_chunk = Chunk::from_native_row_ids(
chunks[0].id,
Expand Down Expand Up @@ -1247,7 +1245,7 @@ mod tests {
)];
let expected_components = [(
MyPoint::descriptor(),
arrow2_util::arrays_to_list_array_opt(&[&*points2].map(Some)).unwrap(),
arrow_util::arrays_to_list_array_opt(&[&*points2].map(Some)).unwrap(),
)];
let expected_chunk = Chunk::from_native_row_ids(
chunks[1].id,
Expand Down Expand Up @@ -1282,9 +1280,9 @@ mod tests {
.with(timeline1, 44)
.with(timeline2, 1001);

let points1 = MyPoint::to_arrow2([MyPoint::new(1.0, 2.0), MyPoint::new(3.0, 4.0)])?;
let points2 = MyPoint::to_arrow2([MyPoint::new(10.0, 20.0), MyPoint::new(30.0, 40.0)])?;
let points3 = MyPoint::to_arrow2([MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)])?;
let points1 = MyPoint::to_arrow([MyPoint::new(1.0, 2.0), MyPoint::new(3.0, 4.0)])?;
let points2 = MyPoint::to_arrow([MyPoint::new(10.0, 20.0), MyPoint::new(30.0, 40.0)])?;
let points3 = MyPoint::to_arrow([MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)])?;

let components1 = [(MyPoint::descriptor(), points1.clone())];
let components2 = [(MyPoint::descriptor(), points2.clone())];
Expand Down Expand Up @@ -1334,7 +1332,7 @@ mod tests {
)];
let expected_components = [(
MyPoint::descriptor(),
arrow2_util::arrays_to_list_array_opt(&[&*points1].map(Some)).unwrap(),
arrow_util::arrays_to_list_array_opt(&[&*points1].map(Some)).unwrap(),
)];
let expected_chunk = Chunk::from_native_row_ids(
chunks[0].id,
Expand Down Expand Up @@ -1372,7 +1370,7 @@ mod tests {
];
let expected_components = [(
MyPoint::descriptor(),
arrow2_util::arrays_to_list_array_opt(&[&*points2, &*points3].map(Some)).unwrap(),
arrow_util::arrays_to_list_array_opt(&[&*points2, &*points3].map(Some)).unwrap(),
)];
let expected_chunk = Chunk::from_native_row_ids(
chunks[1].id,
Expand Down Expand Up @@ -1402,10 +1400,10 @@ mod tests {
let timepoint2 = TimePoint::default().with(timeline1, 43);
let timepoint3 = TimePoint::default().with(timeline1, 44);

let points1 = MyPoint::to_arrow2([MyPoint::new(1.0, 2.0), MyPoint::new(3.0, 4.0)])?;
let points1 = MyPoint::to_arrow([MyPoint::new(1.0, 2.0), MyPoint::new(3.0, 4.0)])?;
let points2 =
MyPoint64::to_arrow2([MyPoint64::new(10.0, 20.0), MyPoint64::new(30.0, 40.0)])?;
let points3 = MyPoint::to_arrow2([MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)])?;
MyPoint64::to_arrow([MyPoint64::new(10.0, 20.0), MyPoint64::new(30.0, 40.0)])?;
let points3 = MyPoint::to_arrow([MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)])?;

let components1 = [(MyPoint::descriptor(), points1.clone())];
let components2 = [(MyPoint::descriptor(), points2.clone())]; // same name, different datatype
Expand Down Expand Up @@ -1455,7 +1453,7 @@ mod tests {
)];
let expected_components = [(
MyPoint::descriptor(),
arrow2_util::arrays_to_list_array_opt(&[&*points1, &*points3].map(Some)).unwrap(),
arrow_util::arrays_to_list_array_opt(&[&*points1, &*points3].map(Some)).unwrap(),
)];
let expected_chunk = Chunk::from_native_row_ids(
chunks[0].id,
Expand Down Expand Up @@ -1483,7 +1481,7 @@ mod tests {
)];
let expected_components = [(
MyPoint::descriptor(),
arrow2_util::arrays_to_list_array_opt(&[&*points2].map(Some)).unwrap(),
arrow_util::arrays_to_list_array_opt(&[&*points2].map(Some)).unwrap(),
)];
let expected_chunk = Chunk::from_native_row_ids(
chunks[1].id,
Expand Down Expand Up @@ -1527,11 +1525,11 @@ mod tests {
.with(timeline2, 1003)
.with(timeline1, 45);

let points1 = MyPoint::to_arrow2([MyPoint::new(1.0, 2.0), MyPoint::new(3.0, 4.0)])?;
let points2 = MyPoint::to_arrow2([MyPoint::new(10.0, 20.0), MyPoint::new(30.0, 40.0)])?;
let points3 = MyPoint::to_arrow2([MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)])?;
let points1 = MyPoint::to_arrow([MyPoint::new(1.0, 2.0), MyPoint::new(3.0, 4.0)])?;
let points2 = MyPoint::to_arrow([MyPoint::new(10.0, 20.0), MyPoint::new(30.0, 40.0)])?;
let points3 = MyPoint::to_arrow([MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)])?;
let points4 =
MyPoint::to_arrow2([MyPoint::new(1000.0, 2000.0), MyPoint::new(3000.0, 4000.0)])?;
MyPoint::to_arrow([MyPoint::new(1000.0, 2000.0), MyPoint::new(3000.0, 4000.0)])?;

let components1 = [(MyPoint::descriptor(), points1.clone())];
let components2 = [(MyPoint::descriptor(), points2.clone())];
Expand Down Expand Up @@ -1594,7 +1592,7 @@ mod tests {
];
let expected_components = [(
MyPoint::descriptor(),
arrow2_util::arrays_to_list_array_opt(
arrow_util::arrays_to_list_array_opt(
&[&*points1, &*points2, &*points3, &*points4].map(Some),
)
.unwrap(),
Expand Down Expand Up @@ -1641,11 +1639,11 @@ mod tests {
.with(timeline2, 1003)
.with(timeline1, 45);

let points1 = MyPoint::to_arrow2([MyPoint::new(1.0, 2.0), MyPoint::new(3.0, 4.0)])?;
let points2 = MyPoint::to_arrow2([MyPoint::new(10.0, 20.0), MyPoint::new(30.0, 40.0)])?;
let points3 = MyPoint::to_arrow2([MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)])?;
let points1 = MyPoint::to_arrow([MyPoint::new(1.0, 2.0), MyPoint::new(3.0, 4.0)])?;
let points2 = MyPoint::to_arrow([MyPoint::new(10.0, 20.0), MyPoint::new(30.0, 40.0)])?;
let points3 = MyPoint::to_arrow([MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)])?;
let points4 =
MyPoint::to_arrow2([MyPoint::new(1000.0, 2000.0), MyPoint::new(3000.0, 4000.0)])?;
MyPoint::to_arrow([MyPoint::new(1000.0, 2000.0), MyPoint::new(3000.0, 4000.0)])?;

let components1 = [(MyPoint::descriptor(), points1.clone())];
let components2 = [(MyPoint::descriptor(), points2.clone())];
Expand Down Expand Up @@ -1708,7 +1706,7 @@ mod tests {
];
let expected_components = [(
MyPoint::descriptor(),
arrow2_util::arrays_to_list_array_opt(&[&*points1, &*points2, &*points3].map(Some))
arrow_util::arrays_to_list_array_opt(&[&*points1, &*points2, &*points3].map(Some))
.unwrap(),
)];
let expected_chunk = Chunk::from_native_row_ids(
Expand Down Expand Up @@ -1747,7 +1745,7 @@ mod tests {
];
let expected_components = [(
MyPoint::descriptor(),
arrow2_util::arrays_to_list_array_opt(&[&*points4].map(Some)).unwrap(),
arrow_util::arrays_to_list_array_opt(&[&*points4].map(Some)).unwrap(),
)];
let expected_chunk = Chunk::from_native_row_ids(
chunks[1].id,
Expand Down
27 changes: 27 additions & 0 deletions crates/store/re_chunk/src/chunk.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -155,6 +155,33 @@ impl std::ops::DerefMut for ChunkComponents {
}
}

impl FromIterator<(ComponentDescriptor, ArrowListArray)> for ChunkComponents {
#[inline]
fn from_iter<T: IntoIterator<Item = (ComponentDescriptor, ArrowListArray)>>(iter: T) -> Self {
let mut this = Self::default();
{
for (component_desc, list_array) in iter {
this.insert_descriptor(component_desc, list_array.into());
}
}
this
}
}

// TODO(cmc): Kinda disgusting but it makes our lives easier during the interim, as long as we're
// in this weird halfway in-between state where we still have a bunch of things indexed by name only.
impl FromIterator<(ComponentName, ArrowListArray)> for ChunkComponents {
#[inline]
fn from_iter<T: IntoIterator<Item = (ComponentName, ArrowListArray)>>(iter: T) -> Self {
iter.into_iter()
.map(|(component_name, list_array)| {
let component_desc = ComponentDescriptor::new(component_name);
(component_desc, list_array)
})
.collect()
}
}

impl FromIterator<(ComponentDescriptor, Arrow2ListArray<i32>)> for ChunkComponents {
#[inline]
fn from_iter<T: IntoIterator<Item = (ComponentDescriptor, Arrow2ListArray<i32>)>>(
Expand Down

0 comments on commit 55df031

Please sign in to comment.