re_chunk/
helpers.rs

1use std::sync::Arc;
2
3use arrow::array::{Array as _, ArrayRef as ArrowArrayRef};
4use re_log::debug_assert;
5use re_log_types::{TimeInt, TimelineName};
6use re_types_core::{Component, ComponentIdentifier};
7
8use crate::{Chunk, ChunkResult, RowId};
9
10// --- Helpers ---
11
12impl Chunk {
13    // --- Batch ---
14
15    /// Returns the raw data for the specified component.
16    ///
17    /// Returns an error if the row index is out of bounds.
18    #[inline]
19    pub fn component_batch_raw(
20        &self,
21        component: ComponentIdentifier,
22        row_index: usize,
23    ) -> Option<ChunkResult<ArrowArrayRef>> {
24        let list_array = self.components.get_array(component)?;
25        if list_array.len() > row_index {
26            list_array
27                .is_valid(row_index)
28                .then(|| Ok(list_array.value(row_index)))
29        } else {
30            Some(Err(crate::ChunkError::IndexOutOfBounds {
31                kind: "row".to_owned(),
32                len: list_array.len(),
33                index: row_index,
34            }))
35        }
36    }
37
38    /// Returns the deserialized data for the specified component.
39    ///
40    /// Returns an error if the data cannot be deserialized, or if the row index is out of bounds.
41    #[inline]
42    pub fn component_batch<C: Component>(
43        &self,
44        component: ComponentIdentifier,
45        row_index: usize,
46    ) -> Option<ChunkResult<Vec<C>>> {
47        let res = self.component_batch_raw(component, row_index)?;
48
49        let array = match res {
50            Ok(array) => array,
51            Err(err) => return Some(Err(err)),
52        };
53
54        let data = C::from_arrow(&*array);
55        Some(data.map_err(Into::into))
56    }
57
58    // --- Instance ---
59
60    /// Returns the raw data for the specified component at the given instance index.
61    ///
62    /// Returns an error if either the row index or instance index are out of bounds.
63    #[inline]
64    pub fn component_instance_raw(
65        &self,
66        component: ComponentIdentifier,
67        row_index: usize,
68        instance_index: usize,
69    ) -> Option<ChunkResult<ArrowArrayRef>> {
70        let res = self.component_batch_raw(component, row_index)?;
71
72        let array = match res {
73            Ok(array) => array,
74            Err(err) => return Some(Err(err)),
75        };
76
77        if array.len() > instance_index {
78            Some(Ok(array.slice(instance_index, 1)))
79        } else {
80            Some(Err(crate::ChunkError::IndexOutOfBounds {
81                kind: "instance".to_owned(),
82                len: array.len(),
83                index: instance_index,
84            }))
85        }
86    }
87
88    /// Returns the component data of the specified instance.
89    ///
90    /// Returns an error if the data cannot be deserialized, or if either the row index or instance index
91    /// are out of bounds.
92    #[inline]
93    pub fn component_instance<C: Component>(
94        &self,
95        component: ComponentIdentifier,
96        row_index: usize,
97        instance_index: usize,
98    ) -> Option<ChunkResult<C>> {
99        let res = self.component_instance_raw(component, row_index, instance_index)?;
100
101        let array = match res {
102            Ok(array) => array,
103            Err(err) => return Some(Err(err)),
104        };
105
106        match C::from_arrow(&*array) {
107            Ok(data) => data.into_iter().next().map(Ok), // NOTE: It's already sliced!
108            Err(err) => Some(Err(err.into())),
109        }
110    }
111
112    // --- Mono ---
113
114    /// Returns the raw data for the specified component, assuming a mono-batch.
115    ///
116    /// Returns an error if either the row index is out of bounds, or the underlying batch is not
117    /// of unit length.
118    #[inline]
119    pub fn component_mono_raw(
120        &self,
121        component: ComponentIdentifier,
122        row_index: usize,
123    ) -> Option<ChunkResult<ArrowArrayRef>> {
124        let res = self.component_batch_raw(component, row_index)?;
125
126        let array = match res {
127            Ok(array) => array,
128            Err(err) => return Some(Err(err)),
129        };
130
131        if array.len() == 1 {
132            Some(Ok(array.slice(0, 1)))
133        } else {
134            Some(Err(crate::ChunkError::IndexOutOfBounds {
135                kind: "mono".to_owned(),
136                len: array.len(),
137                index: 0,
138            }))
139        }
140    }
141
142    /// Returns the deserialized data for the specified component, assuming a mono-batch.
143    ///
144    /// Returns an error if the data cannot be deserialized, or if either the row index is out of bounds,
145    /// or the underlying batch is not of unit length.
146    #[inline]
147    pub fn component_mono<C: Component>(
148        &self,
149        component: ComponentIdentifier,
150        row_index: usize,
151    ) -> Option<ChunkResult<C>> {
152        let res = self.component_mono_raw(component, row_index)?;
153
154        let array = match res {
155            Ok(array) => array,
156            Err(err) => return Some(Err(err)),
157        };
158
159        match C::from_arrow(&*array) {
160            Ok(data) => data.into_iter().next().map(Ok), // NOTE: It's already sliced!
161            Err(err) => Some(Err(err.into())),
162        }
163    }
164}
165
166// --- Unit ---
167
168/// A simple type alias for an `Arc<Chunk>`.
169pub type ChunkShared = Arc<Chunk>;
170
171/// A [`ChunkShared`] that is guaranteed to always contain a single row's worth of data.
172#[derive(Debug, Clone, PartialEq)]
173pub struct UnitChunkShared(ChunkShared);
174
175impl std::ops::Deref for UnitChunkShared {
176    type Target = Chunk;
177
178    #[inline]
179    fn deref(&self) -> &Self::Target {
180        &self.0
181    }
182}
183
184impl re_byte_size::SizeBytes for UnitChunkShared {
185    #[inline]
186    fn heap_size_bytes(&self) -> u64 {
187        Chunk::heap_size_bytes(&self.0)
188    }
189}
190
191impl Chunk {
192    /// Turns the chunk into a [`UnitChunkShared`], if possible.
193    #[inline]
194    pub fn to_unit(self: &ChunkShared) -> Option<UnitChunkShared> {
195        (self.num_rows() == 1).then(|| UnitChunkShared(Arc::clone(self)))
196    }
197
198    /// Turns the chunk into a [`UnitChunkShared`], if possible.
199    #[inline]
200    pub fn into_unit(self) -> Option<UnitChunkShared> {
201        (self.num_rows() == 1).then(|| UnitChunkShared(Arc::new(self)))
202    }
203}
204
205impl UnitChunkShared {
206    /// Turns the unit chunk back into a standard [`Chunk`].
207    #[inline]
208    pub fn into_chunk(self) -> ChunkShared {
209        self.0
210    }
211}
212
213impl UnitChunkShared {
214    /// Returns the index (`(TimeInt, RowId)` pair) of the single row within, on the given timeline.
215    ///
216    /// Returns the single static index if the chunk is static.
217    #[inline]
218    pub fn index(&self, timeline: &TimelineName) -> Option<(TimeInt, RowId)> {
219        debug_assert!(self.num_rows() == 1);
220        if self.is_static() {
221            self.row_ids()
222                .next()
223                .map(|row_id| (TimeInt::STATIC, row_id))
224        } else {
225            let time_column = self.timelines.get(timeline)?;
226            let time = time_column.times().next()?;
227            self.row_ids().next().map(|row_id| (time, row_id))
228        }
229    }
230
231    /// Returns the [`RowId`] of the single row within, on the given timeline.
232    ///
233    /// Returns the single static `RowId` if the chunk is static.
234    // TODO(emilk): this is infallible, so remove the `Option` returntype
235    #[inline]
236    pub fn row_id(&self) -> Option<RowId> {
237        debug_assert!(self.num_rows() == 1);
238        self.row_ids().next()
239    }
240
241    /// Returns the number of instances of the single row within for a given component.
242    #[inline]
243    pub fn num_instances(&self, component: ComponentIdentifier) -> u64 {
244        debug_assert!(self.num_rows() == 1);
245        self.components
246            .values()
247            .filter(|column| column.descriptor.component == component)
248            .map(|column| {
249                let array = column.list_array.value(0);
250                array.nulls().map_or_else(
251                    || array.len(),
252                    |validity| validity.len() - validity.null_count(),
253                )
254            })
255            .max()
256            .unwrap_or(0) as u64
257    }
258}
259
260// --- Unit helpers ---
261
262impl UnitChunkShared {
263    // --- Batch ---
264
265    /// Returns the raw data for the specified component.
266    #[inline]
267    pub fn component_batch_raw(&self, component: ComponentIdentifier) -> Option<ArrowArrayRef> {
268        debug_assert!(self.num_rows() == 1);
269        let list_array = self.components.get_array(component)?;
270        list_array.is_valid(0).then(|| list_array.value(0))
271    }
272
273    /// Like `component_batch_raw`, but also returns `None` if the batch is empty.
274    #[inline]
275    pub fn non_empty_component_batch_raw(
276        &self,
277        component: ComponentIdentifier,
278    ) -> Option<(Option<RowId>, ArrowArrayRef)> {
279        let batch = self.component_batch_raw(component)?;
280        if batch.is_empty() {
281            None
282        } else {
283            Some((self.row_id(), batch))
284        }
285    }
286
287    /// Returns the deserialized data for the specified component.
288    ///
289    /// Returns an error if the data cannot be deserialized.
290    /// In debug builds, panics if the descriptor doesn't have the same type as the component type.
291    #[inline]
292    pub fn component_batch<C: Component>(
293        &self,
294        component: ComponentIdentifier,
295    ) -> Option<ChunkResult<Vec<C>>> {
296        let data = C::from_arrow(&*self.component_batch_raw(component)?);
297        Some(data.map_err(Into::into))
298    }
299
300    // --- Instance ---
301
302    /// Returns the raw data for the specified component at the given instance index.
303    ///
304    /// Returns an error if the instance index is out of bounds.
305    #[inline]
306    pub fn component_instance_raw(
307        &self,
308        component: ComponentIdentifier,
309        instance_index: usize,
310    ) -> Option<ChunkResult<ArrowArrayRef>> {
311        let array = self.component_batch_raw(component)?;
312        if array.len() > instance_index {
313            Some(Ok(array.slice(instance_index, 1)))
314        } else {
315            Some(Err(crate::ChunkError::IndexOutOfBounds {
316                kind: "instance".to_owned(),
317                len: array.len(),
318                index: instance_index,
319            }))
320        }
321    }
322
323    /// Returns the deserialized data for the specified component at the given instance index.
324    ///
325    /// Returns an error if the data cannot be deserialized, or if the instance index is out of bounds.
326    /// In debug builds, panics if the descriptor doesn't have the same type as the component type.
327    #[inline]
328    pub fn component_instance<C: Component>(
329        &self,
330        component: ComponentIdentifier,
331        instance_index: usize,
332    ) -> Option<ChunkResult<C>> {
333        let res = self.component_instance_raw(component, instance_index)?;
334
335        let array = match res {
336            Ok(array) => array,
337            Err(err) => return Some(Err(err)),
338        };
339
340        match C::from_arrow(&*array) {
341            Ok(data) => data.into_iter().next().map(Ok), // NOTE: It's already sliced!
342            Err(err) => Some(Err(err.into())),
343        }
344    }
345
346    // --- Mono ---
347
348    /// Returns the raw data for the specified component, assuming a mono-batch.
349    ///
350    /// Returns an error if the underlying batch is not of unit length.
351    #[inline]
352    pub fn component_mono_raw(
353        &self,
354        component: ComponentIdentifier,
355    ) -> Option<ChunkResult<ArrowArrayRef>> {
356        let array = self.component_batch_raw(component)?;
357        if array.len() == 1 {
358            Some(Ok(array.slice(0, 1)))
359        } else {
360            Some(Err(crate::ChunkError::IndexOutOfBounds {
361                kind: "mono".to_owned(),
362                len: array.len(),
363                index: 0,
364            }))
365        }
366    }
367
368    /// Returns the deserialized data for the specified component, assuming a mono-batch.
369    ///
370    /// Returns an error if the data cannot be deserialized, or if the underlying batch is not of unit length.
371    #[inline]
372    pub fn component_mono<C: Component>(
373        &self,
374        component: ComponentIdentifier,
375    ) -> Option<ChunkResult<C>> {
376        let res = self.component_mono_raw(component)?;
377
378        let array = match res {
379            Ok(array) => array,
380            Err(err) => return Some(Err(err)),
381        };
382
383        match C::from_arrow(&*array) {
384            Ok(data) => data.into_iter().next().map(Ok), // NOTE: It's already sliced!
385            Err(err) => Some(Err(err.into())),
386        }
387    }
388}
re_chunk/helpers.rs

re_chunk/
helpers.rs
