From d17a6aa82cd46a13b03d3e6d5c86874b31f62966 Mon Sep 17 00:00:00 2001
From: Emil Ernerfeldt <emil.ernerfeldt@gmail.com>
Date: Wed, 18 Dec 2024 20:51:42 +0100
Subject: [PATCH] Convert max-image-subscriber

---
 crates/store/re_chunk/src/iter.rs             | 22 +++++++++++++++++++
 crates/store/re_grpc_client/src/lib.rs        |  2 +-
 .../src/max_image_dimension_subscriber.rs     | 10 +++++----
 3 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/crates/store/re_chunk/src/iter.rs b/crates/store/re_chunk/src/iter.rs
index 1534d9ac0a28..de5d045f4ecb 100644
--- 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, FixedSizeListArray as Arrow2FixedSizeListArray,
@@ -212,6 +213,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
index f324862dc3a8..b38b16520016 100644
--- 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
index 735cba3f50ff..75270f0558f3 100644
--- 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();