From 6c95b565329bd378e31e354ab89df698cfa8bd7d Mon Sep 17 00:00:00 2001
From: Emil Ernerfeldt <emil.ernerfeldt@gmail.com>
Date: Thu, 16 Jan 2025 08:30:13 +0100
Subject: [PATCH] Less use of TransportChunk

---
 .../re_log_encoding/src/codec/wire/encoder.rs    | 16 +++++++++-------
 rerun_py/src/remote.rs                           | 15 ++++++---------
 2 files changed, 15 insertions(+), 16 deletions(-)

diff --git a/crates/store/re_log_encoding/src/codec/wire/encoder.rs b/crates/store/re_log_encoding/src/codec/wire/encoder.rs
index aaf10c288246..1c8b657fda50 100644
--- a/crates/store/re_log_encoding/src/codec/wire/encoder.rs
+++ b/crates/store/re_log_encoding/src/codec/wire/encoder.rs
@@ -1,18 +1,20 @@
-use crate::codec::arrow::write_arrow_to_bytes;
-use crate::codec::CodecError;
-use re_chunk::TransportChunk;
 use re_protos::common::v0::RerunChunk;
 use re_protos::remote_store::v0::DataframePart;
 
+use arrow::array::RecordBatch as ArrowRecordBatch;
+
+use crate::codec::arrow::write_arrow_to_bytes;
+use crate::codec::CodecError;
+
 /// Encode a transport chunk into a byte stream.
 fn encode(
     version: re_protos::common::v0::EncoderVersion,
-    chunk: &TransportChunk,
+    batch: &ArrowRecordBatch,
 ) -> Result<Vec<u8>, CodecError> {
     match version {
         re_protos::common::v0::EncoderVersion::V0 => {
             let mut data: Vec<u8> = Vec::new();
-            write_arrow_to_bytes(&mut data, chunk)?;
+            write_arrow_to_bytes(&mut data, batch)?;
             Ok(data)
         }
     }
@@ -23,7 +25,7 @@ pub trait Encode<O> {
     fn encode(&self) -> Result<O, CodecError>;
 }
 
-impl Encode<DataframePart> for TransportChunk {
+impl Encode<DataframePart> for ArrowRecordBatch {
     fn encode(&self) -> Result<DataframePart, CodecError> {
         let payload = encode(re_protos::common::v0::EncoderVersion::V0, self)?;
         Ok(DataframePart {
@@ -33,7 +35,7 @@ impl Encode<DataframePart> for TransportChunk {
     }
 }
 
-impl Encode<RerunChunk> for TransportChunk {
+impl Encode<RerunChunk> for ArrowRecordBatch {
     fn encode(&self) -> Result<RerunChunk, CodecError> {
         let payload = encode(re_protos::common::v0::EncoderVersion::V0, self)?;
         Ok(RerunChunk {
diff --git a/rerun_py/src/remote.rs b/rerun_py/src/remote.rs
index 9de107726613..5959a385b6a9 100644
--- a/rerun_py/src/remote.rs
+++ b/rerun_py/src/remote.rs
@@ -1,4 +1,4 @@
-#![allow(unsafe_op_in_unsafe_fn)]
+#![allow(unsafe_op_in_unsafe_fn)] // False positive due to #[pyfunction] macro
 
 use std::collections::BTreeSet;
 
@@ -8,15 +8,16 @@ use arrow::{
     ffi_stream::ArrowArrayStreamReader,
     pyarrow::PyArrowType,
 };
-// False positive due to #[pyfunction] macro
 use pyo3::{
     exceptions::{PyRuntimeError, PyTypeError, PyValueError},
     prelude::*,
     types::PyDict,
     Bound, PyResult,
 };
+use tokio_stream::StreamExt;
+
 use re_arrow_util::ArrowArrayDowncastRef as _;
-use re_chunk::{Chunk, TransportChunk};
+use re_chunk::Chunk;
 use re_chunk_store::ChunkStore;
 use re_dataframe::{ChunkStoreHandle, QueryExpression, SparseFillStrategy, ViewContentsSelector};
 use re_grpc_client::TonicStatusError;
@@ -32,7 +33,6 @@ use re_protos::{
     TypeConversionError,
 };
 use re_sdk::{ApplicationId, ComponentName, StoreId, StoreKind, Time, Timeline};
-use tokio_stream::StreamExt;
 
 use crate::dataframe::{ComponentLike, PyRecording, PyRecordingHandle, PyRecordingView, PySchema};
 
@@ -321,8 +321,7 @@ impl PyStorageNodeClient {
                         ));
                     }
 
-                    let metadata_tc = TransportChunk::from(metadata);
-                    metadata_tc
+                    metadata
                         .encode()
                         .map_err(|err| PyRuntimeError::new_err(err.to_string()))
                 })
@@ -388,11 +387,9 @@ impl PyStorageNodeClient {
                 ));
             }
 
-            let metadata_tc = TransportChunk::from(metadata);
-
             let request = UpdateCatalogRequest {
                 metadata: Some(
-                    metadata_tc
+                    metadata
                         .encode()
                         .map_err(|err| PyRuntimeError::new_err(err.to_string()))?,
                 ),