Add arrow feature to re_chunk and conversions to RecordBatch (#7355)

### What
Basic type conversions from TransportChunk to RecordBatch and back.

Adding the round-trip test turned up an interesting issue.

TransportChunk <-> RecordBatch fails to round-trip successfully because
we lose the ExtensionType encapsulation that used to be encoded by
arrow2.

While on the surface this isn't immediately problematic, as we don't
care about ExtensionTypes, the discussion indicates there are in fact
going to be very real pain points when it comes to writing semantic data
processing engines using arrow-rs. This is because the metadata is
attached to the FIELD, not the DATATYPE, and there exist many processing
contexts where the context of that field itself is lost.
apache/arrow-rs#4472

### Checklist
* [x] I have read and agree to [Contributor
Guide](https://github.com/rerun-io/rerun/blob/main/CONTRIBUTING.md) and
the [Code of
Conduct](https://github.com/rerun-io/rerun/blob/main/CODE_OF_CONDUCT.md)
* [x] I've included a screenshot or gif (if applicable)
* [x] I have tested the web demo (if applicable):
* Using examples from latest `main` build:
[rerun.io/viewer](https://rerun.io/viewer/pr/7355?manifest_url=https://app.rerun.io/version/main/examples_manifest.json)
* Using full set of examples from `nightly` build:
[rerun.io/viewer](https://rerun.io/viewer/pr/7355?manifest_url=https://app.rerun.io/version/nightly/examples_manifest.json)
* [x] The PR title and labels are set such as to maximize their
usefulness for the next release's CHANGELOG
* [x] If applicable, add a new check to the [release
checklist](https://github.com/rerun-io/rerun/blob/main/tests/python/release_checklist)!
* [x] If have noted any breaking changes to the log API in
`CHANGELOG.md` and the migration guide

- [PR Build Summary](https://build.rerun.io/pr/7355)
- [Recent benchmark results](https://build.rerun.io/graphs/crates.html)
- [Wasm size tracking](https://build.rerun.io/graphs/sizes.html)

To run all checks from `main`, comment on the PR with `@rerun-bot
full-check`.

Author: jleibs

Cargo.lock

@@ -4784,6 +4784,7 @@ version = "0.19.0-alpha.1+dev"
 dependencies = [
  "ahash",
  "anyhow",
+ "arrow",
  "bytemuck",
  "criterion",
  "crossbeam",

crates/store/re_chunk/Cargo.toml

@@ -30,6 +30,9 @@ serde = [
   "re_types_core/serde",
 ]
 
+## Enable conversion to and from arrow-rs types
+arrow = ["arrow2/arrow", "dep:arrow"]
+
 
 [dependencies]
 
@@ -60,6 +63,7 @@ thiserror.workspace = true
 
 # Optional dependencies:
 serde = { workspace = true, optional = true, features = ["derive", "rc"] }
+arrow = { workspace = true, optional = true }
 
 # Native dependencies:
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]

crates/store/re_chunk/src/arrow.rs

@@ -0,0 +1,67 @@
+use arrow::{
+    array::{make_array, RecordBatch},
+    datatypes::{Field, Schema},
+    error::ArrowError,
+};
+
+use crate::TransportChunk;
+
+impl TransportChunk {
+    /// Create an arrow-rs [`RecordBatch`] containing the data from this [`TransportChunk`].
+    ///
+    /// This is a "fairly" cheap operation, as it does not copy the underlying arrow data,
+    /// but does incur overhead of generating an alternative representation of the arrow-
+    /// related rust structures that refer to those data buffers.
+    pub fn try_to_arrow_record_batch(&self) -> Result<RecordBatch, ArrowError> {
+        let fields: Vec<Field> = self
+            .schema
+            .fields
+            .iter()
+            .map(|f| f.clone().into())
+            .collect();
+
+        let metadata = self.schema.metadata.clone().into_iter().collect();
+
+        let schema = Schema::new(fields).with_metadata(metadata);
+
+        let columns: Vec<_> = self
+            .data
+            .columns()
+            .iter()
+            .map(|arr2_array| {
+                let data = arrow2::array::to_data(arr2_array.as_ref());
+                make_array(data)
+            })
+            .collect();
+
+        RecordBatch::try_new(std::sync::Arc::new(schema), columns)
+    }
+
+    /// Create a [`TransportChunk`] from an arrow-rs [`RecordBatch`].
+    ///
+    /// This is a "fairly" cheap operation, as it does not copy the underlying arrow data,
+    /// but does incur overhead of generating an alternative representation of the arrow-
+    /// related rust structures that refer to those data buffers.
+    pub fn from_arrow_record_batch(batch: &RecordBatch) -> Self {
+        let fields: Vec<arrow2::datatypes::Field> = batch
+            .schema()
+            .fields
+            .iter()
+            .map(|f| f.clone().into())
+            .collect();
+
+        let metadata = batch.schema().metadata.clone().into_iter().collect();
+
+        let schema = arrow2::datatypes::Schema::from(fields).with_metadata(metadata);
+
+        let columns: Vec<_> = batch
+            .columns()
+            .iter()
+            .map(|array| arrow2::array::from_data(&array.to_data()))
+            .collect();
+
+        let data = arrow2::chunk::Chunk::new(columns);
+
+        Self { schema, data }
+    }
+}

crates/store/re_chunk/src/chunk.rs

@@ -152,6 +152,71 @@ impl Chunk {
             }
             && *components == rhs.components
     }
+
+    /// Check for equality while ignoring possible `Extension` type information
+    ///
+    /// This is necessary because `arrow2` loses the `Extension` datatype
+    /// when deserializing back from the `arrow_schema::DataType` representation.
+    ///
+    /// In theory we could fix this, but as we're moving away from arrow2 anyways
+    /// it's unlikely worth the effort.
+    pub fn are_equal_ignoring_extension_types(&self, other: &Self) -> bool {
+        let Self {
+            id,
+            entity_path,
+            heap_size_bytes: _,
+            is_sorted,
+            row_ids,
+            timelines,
+            components,
+        } = self;
+
+        let row_ids_no_extension = arrow2::array::StructArray::new(
+            row_ids.data_type().to_logical_type().clone(),
+            row_ids.values().to_vec(),
+            row_ids.validity().cloned(),
+        );
+
+        let components_no_extension: BTreeMap<_, _> = components
+            .iter()
+            .map(|(name, arr)| {
+                let arr = arrow2::array::ListArray::new(
+                    arr.data_type().to_logical_type().clone(),
+                    arr.offsets().clone(),
+                    arr.values().clone(),
+                    arr.validity().cloned(),
+                );
+                (name, arr)
+            })
+            .collect();
+
+        let other_components_no_extension: BTreeMap<_, _> = other
+            .components
+            .iter()
+            .map(|(name, arr)| {
+                let arr = arrow2::array::ListArray::new(
+                    arr.data_type().to_logical_type().clone(),
+                    arr.offsets().clone(),
+                    arr.values().clone(),
+                    arr.validity().cloned(),
+                );
+                (name, arr)
+            })
+            .collect();
+
+        let other_row_ids_no_extension = arrow2::array::StructArray::new(
+            other.row_ids.data_type().to_logical_type().clone(),
+            other.row_ids.values().to_vec(),
+            other.row_ids.validity().cloned(),
+        );
+
+        *id == other.id
+            && *entity_path == other.entity_path
+            && *is_sorted == other.is_sorted
+            && row_ids_no_extension == other_row_ids_no_extension
+            && *timelines == other.timelines
+            && components_no_extension == other_components_no_extension
+    }
 }
 
 impl Clone for Chunk {

crates/store/re_chunk/src/lib.rs

@@ -20,6 +20,9 @@ pub mod util;
 #[cfg(not(target_arch = "wasm32"))]
 mod batcher;
 
+#[cfg(feature = "arrow")]
+mod arrow;
+
 pub use self::builder::{ChunkBuilder, TimeColumnBuilder};
 pub use self::chunk::{Chunk, ChunkError, ChunkResult, TimeColumn};
 pub use self::helpers::{ChunkShared, UnitChunkShared};

crates/store/re_chunk/src/transport.rs

@@ -711,7 +711,18 @@ mod tests {
 
             for _ in 0..3 {
                 let chunk_in_transport = chunk_before.to_transport()?;
-                let chunk_after = Chunk::from_transport(&chunk_in_transport)?;
+                let chunk_roundtrip;
+                #[cfg(feature = "arrow")]
+                {
+                    let chunk_in_record_batch = chunk_in_transport.try_to_arrow_record_batch()?;
+                    chunk_roundtrip =
+                        TransportChunk::from_arrow_record_batch(&chunk_in_record_batch);
+                }
+                #[cfg(not(feature = "arrow"))]
+                {
+                    chunk_roundtrip = &chunk_in_transport;
+                }
+                let chunk_after = Chunk::from_transport(&chunk_roundtrip)?;
 
                 assert_eq!(
                     chunk_in_transport.entity_path()?,
@@ -762,7 +773,14 @@ mod tests {
                 eprintln!("{chunk_in_transport}");
                 eprintln!("{chunk_after}");
 
-                assert_eq!(chunk_before, chunk_after);
+                #[cfg(not(feature = "arrow"))]
+                {
+                    // This will fail when round-tripping all the way to record-batch
+                    // the below check should always pass regardless.
+                    assert_eq!(chunk_before, &chunk_after);
+                }
+
+                assert!(chunk_before.are_equal_ignoring_extension_types(&chunk_after));
 
                 chunk_before = chunk_after;
             }