1use std::hash::{Hash as _, Hasher as _};
2use std::sync::Arc;
3use std::time::{Duration, Instant};
4
5use arrow::buffer::ScalarBuffer as ArrowScalarBuffer;
6use nohash_hasher::IntMap;
7use re_arrow_util::arrays_to_list_array_opt;
8use re_byte_size::SizeBytes as _;
9use re_log_types::{AbsoluteTimeRange, EntityPath, TimeInt, TimePoint, Timeline, TimelineName};
10use re_quota_channel::{Receiver, Sender};
11use re_types_core::{ComponentIdentifier, SerializedComponentBatch, SerializedComponentColumn};
12
13use crate::chunk::ChunkComponents;
14use crate::{Chunk, ChunkId, ChunkResult, RowId, TimeColumn};
15
16#[derive(Debug, thiserror::Error)]
20pub enum BatcherFlushError {
21 #[error("Batcher stopped before flushing completed")]
22 Closed,
23
24 #[error("Batcher flush timed out - not all messages were sent.")]
25 Timeout,
26}
27
28#[derive(thiserror::Error, Debug)]
30pub enum ChunkBatcherError {
31 #[error("Failed to parse config: '{name}={value}': {err}")]
33 ParseConfig {
34 name: &'static str,
35 value: String,
36 err: Box<dyn std::error::Error + Send + Sync>,
37 },
38
39 #[error("Failed to spawn background thread '{name}': {err}")]
41 SpawnThread {
42 name: &'static str,
43 err: Box<dyn std::error::Error + Send + Sync>,
44 },
45}
46
47pub type ChunkBatcherResult<T> = Result<T, ChunkBatcherError>;
48
49#[derive(Clone, Default)]
51pub struct BatcherHooks {
52 #[expect(clippy::type_complexity)]
59 pub on_insert: Option<Arc<dyn Fn(&[PendingRow]) + Send + Sync>>,
60
61 #[expect(clippy::type_complexity)]
66 pub on_config_change: Option<Arc<dyn Fn(&ChunkBatcherConfig) + Send + Sync>>,
67
68 pub on_release: Option<re_log_types::ArrowRecordBatchReleaseCallback>,
74}
75
76impl BatcherHooks {
77 pub const NONE: Self = Self {
78 on_insert: None,
79 on_config_change: None,
80 on_release: None,
81 };
82}
83
84impl PartialEq for BatcherHooks {
85 fn eq(&self, other: &Self) -> bool {
86 let Self {
87 on_insert,
88 on_config_change,
89 on_release,
90 } = self;
91
92 let on_insert_eq = match (on_insert, &other.on_insert) {
93 (Some(a), Some(b)) => Arc::ptr_eq(a, b),
94 (None, None) => true,
95 _ => false,
96 };
97
98 let on_config_change_eq = match (on_config_change, &other.on_config_change) {
99 (Some(a), Some(b)) => Arc::ptr_eq(a, b),
100 (None, None) => true,
101 _ => false,
102 };
103
104 on_insert_eq && on_config_change_eq && on_release == &other.on_release
105 }
106}
107
108impl std::fmt::Debug for BatcherHooks {
109 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
110 let Self {
111 on_insert,
112 on_config_change,
113 on_release,
114 } = self;
115 f.debug_struct("BatcherHooks")
116 .field("on_insert", &on_insert.as_ref().map(|_| "…"))
117 .field("on_config_change", &on_config_change.as_ref().map(|_| "…"))
118 .field("on_release", &on_release)
119 .finish()
120 }
121}
122
123#[derive(Clone, Copy, Debug, PartialEq, Eq, re_byte_size::SizeBytes)]
129pub struct ChunkBatcherConfig {
130 pub flush_tick: Duration,
137
138 pub flush_num_bytes: u64,
142
143 pub flush_num_rows: u64,
145
146 pub chunk_max_rows_if_unsorted: u64,
149
150 pub max_bytes_in_flight: u64,
157}
158
159impl Default for ChunkBatcherConfig {
160 fn default() -> Self {
161 Self::DEFAULT
162 }
163}
164
165impl ChunkBatcherConfig {
166 pub const DEFAULT: Self = Self {
168 flush_tick: Duration::from_millis(200),
169 flush_num_bytes: 2 * 1024 * 1024, flush_num_rows: u64::MAX,
171 chunk_max_rows_if_unsorted: 8192,
172 max_bytes_in_flight: 100 * 1024 * 1024, };
174
175 pub const LOW_LATENCY: Self = Self {
177 flush_tick: Duration::from_millis(8), ..Self::DEFAULT
179 };
180
181 pub const ALWAYS_TEST_ONLY: Self = Self {
192 flush_tick: Duration::MAX,
193 flush_num_bytes: 0,
194 flush_num_rows: 0,
195 chunk_max_rows_if_unsorted: 256,
196 ..Self::DEFAULT
197 };
198
199 pub const NEVER: Self = Self {
201 flush_tick: Duration::MAX,
202 flush_num_bytes: u64::MAX,
203 flush_num_rows: u64::MAX,
204 chunk_max_rows_if_unsorted: 256,
205 ..Self::DEFAULT
206 };
207
208 #[inline]
214 pub fn always_flushes(&self) -> bool {
215 self.flush_num_rows == 0 || self.flush_num_bytes == 0
216 }
217
218 pub const ENV_FLUSH_TICK: &'static str = "RERUN_FLUSH_TICK_SECS";
220
221 pub const ENV_FLUSH_NUM_BYTES: &'static str = "RERUN_FLUSH_NUM_BYTES";
223
224 pub const ENV_FLUSH_NUM_ROWS: &'static str = "RERUN_FLUSH_NUM_ROWS";
226
227 pub const ENV_CHUNK_MAX_ROWS_IF_UNSORTED: &'static str = "RERUN_CHUNK_MAX_ROWS_IF_UNSORTED";
231
232 #[deprecated(note = "use `RERUN_CHUNK_MAX_ROWS_IF_UNSORTED` instead")]
234 const ENV_MAX_CHUNK_ROWS_IF_UNSORTED: &'static str = "RERUN_MAX_CHUNK_ROWS_IF_UNSORTED";
235
236 #[inline]
241 pub fn from_env() -> ChunkBatcherResult<Self> {
242 Self::default().apply_env()
243 }
244
245 pub fn apply_env(&self) -> ChunkBatcherResult<Self> {
250 let mut new = *self;
251
252 if let Ok(s) = std::env::var(Self::ENV_FLUSH_TICK) {
253 let flush_duration_secs: f64 =
254 s.parse().map_err(|err| ChunkBatcherError::ParseConfig {
255 name: Self::ENV_FLUSH_TICK,
256 value: s.clone(),
257 err: Box::new(err),
258 })?;
259
260 new.flush_tick = Duration::try_from_secs_f64(flush_duration_secs).map_err(|err| {
261 ChunkBatcherError::ParseConfig {
262 name: Self::ENV_FLUSH_TICK,
263 value: s.clone(),
264 err: Box::new(err),
265 }
266 })?;
267 }
268
269 if let Ok(s) = std::env::var(Self::ENV_FLUSH_NUM_BYTES) {
270 if let Some(num_bytes) = re_format::parse_bytes(&s) {
271 new.flush_num_bytes = num_bytes.unsigned_abs();
273 } else {
274 new.flush_num_bytes = s.parse().map_err(|err| ChunkBatcherError::ParseConfig {
276 name: Self::ENV_FLUSH_NUM_BYTES,
277 value: s.clone(),
278 err: Box::new(err),
279 })?;
280 }
281 }
282
283 if let Ok(s) = std::env::var(Self::ENV_FLUSH_NUM_ROWS) {
284 new.flush_num_rows = s.parse().map_err(|err| ChunkBatcherError::ParseConfig {
285 name: Self::ENV_FLUSH_NUM_ROWS,
286 value: s.clone(),
287 err: Box::new(err),
288 })?;
289 }
290
291 if let Ok(s) = std::env::var(Self::ENV_CHUNK_MAX_ROWS_IF_UNSORTED) {
292 new.chunk_max_rows_if_unsorted =
293 s.parse().map_err(|err| ChunkBatcherError::ParseConfig {
294 name: Self::ENV_CHUNK_MAX_ROWS_IF_UNSORTED,
295 value: s.clone(),
296 err: Box::new(err),
297 })?;
298 }
299
300 #[expect(deprecated)]
302 if let Ok(s) = std::env::var(Self::ENV_MAX_CHUNK_ROWS_IF_UNSORTED) {
303 new.chunk_max_rows_if_unsorted =
304 s.parse().map_err(|err| ChunkBatcherError::ParseConfig {
305 name: Self::ENV_MAX_CHUNK_ROWS_IF_UNSORTED,
306 value: s.clone(),
307 err: Box::new(err),
308 })?;
309 }
310
311 Ok(new)
312 }
313}
314
315#[test]
316fn chunk_batcher_config() {
317 #![expect(unsafe_code)] unsafe {
322 std::env::set_var("RERUN_FLUSH_TICK_SECS", "0.3");
323 std::env::set_var("RERUN_FLUSH_NUM_BYTES", "42");
324 std::env::set_var("RERUN_FLUSH_NUM_ROWS", "666");
325 std::env::set_var("RERUN_CHUNK_MAX_ROWS_IF_UNSORTED", "7777");
326 }
327
328 let config = ChunkBatcherConfig::from_env().unwrap();
329 let expected = ChunkBatcherConfig {
330 flush_tick: Duration::from_millis(300),
331 flush_num_bytes: 42,
332 flush_num_rows: 666,
333 chunk_max_rows_if_unsorted: 7777,
334 ..Default::default()
335 };
336 assert_eq!(expected, config);
337
338 unsafe {
340 std::env::set_var("RERUN_MAX_CHUNK_ROWS_IF_UNSORTED", "9999");
341 }
342
343 let config = ChunkBatcherConfig::from_env().unwrap();
344 let expected = ChunkBatcherConfig {
345 flush_tick: Duration::from_millis(300),
346 flush_num_bytes: 42,
347 flush_num_rows: 666,
348 chunk_max_rows_if_unsorted: 9999,
349 ..Default::default()
350 };
351 assert_eq!(expected, config);
352}
353
354#[derive(Clone)]
392pub struct ChunkBatcher {
393 inner: Arc<ChunkBatcherInner>,
394}
395
396struct ChunkBatcherInner {
399 tx_cmds: Sender<Command>,
403 rx_chunks: Option<Receiver<Chunk>>,
405 cmds_to_chunks_handle: Option<std::thread::JoinHandle<()>>,
406}
407
408impl Drop for ChunkBatcherInner {
409 fn drop(&mut self) {
410 if let Some(rx_chunks) = self.rx_chunks.take()
413 && !rx_chunks.is_empty()
414 {
415 re_log::warn!("Dropping data");
416 }
417
418 self.tx_cmds.send(Command::Shutdown).ok();
421 if let Some(handle) = self.cmds_to_chunks_handle.take() {
422 handle.join().ok();
423 }
424 }
425}
426
427#[derive(re_byte_size::SizeBytes)]
428enum Command {
429 AppendChunk(Chunk),
430 AppendRow(EntityPath, PendingRow),
431 Flush {
432 #[size_bytes(ignore)]
433 on_done: crossbeam::channel::Sender<()>,
434 },
435 UpdateConfig(ChunkBatcherConfig),
436 Shutdown,
437}
438
439impl Command {
440 fn flush() -> (Self, crossbeam::channel::Receiver<()>) {
441 let (tx, rx) = crossbeam::channel::bounded(1); (Self::Flush { on_done: tx }, rx)
443 }
444}
445
446impl ChunkBatcher {
447 #[must_use = "Batching threads will automatically shutdown when this object is dropped"]
452 pub fn new(config: ChunkBatcherConfig, hooks: BatcherHooks) -> ChunkBatcherResult<Self> {
453 let (tx_cmds, rx_cmd) =
454 re_quota_channel::channel("batcher_input", config.max_bytes_in_flight / 2);
455 let (tx_chunk, rx_chunks) =
456 re_quota_channel::channel("batcher_output", config.max_bytes_in_flight / 2);
457
458 let cmds_to_chunks_handle = {
459 const NAME: &str = "ChunkBatcher::cmds_to_chunks";
460 std::thread::Builder::new()
461 .name(NAME.into())
462 .spawn(move || batching_thread(config, hooks, rx_cmd, tx_chunk))
463 .map_err(|err| ChunkBatcherError::SpawnThread {
464 name: NAME,
465 err: Box::new(err),
466 })?
467 };
468
469 re_log::debug!(?config, "creating new chunk batcher");
470
471 let inner = ChunkBatcherInner {
472 tx_cmds,
473 rx_chunks: Some(rx_chunks),
474 cmds_to_chunks_handle: Some(cmds_to_chunks_handle),
475 };
476
477 Ok(Self {
478 inner: Arc::new(inner),
479 })
480 }
481
482 pub fn push_chunk(&self, chunk: Chunk) {
485 self.inner.push_chunk(chunk);
486 }
487
488 #[inline]
494 pub fn push_row(&self, entity_path: EntityPath, row: PendingRow) {
495 self.inner.push_row(entity_path, row);
496 }
497
498 #[inline]
503 pub fn flush_async(&self) {
504 self.inner.flush_async();
505 }
506
507 #[inline]
511 pub fn flush_blocking(&self, timeout: Duration) -> Result<(), BatcherFlushError> {
512 self.inner.flush_blocking(timeout)
513 }
514
515 pub fn update_config(&self, config: ChunkBatcherConfig) {
517 self.inner.update_config(config);
518 }
519
520 pub fn chunks(&self) -> Receiver<Chunk> {
528 #[expect(clippy::unwrap_used)]
531 self.inner.rx_chunks.clone().unwrap()
532 }
533}
534
535impl ChunkBatcherInner {
536 fn push_chunk(&self, chunk: Chunk) {
537 self.send_cmd(Command::AppendChunk(chunk));
538 }
539
540 fn push_row(&self, entity_path: EntityPath, row: PendingRow) {
541 self.send_cmd(Command::AppendRow(entity_path, row));
542 }
543
544 fn flush_async(&self) {
545 let (flush_cmd, _) = Command::flush();
546 self.send_cmd(flush_cmd);
547 }
548
549 fn flush_blocking(&self, timeout: Duration) -> Result<(), BatcherFlushError> {
550 use crossbeam::channel::RecvTimeoutError;
551
552 let (flush_cmd, on_done) = Command::flush();
553 self.send_cmd(flush_cmd);
554
555 on_done.recv_timeout(timeout).map_err(|err| match err {
556 RecvTimeoutError::Timeout => BatcherFlushError::Timeout,
557 RecvTimeoutError::Disconnected => BatcherFlushError::Closed,
558 })
559 }
560
561 fn update_config(&self, config: ChunkBatcherConfig) {
562 self.send_cmd(Command::UpdateConfig(config));
563 }
564
565 fn send_cmd(&self, cmd: Command) {
566 self.tx_cmds.send(cmd).ok();
569 }
570}
571
572#[expect(clippy::needless_pass_by_value)]
573fn batching_thread(
574 mut config: ChunkBatcherConfig,
575 hooks: BatcherHooks,
576 rx_cmd: Receiver<Command>,
577 tx_chunk: Sender<Chunk>,
578) {
579 let mut rx_tick = crossbeam::channel::tick(config.flush_tick);
580
581 struct Accumulator {
582 latest: Instant,
583 entity_path: EntityPath,
584 pending_rows: Vec<PendingRow>,
585 pending_num_bytes: u64,
586 }
587
588 impl Accumulator {
589 fn new(entity_path: EntityPath) -> Self {
590 Self {
591 entity_path,
592 latest: Instant::now(),
593 pending_rows: Default::default(),
594 pending_num_bytes: Default::default(),
595 }
596 }
597
598 fn reset(&mut self) {
599 self.latest = Instant::now();
600 self.pending_rows.clear();
601 self.pending_num_bytes = 0;
602 }
603 }
604
605 let mut accs: IntMap<EntityPath, Accumulator> = IntMap::default();
606
607 fn do_push_row(acc: &mut Accumulator, row: PendingRow) {
608 acc.pending_num_bytes += row.total_size_bytes();
609 acc.pending_rows.push(row);
610 }
611
612 fn do_flush_all(
613 acc: &mut Accumulator,
614 tx_chunk: &Sender<Chunk>,
615 reason: &str,
616 chunk_max_rows_if_unsorted: u64,
617 ) {
618 let rows = std::mem::take(&mut acc.pending_rows);
619 if rows.is_empty() {
620 return;
621 }
622
623 re_log::trace!(
624 "Flushing {} rows and {} bytes. Reason: {reason}",
625 rows.len(),
626 re_format::format_bytes(acc.pending_num_bytes as _)
627 );
628
629 let chunks =
630 PendingRow::many_into_chunks(acc.entity_path.clone(), chunk_max_rows_if_unsorted, rows);
631 for chunk in chunks {
632 let chunk = match chunk {
633 Ok(chunk) => chunk,
634 Err(err) => {
635 re_log::error!(%err, "corrupt chunk detected, dropping");
636 continue;
637 }
638 };
639
640 if !chunk.components.is_empty() {
644 tx_chunk.send(chunk).ok();
646 } else {
647 re_log::warn_once!(
648 "Dropping chunk without components. Entity path: {}",
649 chunk.entity_path()
650 );
651 }
652 }
653
654 acc.reset();
655 }
656
657 re_log::trace!(
658 "Flushing every: {:.2}s, {} rows, {}",
659 config.flush_tick.as_secs_f64(),
660 config.flush_num_rows,
661 re_format::format_bytes(config.flush_num_bytes as _),
662 );
663 if let Some(on_config_change) = hooks.on_config_change.as_ref() {
665 on_config_change(&config);
666 }
667
668 let mut skip_next_tick = false;
671
672 loop {
673 crossbeam::select! {
674 recv(rx_cmd.inner()) -> cmd => {
675 let Ok(cmd) = cmd else {
676 break;
679 };
680
681 let re_quota_channel::SizedMessage { msg, size_bytes } = cmd;
682
683 rx_cmd.manual_on_receive(size_bytes);
684
685 match msg {
686 Command::AppendChunk(chunk) => {
687 if !chunk.components.is_empty() {
691 tx_chunk.send(chunk).ok();
693 } else {
694 re_log::warn_once!(
695 "Dropping chunk without components. Entity path: {}",
696 chunk.entity_path()
697 );
698 }
699 },
700 Command::AppendRow(entity_path, row) => {
701 let acc = accs.entry(entity_path.clone())
702 .or_insert_with(|| Accumulator::new(entity_path));
703 do_push_row(acc, row);
704
705 if let Some(config) = hooks.on_insert.as_ref() {
706 config(&acc.pending_rows);
707 }
708
709 if acc.pending_rows.len() as u64 >= config.flush_num_rows {
710 do_flush_all(acc, &tx_chunk, "rows", config.chunk_max_rows_if_unsorted);
711 skip_next_tick = true;
712 } else if acc.pending_num_bytes >= config.flush_num_bytes {
713 do_flush_all(acc, &tx_chunk, "bytes", config.chunk_max_rows_if_unsorted);
714 skip_next_tick = true;
715 }
716 },
717
718 Command::Flush{ on_done } => {
719 skip_next_tick = true;
720 for acc in accs.values_mut() {
721 do_flush_all(acc, &tx_chunk, "manual", config.chunk_max_rows_if_unsorted);
722 }
723 re_quota_channel::send_crossbeam(&on_done, ()).ok();
724 },
725
726 Command::UpdateConfig(new_config) => {
727 if config.max_bytes_in_flight != new_config.max_bytes_in_flight {
729 re_log::warn!("Cannot change max_bytes_in_flight after batcher has been created. \
730 Previous max_bytes_in_flight: {:?}, new max_bytes_in_flight: {:?}",
731 config.max_bytes_in_flight, new_config.max_bytes_in_flight);
732 }
733
734 re_log::trace!("Updated batcher config: {:?}", new_config);
735 if let Some(on_config_change) = hooks.on_config_change.as_ref() {
736 on_config_change(&new_config);
737 }
738
739 config = new_config;
740 rx_tick = crossbeam::channel::tick(config.flush_tick);
741 }
742
743 Command::Shutdown => break,
744 }
745 },
746
747 recv(rx_tick) -> _ => {
748 if skip_next_tick {
749 skip_next_tick = false;
750 } else {
751 for acc in accs.values_mut() {
753 do_flush_all(acc, &tx_chunk, "tick", config.chunk_max_rows_if_unsorted);
754 }
755 }
756 },
757 };
758 }
759
760 drop(rx_cmd);
761 for acc in accs.values_mut() {
762 do_flush_all(
763 acc,
764 &tx_chunk,
765 "shutdown",
766 config.chunk_max_rows_if_unsorted,
767 );
768 }
769 drop(tx_chunk);
770
771 }
775
776#[derive(Debug, Clone, re_byte_size::SizeBytes)]
782pub struct PendingRow {
783 pub row_id: RowId,
786
787 pub timepoint: TimePoint,
789
790 pub components: IntMap<ComponentIdentifier, SerializedComponentBatch>,
794}
795
796impl PendingRow {
797 #[inline]
798 pub fn new(
799 timepoint: TimePoint,
800 components: IntMap<ComponentIdentifier, SerializedComponentBatch>,
801 ) -> Self {
802 Self {
803 row_id: RowId::new(),
804 timepoint,
805 components,
806 }
807 }
808
809 #[inline]
810 pub fn from_iter(
811 timepoint: TimePoint,
812 components: impl IntoIterator<Item = SerializedComponentBatch>,
813 ) -> Self {
814 Self::new(
815 timepoint,
816 components
817 .into_iter()
818 .map(|component| (component.descriptor.component, component))
819 .collect(),
820 )
821 }
822}
823
824impl PendingRow {
825 pub fn into_chunk(self, entity_path: EntityPath) -> ChunkResult<Chunk> {
832 let Self {
833 row_id,
834 timepoint,
835 components,
836 } = self;
837
838 let timelines = timepoint
839 .into_iter()
840 .map(|(timeline_name, cell)| {
841 let times = ArrowScalarBuffer::from(vec![cell.as_i64()]);
842 let time_column =
843 TimeColumn::new(Some(true), Timeline::new(timeline_name, cell.typ()), times);
844 (timeline_name, time_column)
845 })
846 .collect();
847
848 let mut per_desc = ChunkComponents::default();
849 for (_component, batch) in components {
850 let list_array = arrays_to_list_array_opt(&[Some(&*batch.array as _)]);
851 if let Some(list_array) = list_array {
852 per_desc.insert(SerializedComponentColumn::new(list_array, batch.descriptor));
853 }
854 }
855
856 Chunk::from_native_row_ids(
857 ChunkId::new(),
858 entity_path,
859 Some(true),
860 &[row_id],
861 timelines,
862 per_desc,
863 )
864 }
865
866 pub fn many_into_chunks(
880 entity_path: EntityPath,
881 chunk_max_rows_if_unsorted: u64,
882 mut rows: Vec<Self>,
883 ) -> impl Iterator<Item = ChunkResult<Chunk>> {
884 re_tracing::profile_function!();
885
886 {
889 re_tracing::profile_scope!("sort rows");
890 rows.sort_by_key(|row| row.row_id);
891 }
892
893 let mut per_timeline_set: IntMap<u64 , Vec<Self>> = Default::default();
895 {
896 re_tracing::profile_scope!("compute timeline sets");
897
898 for row in rows {
901 let mut hasher = ahash::AHasher::default();
902 row.timepoint.timeline_names().for_each(|timeline| {
903 <TimelineName as std::hash::Hash>::hash(timeline, &mut hasher);
904 });
905
906 per_timeline_set
907 .entry(hasher.finish())
908 .or_default()
909 .push(row);
910 }
911 }
912
913 per_timeline_set.into_values().flat_map(move |rows| {
914 re_tracing::profile_scope!("iterate per timeline set");
915
916 let mut per_datatype_set: IntMap<u64 , Vec<Self>> =
918 Default::default();
919 {
920 re_tracing::profile_scope!("compute datatype sets");
921
922 for row in rows {
935 let mut hasher = ahash::AHasher::default();
936 row.components
937 .values()
938 .for_each(|batch| batch.array.data_type().hash(&mut hasher));
939 per_datatype_set
940 .entry(hasher.finish())
941 .or_default()
942 .push(row);
943 }
944 }
945
946 let entity_path = entity_path.clone();
948 per_datatype_set.into_values().flat_map(move |rows| {
949 re_tracing::profile_scope!("iterate per datatype set");
950
951 let mut row_ids: Vec<RowId> = Vec::with_capacity(rows.len());
952 let mut timelines: IntMap<TimelineName, PendingTimeColumn> = IntMap::default();
953
954 let mut all_components: IntMap<ComponentIdentifier, _> = IntMap::default();
957 for row in &rows {
958 for (component, batch) in &row.components {
959 all_components
960 .entry(*component)
961 .or_insert_with(|| (batch.descriptor.clone(), Vec::new()));
962 }
963 }
964
965 let mut chunks = Vec::new();
966
967 let mut components = all_components.clone();
968 for row in &rows {
969 let Self {
970 row_id,
971 timepoint: row_timepoint,
972 components: row_components,
973 } = row;
974
975 for (&timeline_name, cell) in row_timepoint {
979 let time_column = timelines.entry(timeline_name).or_insert_with(|| {
980 PendingTimeColumn::new(Timeline::new(timeline_name, cell.typ()))
981 });
982
983 if !row_ids.is_empty() && row_ids.len() as u64 >= chunk_max_rows_if_unsorted
985 && !time_column.is_sorted
986 {
987 chunks.push(Chunk::from_native_row_ids(
988 ChunkId::new(),
989 entity_path.clone(),
990 Some(true),
991 &std::mem::take(&mut row_ids),
992 std::mem::take(&mut timelines)
993 .into_iter()
994 .map(|(name, time_column)| (name, time_column.finish()))
995 .collect(),
996 {
997 let mut per_component = ChunkComponents::default();
998 for (_component, (desc, arrays)) in
999 std::mem::take(&mut components)
1000 {
1001 let list_array = arrays_to_list_array_opt(&arrays);
1002 if let Some(list_array) = list_array {
1003 per_component.insert(SerializedComponentColumn::new(
1004 list_array, desc,
1005 ));
1006 }
1007 }
1008 per_component
1009 },
1010 ));
1011
1012 components = all_components.clone();
1013 }
1014 }
1015
1016 row_ids.push(*row_id);
1017
1018 for (&timeline_name, &cell) in row_timepoint {
1019 let time_column = timelines.entry(timeline_name).or_insert_with(|| {
1020 PendingTimeColumn::new(Timeline::new(timeline_name, cell.typ()))
1021 });
1022 time_column.push(cell.into());
1023 }
1024
1025 for (component, (_desc, arrays)) in &mut components {
1026 arrays.push(
1029 row_components
1030 .get(component)
1031 .map(|batch| &*batch.array as _),
1032 );
1033 }
1034 }
1035
1036 chunks.push(Chunk::from_native_row_ids(
1037 ChunkId::new(),
1038 entity_path.clone(),
1039 Some(true),
1040 &std::mem::take(&mut row_ids),
1041 timelines
1042 .into_iter()
1043 .map(|(timeline, time_column)| (timeline, time_column.finish()))
1044 .collect(),
1045 {
1046 let mut per_desc = ChunkComponents::default();
1047 for (_component, (desc, arrays)) in components {
1048 let list_array = arrays_to_list_array_opt(&arrays);
1049 if let Some(list_array) = list_array {
1050 per_desc.insert(SerializedComponentColumn::new(list_array, desc));
1051 }
1052 }
1053 per_desc
1054 },
1055 ));
1056
1057 chunks
1058 })
1059 })
1060 }
1061}
1062
1063struct PendingTimeColumn {
1067 timeline: Timeline,
1068 times: Vec<i64>,
1069 is_sorted: bool,
1070 time_range: AbsoluteTimeRange,
1071}
1072
1073impl PendingTimeColumn {
1074 fn new(timeline: Timeline) -> Self {
1075 Self {
1076 timeline,
1077 times: Default::default(),
1078 is_sorted: true,
1079 time_range: AbsoluteTimeRange::EMPTY,
1080 }
1081 }
1082
1083 fn push(&mut self, time: TimeInt) {
1085 let Self {
1086 timeline: _,
1087 times,
1088 is_sorted,
1089 time_range,
1090 } = self;
1091
1092 *is_sorted &= times.last().copied().unwrap_or(TimeInt::MIN.as_i64()) <= time.as_i64();
1093 time_range.set_min(TimeInt::min(time_range.min(), time));
1094 time_range.set_max(TimeInt::max(time_range.max(), time));
1095 times.push(time.as_i64());
1096 }
1097
1098 fn finish(self) -> TimeColumn {
1099 let Self {
1100 timeline,
1101 times,
1102 is_sorted,
1103 time_range,
1104 } = self;
1105
1106 TimeColumn {
1107 timeline,
1108 times: ArrowScalarBuffer::from(times),
1109 is_sorted,
1110 time_range,
1111 }
1112 }
1113}
1114
1115#[cfg(test)]
1122mod tests {
1123 use crossbeam::channel::TryRecvError;
1124 use re_log_types::example_components::{MyIndex, MyLabel, MyPoint, MyPoint64, MyPoints};
1125 use re_types_core::{ComponentDescriptor, Loggable as _};
1126
1127 use super::*;
1128
1129 #[test]
1131 fn simple() -> anyhow::Result<()> {
1132 let batcher = ChunkBatcher::new(ChunkBatcherConfig::NEVER, BatcherHooks::NONE)?;
1133
1134 let timeline1 = Timeline::new_duration("log_time");
1135
1136 let timepoint1 = TimePoint::default().with(timeline1, 42);
1137 let timepoint2 = TimePoint::default().with(timeline1, 43);
1138 let timepoint3 = TimePoint::default().with(timeline1, 44);
1139
1140 let points1 = MyPoint::to_arrow([MyPoint::new(1.0, 2.0), MyPoint::new(3.0, 4.0)])?;
1141 let points2 = MyPoint::to_arrow([MyPoint::new(10.0, 20.0), MyPoint::new(30.0, 40.0)])?;
1142 let points3 = MyPoint::to_arrow([MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)])?;
1143
1144 let labels1 = MyLabel::to_arrow([MyLabel("a".into()), MyLabel("b".into())])?;
1145 let labels2 = MyLabel::to_arrow([MyLabel("c".into()), MyLabel("d".into())])?;
1146 let labels3 = MyLabel::to_arrow([MyLabel("e".into()), MyLabel("d".into())])?;
1147
1148 let indices1 = MyIndex::to_arrow([MyIndex(0), MyIndex(1)])?;
1149 let indices2 = MyIndex::to_arrow([MyIndex(2), MyIndex(3)])?;
1150 let indices3 = MyIndex::to_arrow([MyIndex(4), MyIndex(5)])?;
1151
1152 let components1 = [
1153 SerializedComponentBatch::new(points1.clone(), MyPoints::descriptor_points()),
1154 SerializedComponentBatch::new(labels1.clone(), MyPoints::descriptor_labels()),
1155 SerializedComponentBatch::new(indices1.clone(), MyIndex::partial_descriptor()),
1156 ];
1157 let components2 = [
1158 SerializedComponentBatch::new(points2.clone(), MyPoints::descriptor_points()),
1159 SerializedComponentBatch::new(labels2.clone(), MyPoints::descriptor_labels()),
1160 SerializedComponentBatch::new(indices2.clone(), MyIndex::partial_descriptor()),
1161 ];
1162 let components3 = [
1163 SerializedComponentBatch::new(points3.clone(), MyPoints::descriptor_points()),
1164 SerializedComponentBatch::new(labels3.clone(), MyPoints::descriptor_labels()),
1165 SerializedComponentBatch::new(indices3.clone(), MyIndex::partial_descriptor()),
1166 ];
1167
1168 let row1 = PendingRow::from_iter(timepoint1.clone(), components1);
1169 let row2 = PendingRow::from_iter(timepoint2.clone(), components2);
1170 let row3 = PendingRow::from_iter(timepoint3.clone(), components3);
1171
1172 let entity_path1: EntityPath = "a/b/c".into();
1173 batcher.push_row(entity_path1.clone(), row1.clone());
1174 batcher.push_row(entity_path1.clone(), row2.clone());
1175 batcher.push_row(entity_path1.clone(), row3.clone());
1176
1177 let chunks_rx = batcher.chunks();
1178 drop(batcher); let mut chunks = Vec::new();
1181 loop {
1182 let chunk = match chunks_rx.try_recv() {
1183 Ok(chunk) => chunk,
1184 Err(TryRecvError::Empty) => panic!("expected chunk, got none"),
1185 Err(TryRecvError::Disconnected) => break,
1186 };
1187 chunks.push(chunk);
1188 }
1189
1190 chunks.sort_by_key(|chunk| chunk.row_id_range().unwrap().0);
1191
1192 eprintln!("Chunks:");
1194 for chunk in &chunks {
1195 eprintln!("{chunk}");
1196 }
1197
1198 assert_eq!(1, chunks.len());
1199
1200 {
1201 let expected_row_ids = vec![row1.row_id, row2.row_id, row3.row_id];
1202 let expected_timelines = [(
1203 *timeline1.name(),
1204 TimeColumn::new(Some(true), timeline1, vec![42, 43, 44].into()),
1205 )];
1206 let expected_components = [
1207 (
1208 MyPoints::descriptor_points(),
1209 arrays_to_list_array_opt(&[&*points1, &*points2, &*points3].map(Some)).unwrap(),
1210 ), (
1212 MyPoints::descriptor_labels(),
1213 arrays_to_list_array_opt(&[&*labels1, &*labels2, &*labels3].map(Some)).unwrap(),
1214 ), (
1216 MyIndex::partial_descriptor(),
1217 arrays_to_list_array_opt(&[&*indices1, &*indices2, &*indices3].map(Some))
1218 .unwrap(),
1219 ), ];
1221 let expected_chunk = Chunk::from_native_row_ids(
1222 chunks[0].id,
1223 entity_path1.clone(),
1224 None,
1225 &expected_row_ids,
1226 expected_timelines.into_iter().collect(),
1227 expected_components.into_iter().collect(),
1228 )?;
1229
1230 eprintln!("Expected:\n{expected_chunk}");
1231 eprintln!("Got:\n{}", chunks[0]);
1232 assert_eq!(expected_chunk, chunks[0]);
1233 }
1234
1235 Ok(())
1236 }
1237
1238 #[test]
1239 #[expect(clippy::len_zero)]
1240 fn simple_but_hashes_might_not_match() -> anyhow::Result<()> {
1241 let batcher = ChunkBatcher::new(ChunkBatcherConfig::NEVER, BatcherHooks::NONE)?;
1242
1243 let timeline1 = Timeline::new_duration("log_time");
1244
1245 let timepoint1 = TimePoint::default().with(timeline1, 42);
1246 let timepoint2 = TimePoint::default().with(timeline1, 43);
1247 let timepoint3 = TimePoint::default().with(timeline1, 44);
1248
1249 let points1 = MyPoint::to_arrow([MyPoint::new(1.0, 2.0), MyPoint::new(3.0, 4.0)])?;
1250 let points2 = MyPoint::to_arrow([MyPoint::new(10.0, 20.0), MyPoint::new(30.0, 40.0)])?;
1251 let points3 = MyPoint::to_arrow([MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)])?;
1252
1253 let labels1 = MyLabel::to_arrow([MyLabel("a".into()), MyLabel("b".into())])?;
1254 let labels2 = MyLabel::to_arrow([MyLabel("c".into()), MyLabel("d".into())])?;
1255 let labels3 = MyLabel::to_arrow([MyLabel("e".into()), MyLabel("d".into())])?;
1256
1257 let indices1 = MyIndex::to_arrow([MyIndex(0), MyIndex(1)])?;
1258 let indices2 = MyIndex::to_arrow([MyIndex(2), MyIndex(3)])?;
1259 let indices3 = MyIndex::to_arrow([MyIndex(4), MyIndex(5)])?;
1260
1261 let components1 = [
1262 SerializedComponentBatch::new(indices1.clone(), MyIndex::partial_descriptor()),
1263 SerializedComponentBatch::new(points1.clone(), MyPoints::descriptor_points()),
1264 SerializedComponentBatch::new(labels1.clone(), MyPoints::descriptor_labels()),
1265 ];
1266 let components2 = [
1267 SerializedComponentBatch::new(points2.clone(), MyPoints::descriptor_points()),
1268 SerializedComponentBatch::new(labels2.clone(), MyPoints::descriptor_labels()),
1269 SerializedComponentBatch::new(indices2.clone(), MyIndex::partial_descriptor()),
1270 ];
1271 let components3 = [
1272 SerializedComponentBatch::new(labels3.clone(), MyPoints::descriptor_labels()),
1273 SerializedComponentBatch::new(indices3.clone(), MyIndex::partial_descriptor()),
1274 SerializedComponentBatch::new(points3.clone(), MyPoints::descriptor_points()),
1275 ];
1276
1277 let row1 = PendingRow::from_iter(timepoint1.clone(), components1);
1278 let row2 = PendingRow::from_iter(timepoint2.clone(), components2);
1279 let row3 = PendingRow::from_iter(timepoint3.clone(), components3);
1280
1281 let entity_path1: EntityPath = "a/b/c".into();
1282 batcher.push_row(entity_path1.clone(), row1.clone());
1283 batcher.push_row(entity_path1.clone(), row2.clone());
1284 batcher.push_row(entity_path1.clone(), row3.clone());
1285
1286 let chunks_rx = batcher.chunks();
1287 drop(batcher); let mut chunks = Vec::new();
1290 loop {
1291 let chunk = match chunks_rx.try_recv() {
1292 Ok(chunk) => chunk,
1293 Err(TryRecvError::Empty) => panic!("expected chunk, got none"),
1294 Err(TryRecvError::Disconnected) => break,
1295 };
1296 chunks.push(chunk);
1297 }
1298
1299 chunks.sort_by_key(|chunk| chunk.row_id_range().unwrap().0);
1300
1301 eprintln!("Chunks:");
1303 for chunk in &chunks {
1304 eprintln!("{chunk}");
1305 }
1306
1307 assert!(chunks.len() >= 1);
1315
1316 Ok(())
1317 }
1318
1319 #[test]
1320 #[expect(clippy::zero_sized_map_values)]
1321 fn intmap_order_is_deterministic() {
1322 let descriptors = [
1323 MyPoints::descriptor_points(),
1324 MyPoints::descriptor_colors(),
1325 MyPoints::descriptor_labels(),
1326 MyPoint64::partial_descriptor(),
1327 MyIndex::partial_descriptor(),
1328 ];
1329
1330 let expected: IntMap<ComponentDescriptor, ()> =
1331 descriptors.iter().cloned().map(|d| (d, ())).collect();
1332 let expected: Vec<_> = expected.into_keys().collect();
1333
1334 for _ in 0..1_000 {
1335 let got_collect: IntMap<ComponentDescriptor, ()> =
1336 descriptors.clone().into_iter().map(|d| (d, ())).collect();
1337 let got_collect: Vec<_> = got_collect.into_keys().collect();
1338
1339 let mut got_insert: IntMap<ComponentDescriptor, ()> = Default::default();
1340 for d in descriptors.clone() {
1341 got_insert.insert(d, ());
1342 }
1343 let got_insert: Vec<_> = got_insert.into_keys().collect();
1344
1345 assert_eq!(expected, got_collect);
1346 assert_eq!(expected, got_insert);
1347 }
1348 }
1349
1350 #[test]
1352 fn simple_static() -> anyhow::Result<()> {
1353 let batcher = ChunkBatcher::new(ChunkBatcherConfig::NEVER, BatcherHooks::NONE)?;
1354
1355 let static_ = TimePoint::default();
1356
1357 let points1 = MyPoint::to_arrow([MyPoint::new(1.0, 2.0), MyPoint::new(3.0, 4.0)])?;
1358 let points2 = MyPoint::to_arrow([MyPoint::new(10.0, 20.0), MyPoint::new(30.0, 40.0)])?;
1359 let points3 = MyPoint::to_arrow([MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)])?;
1360
1361 let components1 = [SerializedComponentBatch::new(
1362 points1.clone(),
1363 MyPoints::descriptor_points(),
1364 )];
1365 let components2 = [SerializedComponentBatch::new(
1366 points2.clone(),
1367 MyPoints::descriptor_points(),
1368 )];
1369 let components3 = [SerializedComponentBatch::new(
1370 points3.clone(),
1371 MyPoints::descriptor_points(),
1372 )];
1373
1374 let row1 = PendingRow::from_iter(static_.clone(), components1);
1375 let row2 = PendingRow::from_iter(static_.clone(), components2);
1376 let row3 = PendingRow::from_iter(static_.clone(), components3);
1377
1378 let entity_path1: EntityPath = "a/b/c".into();
1379 batcher.push_row(entity_path1.clone(), row1.clone());
1380 batcher.push_row(entity_path1.clone(), row2.clone());
1381 batcher.push_row(entity_path1.clone(), row3.clone());
1382
1383 let chunks_rx = batcher.chunks();
1384 drop(batcher); let mut chunks = Vec::new();
1387 loop {
1388 let chunk = match chunks_rx.try_recv() {
1389 Ok(chunk) => chunk,
1390 Err(TryRecvError::Empty) => panic!("expected chunk, got none"),
1391 Err(TryRecvError::Disconnected) => break,
1392 };
1393 chunks.push(chunk);
1394 }
1395
1396 chunks.sort_by_key(|chunk| chunk.row_id_range().unwrap().0);
1397
1398 eprintln!("Chunks:");
1400 for chunk in &chunks {
1401 eprintln!("{chunk}");
1402 }
1403
1404 assert_eq!(1, chunks.len());
1405
1406 {
1407 let expected_row_ids = vec![row1.row_id, row2.row_id, row3.row_id];
1408 let expected_timelines = [];
1409 let expected_components = [(
1410 MyPoints::descriptor_points(),
1411 arrays_to_list_array_opt(&[&*points1, &*points2, &*points3].map(Some)).unwrap(),
1412 )];
1413 let expected_chunk = Chunk::from_native_row_ids(
1414 chunks[0].id,
1415 entity_path1.clone(),
1416 None,
1417 &expected_row_ids,
1418 expected_timelines.into_iter().collect(),
1419 expected_components.into_iter().collect(),
1420 )?;
1421
1422 eprintln!("Expected:\n{expected_chunk}");
1423 eprintln!("Got:\n{}", chunks[0]);
1424 assert_eq!(expected_chunk, chunks[0]);
1425 }
1426
1427 Ok(())
1428 }
1429
1430 #[test]
1432 fn different_entities() -> anyhow::Result<()> {
1433 let batcher = ChunkBatcher::new(ChunkBatcherConfig::NEVER, BatcherHooks::NONE)?;
1434
1435 let timeline1 = Timeline::new_duration("log_time");
1436
1437 let timepoint1 = TimePoint::default().with(timeline1, 42);
1438 let timepoint2 = TimePoint::default().with(timeline1, 43);
1439 let timepoint3 = TimePoint::default().with(timeline1, 44);
1440
1441 let points1 = MyPoint::to_arrow([MyPoint::new(1.0, 2.0), MyPoint::new(3.0, 4.0)])?;
1442 let points2 = MyPoint::to_arrow([MyPoint::new(10.0, 20.0), MyPoint::new(30.0, 40.0)])?;
1443 let points3 = MyPoint::to_arrow([MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)])?;
1444
1445 let components1 = [SerializedComponentBatch::new(
1446 points1.clone(),
1447 MyPoints::descriptor_points(),
1448 )];
1449 let components2 = [SerializedComponentBatch::new(
1450 points2.clone(),
1451 MyPoints::descriptor_points(),
1452 )];
1453 let components3 = [SerializedComponentBatch::new(
1454 points3.clone(),
1455 MyPoints::descriptor_points(),
1456 )];
1457
1458 let row1 = PendingRow::from_iter(timepoint1.clone(), components1);
1459 let row2 = PendingRow::from_iter(timepoint2.clone(), components2);
1460 let row3 = PendingRow::from_iter(timepoint3.clone(), components3);
1461
1462 let entity_path1: EntityPath = "ent1".into();
1463 let entity_path2: EntityPath = "ent2".into();
1464 batcher.push_row(entity_path1.clone(), row1.clone());
1465 batcher.push_row(entity_path2.clone(), row2.clone());
1466 batcher.push_row(entity_path1.clone(), row3.clone());
1467
1468 let chunks_rx = batcher.chunks();
1469 drop(batcher); let mut chunks = Vec::new();
1472 loop {
1473 let chunk = match chunks_rx.try_recv() {
1474 Ok(chunk) => chunk,
1475 Err(TryRecvError::Empty) => panic!("expected chunk, got none"),
1476 Err(TryRecvError::Disconnected) => break,
1477 };
1478 chunks.push(chunk);
1479 }
1480
1481 chunks.sort_by_key(|chunk| chunk.row_id_range().unwrap().0);
1482
1483 eprintln!("Chunks:");
1485 for chunk in &chunks {
1486 eprintln!("{chunk}");
1487 }
1488
1489 assert_eq!(2, chunks.len());
1490
1491 {
1492 let expected_row_ids = vec![row1.row_id, row3.row_id];
1493 let expected_timelines = [(
1494 *timeline1.name(),
1495 TimeColumn::new(Some(true), timeline1, vec![42, 44].into()),
1496 )];
1497 let expected_components = [(
1498 MyPoints::descriptor_points(),
1499 arrays_to_list_array_opt(&[&*points1, &*points3].map(Some)).unwrap(),
1500 )];
1501 let expected_chunk = Chunk::from_native_row_ids(
1502 chunks[0].id,
1503 entity_path1.clone(),
1504 None,
1505 &expected_row_ids,
1506 expected_timelines.into_iter().collect(),
1507 expected_components.into_iter().collect(),
1508 )?;
1509
1510 eprintln!("Expected:\n{expected_chunk}");
1511 eprintln!("Got:\n{}", chunks[0]);
1512 assert_eq!(expected_chunk, chunks[0]);
1513 }
1514
1515 {
1516 let expected_row_ids = vec![row2.row_id];
1517 let expected_timelines = [(
1518 *timeline1.name(),
1519 TimeColumn::new(Some(true), timeline1, vec![43].into()),
1520 )];
1521 let expected_components = [(
1522 MyPoints::descriptor_points(),
1523 arrays_to_list_array_opt(&[&*points2].map(Some)).unwrap(),
1524 )];
1525 let expected_chunk = Chunk::from_native_row_ids(
1526 chunks[1].id,
1527 entity_path2.clone(),
1528 None,
1529 &expected_row_ids,
1530 expected_timelines.into_iter().collect(),
1531 expected_components.into_iter().collect(),
1532 )?;
1533
1534 eprintln!("Expected:\n{expected_chunk}");
1535 eprintln!("Got:\n{}", chunks[1]);
1536 assert_eq!(expected_chunk, chunks[1]);
1537 }
1538
1539 Ok(())
1540 }
1541
1542 #[test]
1544 fn different_timelines() -> anyhow::Result<()> {
1545 let batcher = ChunkBatcher::new(ChunkBatcherConfig::NEVER, BatcherHooks::NONE)?;
1546
1547 let timeline1 = Timeline::new_duration("log_time");
1548 let timeline2 = Timeline::new_sequence("frame_nr");
1549
1550 let timepoint1 = TimePoint::default().with(timeline1, 42);
1551 let timepoint2 = TimePoint::default()
1552 .with(timeline1, 43)
1553 .with(timeline2, 1000);
1554 let timepoint3 = TimePoint::default()
1555 .with(timeline1, 44)
1556 .with(timeline2, 1001);
1557
1558 let points1 = MyPoint::to_arrow([MyPoint::new(1.0, 2.0), MyPoint::new(3.0, 4.0)])?;
1559 let points2 = MyPoint::to_arrow([MyPoint::new(10.0, 20.0), MyPoint::new(30.0, 40.0)])?;
1560 let points3 = MyPoint::to_arrow([MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)])?;
1561
1562 let components1 = [SerializedComponentBatch::new(
1563 points1.clone(),
1564 MyPoints::descriptor_points(),
1565 )];
1566 let components2 = [SerializedComponentBatch::new(
1567 points2.clone(),
1568 MyPoints::descriptor_points(),
1569 )];
1570 let components3 = [SerializedComponentBatch::new(
1571 points3.clone(),
1572 MyPoints::descriptor_points(),
1573 )];
1574
1575 let row1 = PendingRow::from_iter(timepoint1.clone(), components1);
1576 let row2 = PendingRow::from_iter(timepoint2.clone(), components2);
1577 let row3 = PendingRow::from_iter(timepoint3.clone(), components3);
1578
1579 let entity_path1: EntityPath = "a/b/c".into();
1580 batcher.push_row(entity_path1.clone(), row1.clone());
1581 batcher.push_row(entity_path1.clone(), row2.clone());
1582 batcher.push_row(entity_path1.clone(), row3.clone());
1583
1584 let chunks_rx = batcher.chunks();
1585 drop(batcher); let mut chunks = Vec::new();
1588 loop {
1589 let chunk = match chunks_rx.try_recv() {
1590 Ok(chunk) => chunk,
1591 Err(TryRecvError::Empty) => panic!("expected chunk, got none"),
1592 Err(TryRecvError::Disconnected) => break,
1593 };
1594 chunks.push(chunk);
1595 }
1596
1597 chunks.sort_by_key(|chunk| chunk.row_id_range().unwrap().0);
1598
1599 eprintln!("Chunks:");
1601 for chunk in &chunks {
1602 eprintln!("{chunk}");
1603 }
1604
1605 assert_eq!(2, chunks.len());
1606
1607 {
1608 let expected_row_ids = vec![row1.row_id];
1609 let expected_timelines = [(
1610 *timeline1.name(),
1611 TimeColumn::new(Some(true), timeline1, vec![42].into()),
1612 )];
1613 let expected_components = [(
1614 MyPoints::descriptor_points(),
1615 arrays_to_list_array_opt(&[&*points1].map(Some)).unwrap(),
1616 )];
1617 let expected_chunk = Chunk::from_native_row_ids(
1618 chunks[0].id,
1619 entity_path1.clone(),
1620 None,
1621 &expected_row_ids,
1622 expected_timelines.into_iter().collect(),
1623 expected_components.into_iter().collect(),
1624 )?;
1625
1626 eprintln!("Expected:\n{expected_chunk}");
1627 eprintln!("Got:\n{}", chunks[0]);
1628 assert_eq!(expected_chunk, chunks[0]);
1629 }
1630
1631 {
1632 let expected_row_ids = vec![row2.row_id, row3.row_id];
1633 let expected_timelines = [
1634 (
1635 *timeline1.name(),
1636 TimeColumn::new(Some(true), timeline1, vec![43, 44].into()),
1637 ),
1638 (
1639 *timeline2.name(),
1640 TimeColumn::new(Some(true), timeline2, vec![1000, 1001].into()),
1641 ),
1642 ];
1643 let expected_components = [(
1644 MyPoints::descriptor_points(),
1645 arrays_to_list_array_opt(&[&*points2, &*points3].map(Some)).unwrap(),
1646 )];
1647 let expected_chunk = Chunk::from_native_row_ids(
1648 chunks[1].id,
1649 entity_path1.clone(),
1650 None,
1651 &expected_row_ids,
1652 expected_timelines.into_iter().collect(),
1653 expected_components.into_iter().collect(),
1654 )?;
1655
1656 eprintln!("Expected:\n{expected_chunk}");
1657 eprintln!("Got:\n{}", chunks[1]);
1658 assert_eq!(expected_chunk, chunks[1]);
1659 }
1660
1661 Ok(())
1662 }
1663
1664 #[test]
1666 fn different_datatypes() -> anyhow::Result<()> {
1667 let batcher = ChunkBatcher::new(ChunkBatcherConfig::NEVER, BatcherHooks::NONE)?;
1668
1669 let timeline1 = Timeline::new_duration("log_time");
1670
1671 let timepoint1 = TimePoint::default().with(timeline1, 42);
1672 let timepoint2 = TimePoint::default().with(timeline1, 43);
1673 let timepoint3 = TimePoint::default().with(timeline1, 44);
1674
1675 let points1 = MyPoint::to_arrow([MyPoint::new(1.0, 2.0), MyPoint::new(3.0, 4.0)])?;
1676 let points2 =
1677 MyPoint64::to_arrow([MyPoint64::new(10.0, 20.0), MyPoint64::new(30.0, 40.0)])?;
1678 let points3 = MyPoint::to_arrow([MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)])?;
1679
1680 let components1 = [SerializedComponentBatch::new(
1681 points1.clone(),
1682 MyPoints::descriptor_points(),
1683 )];
1684 let components2 = [SerializedComponentBatch::new(
1685 points2.clone(),
1686 MyPoints::descriptor_points(),
1687 )]; let components3 = [SerializedComponentBatch::new(
1689 points3.clone(),
1690 MyPoints::descriptor_points(),
1691 )];
1692
1693 let row1 = PendingRow::from_iter(timepoint1.clone(), components1);
1694 let row2 = PendingRow::from_iter(timepoint2.clone(), components2);
1695 let row3 = PendingRow::from_iter(timepoint3.clone(), components3);
1696
1697 let entity_path1: EntityPath = "a/b/c".into();
1698 batcher.push_row(entity_path1.clone(), row1.clone());
1699 batcher.push_row(entity_path1.clone(), row2.clone());
1700 batcher.push_row(entity_path1.clone(), row3.clone());
1701
1702 let chunks_rx = batcher.chunks();
1703 drop(batcher); let mut chunks = Vec::new();
1706 loop {
1707 let chunk = match chunks_rx.try_recv() {
1708 Ok(chunk) => chunk,
1709 Err(TryRecvError::Empty) => panic!("expected chunk, got none"),
1710 Err(TryRecvError::Disconnected) => break,
1711 };
1712 chunks.push(chunk);
1713 }
1714
1715 chunks.sort_by_key(|chunk| chunk.row_id_range().unwrap().0);
1716
1717 eprintln!("Chunks:");
1719 for chunk in &chunks {
1720 eprintln!("{chunk}");
1721 }
1722
1723 assert_eq!(2, chunks.len());
1724
1725 {
1726 let expected_row_ids = vec![row1.row_id, row3.row_id];
1727 let expected_timelines = [(
1728 *timeline1.name(),
1729 TimeColumn::new(Some(true), timeline1, vec![42, 44].into()),
1730 )];
1731 let expected_components = [(
1732 MyPoints::descriptor_points(),
1733 arrays_to_list_array_opt(&[&*points1, &*points3].map(Some)).unwrap(),
1734 )];
1735 let expected_chunk = Chunk::from_native_row_ids(
1736 chunks[0].id,
1737 entity_path1.clone(),
1738 None,
1739 &expected_row_ids,
1740 expected_timelines.into_iter().collect(),
1741 expected_components.into_iter().collect(),
1742 )?;
1743
1744 eprintln!("Expected:\n{expected_chunk}");
1745 eprintln!("Got:\n{}", chunks[0]);
1746 assert_eq!(expected_chunk, chunks[0]);
1747 }
1748
1749 {
1750 let expected_row_ids = vec![row2.row_id];
1751 let expected_timelines = [(
1752 *timeline1.name(),
1753 TimeColumn::new(Some(true), timeline1, vec![43].into()),
1754 )];
1755 let expected_components = [(
1756 MyPoints::descriptor_points(),
1757 arrays_to_list_array_opt(&[&*points2].map(Some)).unwrap(),
1758 )];
1759 let expected_chunk = Chunk::from_native_row_ids(
1760 chunks[1].id,
1761 entity_path1.clone(),
1762 None,
1763 &expected_row_ids,
1764 expected_timelines.into_iter().collect(),
1765 expected_components.into_iter().collect(),
1766 )?;
1767
1768 eprintln!("Expected:\n{expected_chunk}");
1769 eprintln!("Got:\n{}", chunks[1]);
1770 assert_eq!(expected_chunk, chunks[1]);
1771 }
1772
1773 Ok(())
1774 }
1775
1776 #[test]
1779 fn unsorted_timeline_below_threshold() -> anyhow::Result<()> {
1780 let batcher = ChunkBatcher::new(
1781 ChunkBatcherConfig {
1782 chunk_max_rows_if_unsorted: 1000,
1783 ..ChunkBatcherConfig::NEVER
1784 },
1785 BatcherHooks::NONE,
1786 )?;
1787
1788 let timeline1 = Timeline::new_duration("log_time");
1789 let timeline2 = Timeline::new_duration("frame_nr");
1790
1791 let timepoint1 = TimePoint::default()
1792 .with(timeline2, 1000)
1793 .with(timeline1, 42);
1794 let timepoint2 = TimePoint::default()
1795 .with(timeline2, 1001)
1796 .with(timeline1, 43);
1797 let timepoint3 = TimePoint::default()
1798 .with(timeline2, 1002)
1799 .with(timeline1, 44);
1800 let timepoint4 = TimePoint::default()
1801 .with(timeline2, 1003)
1802 .with(timeline1, 45);
1803
1804 let points1 = MyPoint::to_arrow([MyPoint::new(1.0, 2.0), MyPoint::new(3.0, 4.0)])?;
1805 let points2 = MyPoint::to_arrow([MyPoint::new(10.0, 20.0), MyPoint::new(30.0, 40.0)])?;
1806 let points3 = MyPoint::to_arrow([MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)])?;
1807 let points4 =
1808 MyPoint::to_arrow([MyPoint::new(1000.0, 2000.0), MyPoint::new(3000.0, 4000.0)])?;
1809
1810 let components1 = [SerializedComponentBatch::new(
1811 points1.clone(),
1812 MyPoints::descriptor_points(),
1813 )];
1814 let components2 = [SerializedComponentBatch::new(
1815 points2.clone(),
1816 MyPoints::descriptor_points(),
1817 )];
1818 let components3 = [SerializedComponentBatch::new(
1819 points3.clone(),
1820 MyPoints::descriptor_points(),
1821 )];
1822 let components4 = [SerializedComponentBatch::new(
1823 points4.clone(),
1824 MyPoints::descriptor_points(),
1825 )];
1826
1827 let row1 = PendingRow::from_iter(timepoint4.clone(), components1);
1828 let row2 = PendingRow::from_iter(timepoint1.clone(), components2);
1829 let row3 = PendingRow::from_iter(timepoint2.clone(), components3);
1830 let row4 = PendingRow::from_iter(timepoint3.clone(), components4);
1831
1832 let entity_path1: EntityPath = "a/b/c".into();
1833 batcher.push_row(entity_path1.clone(), row1.clone());
1834 batcher.push_row(entity_path1.clone(), row2.clone());
1835 batcher.push_row(entity_path1.clone(), row3.clone());
1836 batcher.push_row(entity_path1.clone(), row4.clone());
1837
1838 let chunks_rx = batcher.chunks();
1839 drop(batcher); let mut chunks = Vec::new();
1842 loop {
1843 let chunk = match chunks_rx.try_recv() {
1844 Ok(chunk) => chunk,
1845 Err(TryRecvError::Empty) => panic!("expected chunk, got none"),
1846 Err(TryRecvError::Disconnected) => break,
1847 };
1848 chunks.push(chunk);
1849 }
1850
1851 chunks.sort_by_key(|chunk| chunk.row_id_range().unwrap().0);
1852
1853 eprintln!("Chunks:");
1855 for chunk in &chunks {
1856 eprintln!("{chunk}");
1857 }
1858
1859 assert_eq!(1, chunks.len());
1860
1861 {
1862 let expected_row_ids = vec![row1.row_id, row2.row_id, row3.row_id, row4.row_id];
1863 let expected_timelines = [
1864 (
1865 *timeline1.name(),
1866 TimeColumn::new(Some(false), timeline1, vec![45, 42, 43, 44].into()),
1867 ),
1868 (
1869 *timeline2.name(),
1870 TimeColumn::new(Some(false), timeline2, vec![1003, 1000, 1001, 1002].into()),
1871 ),
1872 ];
1873 let expected_components = [(
1874 MyPoints::descriptor_points(),
1875 arrays_to_list_array_opt(&[&*points1, &*points2, &*points3, &*points4].map(Some))
1876 .unwrap(),
1877 )];
1878 let expected_chunk = Chunk::from_native_row_ids(
1879 chunks[0].id,
1880 entity_path1.clone(),
1881 None,
1882 &expected_row_ids,
1883 expected_timelines.into_iter().collect(),
1884 expected_components.into_iter().collect(),
1885 )?;
1886
1887 eprintln!("Expected:\n{expected_chunk}");
1888 eprintln!("Got:\n{}", chunks[0]);
1889 assert_eq!(expected_chunk, chunks[0]);
1890 }
1891
1892 Ok(())
1893 }
1894
1895 #[test]
1898 fn unsorted_timeline_above_threshold() -> anyhow::Result<()> {
1899 let batcher = ChunkBatcher::new(
1900 ChunkBatcherConfig {
1901 chunk_max_rows_if_unsorted: 3,
1902 ..ChunkBatcherConfig::NEVER
1903 },
1904 BatcherHooks::NONE,
1905 )?;
1906
1907 let timeline1 = Timeline::new_duration("log_time");
1908 let timeline2 = Timeline::new_duration("frame_nr");
1909
1910 let timepoint1 = TimePoint::default()
1911 .with(timeline2, 1000)
1912 .with(timeline1, 42);
1913 let timepoint2 = TimePoint::default()
1914 .with(timeline2, 1001)
1915 .with(timeline1, 43);
1916 let timepoint3 = TimePoint::default()
1917 .with(timeline2, 1002)
1918 .with(timeline1, 44);
1919 let timepoint4 = TimePoint::default()
1920 .with(timeline2, 1003)
1921 .with(timeline1, 45);
1922
1923 let points1 = MyPoint::to_arrow([MyPoint::new(1.0, 2.0), MyPoint::new(3.0, 4.0)])?;
1924 let points2 = MyPoint::to_arrow([MyPoint::new(10.0, 20.0), MyPoint::new(30.0, 40.0)])?;
1925 let points3 = MyPoint::to_arrow([MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)])?;
1926 let points4 =
1927 MyPoint::to_arrow([MyPoint::new(1000.0, 2000.0), MyPoint::new(3000.0, 4000.0)])?;
1928
1929 let components1 = [SerializedComponentBatch::new(
1930 points1.clone(),
1931 MyPoints::descriptor_points(),
1932 )];
1933 let components2 = [SerializedComponentBatch::new(
1934 points2.clone(),
1935 MyPoints::descriptor_points(),
1936 )];
1937 let components3 = [SerializedComponentBatch::new(
1938 points3.clone(),
1939 MyPoints::descriptor_points(),
1940 )];
1941 let components4 = [SerializedComponentBatch::new(
1942 points4.clone(),
1943 MyPoints::descriptor_points(),
1944 )];
1945
1946 let row1 = PendingRow::from_iter(timepoint4.clone(), components1);
1947 let row2 = PendingRow::from_iter(timepoint1.clone(), components2);
1948 let row3 = PendingRow::from_iter(timepoint2.clone(), components3);
1949 let row4 = PendingRow::from_iter(timepoint3.clone(), components4);
1950
1951 let entity_path1: EntityPath = "a/b/c".into();
1952 batcher.push_row(entity_path1.clone(), row1.clone());
1953 batcher.push_row(entity_path1.clone(), row2.clone());
1954 batcher.push_row(entity_path1.clone(), row3.clone());
1955 batcher.push_row(entity_path1.clone(), row4.clone());
1956
1957 let chunks_rx = batcher.chunks();
1958 drop(batcher); let mut chunks = Vec::new();
1961 loop {
1962 let chunk = match chunks_rx.try_recv() {
1963 Ok(chunk) => chunk,
1964 Err(TryRecvError::Empty) => panic!("expected chunk, got none"),
1965 Err(TryRecvError::Disconnected) => break,
1966 };
1967 chunks.push(chunk);
1968 }
1969
1970 chunks.sort_by_key(|chunk| chunk.row_id_range().unwrap().0);
1971
1972 eprintln!("Chunks:");
1974 for chunk in &chunks {
1975 eprintln!("{chunk}");
1976 }
1977
1978 assert_eq!(2, chunks.len());
1979
1980 {
1981 let expected_row_ids = vec![row1.row_id, row2.row_id, row3.row_id];
1982 let expected_timelines = [
1983 (
1984 *timeline1.name(),
1985 TimeColumn::new(Some(false), timeline1, vec![45, 42, 43].into()),
1986 ),
1987 (
1988 *timeline2.name(),
1989 TimeColumn::new(Some(false), timeline2, vec![1003, 1000, 1001].into()),
1990 ),
1991 ];
1992 let expected_components = [(
1993 MyPoints::descriptor_points(),
1994 arrays_to_list_array_opt(&[&*points1, &*points2, &*points3].map(Some)).unwrap(),
1995 )];
1996 let expected_chunk = Chunk::from_native_row_ids(
1997 chunks[0].id,
1998 entity_path1.clone(),
1999 None,
2000 &expected_row_ids,
2001 expected_timelines.into_iter().collect(),
2002 expected_components.into_iter().collect(),
2003 )?;
2004
2005 eprintln!("Expected:\n{expected_chunk}");
2006 eprintln!("Got:\n{}", chunks[0]);
2007 assert_eq!(expected_chunk, chunks[0]);
2008 }
2009
2010 {
2011 let expected_row_ids = vec![row4.row_id];
2012 let expected_timelines = [
2013 (
2014 *timeline1.name(),
2015 TimeColumn::new(Some(true), timeline1, vec![44].into()),
2016 ),
2017 (
2018 *timeline2.name(),
2019 TimeColumn::new(Some(true), timeline2, vec![1002].into()),
2020 ),
2021 ];
2022 let expected_components = [(
2023 MyPoints::descriptor_points(),
2024 arrays_to_list_array_opt(&[&*points4].map(Some)).unwrap(),
2025 )];
2026 let expected_chunk = Chunk::from_native_row_ids(
2027 chunks[1].id,
2028 entity_path1.clone(),
2029 None,
2030 &expected_row_ids,
2031 expected_timelines.into_iter().collect(),
2032 expected_components.into_iter().collect(),
2033 )?;
2034
2035 eprintln!("Expected:\n{expected_chunk}");
2036 eprintln!("Got:\n{}", chunks[1]);
2037 assert_eq!(expected_chunk, chunks[1]);
2038 }
2039
2040 Ok(())
2041 }
2042}