Convert max-image-subscriber

rerun-io · Dec 19, 2024 · a6ba10b · a6ba10b
1 parent 8171d18
commit a6ba10b
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 5 deletions.
diff --git a/crates/store/re_chunk/src/iter.rs b/crates/store/re_chunk/src/iter.rs
@@ -1,5 +1,6 @@
 use std::sync::Arc;
 
+use arrow::array::ArrayRef as ArrowArrayRef;
 use arrow2::{
     array::{
         Array as Arrow2Array, BooleanArray as Arrow2BooleanArray,
@@ -213,6 +214,27 @@ impl Chunk {
     pub fn iter_component_arrays(
         &self,
         component_name: &ComponentName,
+    ) -> impl Iterator<Item = ArrowArrayRef> + '_ {
+        if let Some(list_array) = self.get_first_component(component_name) {
+            Either::Left(list_array.iter().flatten().map(|c| c.into()))
+        } else {
+            Either::Right(std::iter::empty())
+        }
+    }
+
+    /// Returns an iterator over the raw arrays of a [`Chunk`], for a given component.
+    ///
+    /// See also:
+    /// * [`Self::iter_primitive`]
+    /// * [`Self::iter_primitive_array`]
+    /// * [`Self::iter_primitive_array_list`]
+    /// * [`Self::iter_string`]
+    /// * [`Self::iter_buffer`].
+    /// * [`Self::iter_component`].
+    #[inline]
+    pub fn iter_component_arrays_arrow2(
+        &self,
+        component_name: &ComponentName,
     ) -> impl Iterator<Item = Box<dyn Arrow2Array>> + '_ {
         let Some(list_array) = self.get_first_component(component_name) else {
             return Either::Left(std::iter::empty());

diff --git a/crates/store/re_grpc_client/src/lib.rs b/crates/store/re_grpc_client/src/lib.rs
@@ -484,7 +484,7 @@ async fn stream_catalog_async(
             )))?;
 
         let recording_uri_arrays: Vec<Box<dyn Arrow2Array>> = chunk
-            .iter_component_arrays(&"id".into())
+            .iter_component_arrays_arrow2(&"id".into())
             .map(|id| {
                 let rec_id = id
                     .as_any()

diff --git a/crates/viewer/re_view_spatial/src/max_image_dimension_subscriber.rs b/crates/viewer/re_view_spatial/src/max_image_dimension_subscriber.rs
@@ -1,4 +1,3 @@
-use arrow2::array::Array;
 use nohash_hasher::IntMap;
 use once_cell::sync::OnceCell;
 
@@ -104,13 +103,16 @@ impl PerStoreChunkSubscriber for MaxImageDimensionsStoreSubscriber {
     }
 }
 
-fn size_from_blob(blob: &dyn Array, media_type: Option<&dyn Array>) -> Option<[u32; 2]> {
+fn size_from_blob(
+    blob: &dyn arrow::array::Array,
+    media_type: Option<&dyn arrow::array::Array>,
+) -> Option<[u32; 2]> {
     re_tracing::profile_function!();
 
-    let blob = Blob::from_arrow2_opt(blob).ok()?.first()?.clone()?;
+    let blob = Blob::from_arrow_opt(blob).ok()?.first()?.clone()?;
 
     let media_type: Option<MediaType> = media_type
-        .and_then(|media_type| MediaType::from_arrow2_opt(media_type).ok())
+        .and_then(|media_type| MediaType::from_arrow_opt(media_type).ok())
         .and_then(|list| list.first().cloned())
         .flatten();