Encode LogMsg using protobuf

Cargo.lock

@@ -5826,6 +5826,7 @@ dependencies = [
  "re_chunk",
  "re_chunk_store",
  "re_log",
+ "re_log_encoding",
  "re_log_types",
  "re_query",
  "re_tracing",

crates/build/re_protos_builder/src/bin/build_re_remote_store_types.rs

@@ -1,15 +1,15 @@
 //! This binary runs the remote store gRPC service codegen manually.
 //!
-//! It is easiest to call this using `pixi run codegen-protos`,
+//! It is easiest to call this using `pixi run codegen-protos   `,
 //! which will set up the necessary tools.
 
 #![allow(clippy::unwrap_used)]
 
 use camino::Utf8Path;
 
-const PROTOBUF_DEFINITIONS_DIR_PATH: &str = "crates/store/re_protos/proto";
-const PROTOBUF_REMOTE_STORE_V0_RELATIVE_PATH: &str = "rerun/v0/remote_store.proto";
-const RUST_V0_OUTPUT_DIR_PATH: &str = "crates/store/re_protos/src/v0";
+const PROTOS_DIR: &str = "crates/store/re_protos/proto";
+const INPUT_V0: &[&str] = &["rerun/v0/remote_store.proto", "rerun/v0/log_msg.proto"];
+const OUTPUT_V0_RUST_DIR: &str = "crates/store/re_protos/src/v0";
 
 fn main() {
     re_log::setup_logging();
@@ -26,8 +26,8 @@ fn main() {
         "failed to find workspace root"
     );
 
-    let definitions_dir_path = workspace_dir.join(PROTOBUF_DEFINITIONS_DIR_PATH);
-    let rust_generated_output_dir_path = workspace_dir.join(RUST_V0_OUTPUT_DIR_PATH);
+    let definitions_dir_path = workspace_dir.join(PROTOS_DIR);
+    let rust_generated_output_dir_path = workspace_dir.join(OUTPUT_V0_RUST_DIR);
 
     re_log::info!(
         definitions=?definitions_dir_path,
@@ -37,7 +37,7 @@ fn main() {
 
     re_protos_builder::generate_rust_code(
         definitions_dir_path,
-        &[PROTOBUF_REMOTE_STORE_V0_RELATIVE_PATH],
+        INPUT_V0,
         rust_generated_output_dir_path,
     );
 }

crates/store/re_chunk_store/src/lib.rs

@@ -49,16 +49,14 @@ pub use re_chunk::{
     Chunk, ChunkId, ChunkShared, LatestAtQuery, RangeQuery, RangeQueryOptions, RowId,
     UnitChunkShared,
 };
-#[doc(no_inline)]
-pub use re_log_encoding::decoder::VersionPolicy;
+
 #[doc(no_inline)]
 pub use re_log_types::{ResolvedTimeRange, TimeInt, TimeType, Timeline};
 
 pub mod external {
     pub use arrow2;
 
     pub use re_chunk;
-    pub use re_log_encoding;
 }
 
 // ---

crates/store/re_chunk_store/src/store.rs

@@ -691,7 +691,7 @@ impl ChunkStore {
     pub fn from_rrd_filepath(
         store_config: &ChunkStoreConfig,
         path_to_rrd: impl AsRef<std::path::Path>,
-        version_policy: crate::VersionPolicy,
+        version_policy: re_log_encoding::VersionPolicy,
     ) -> anyhow::Result<BTreeMap<StoreId, Self>> {
         let path_to_rrd = path_to_rrd.as_ref();
 
@@ -808,7 +808,7 @@ impl ChunkStore {
     pub fn handle_from_rrd_filepath(
         store_config: &ChunkStoreConfig,
         path_to_rrd: impl AsRef<std::path::Path>,
-        version_policy: crate::VersionPolicy,
+        version_policy: re_log_encoding::VersionPolicy,
     ) -> anyhow::Result<BTreeMap<StoreId, ChunkStoreHandle>> {
         Ok(
             Self::from_rrd_filepath(store_config, path_to_rrd, version_policy)?

crates/store/re_data_loader/src/loader_external.rs

@@ -182,7 +182,7 @@ impl crate::DataLoader for ExternalLoader {
                 // streaming data to stdout.
                 let is_sending_data = Arc::new(AtomicBool::new(false));
 
-                let version_policy = re_log_encoding::decoder::VersionPolicy::Warn;
+                let version_policy = re_log_encoding::VersionPolicy::Warn;
                 let stdout = std::io::BufReader::new(stdout);
                 match re_log_encoding::decoder::Decoder::new(version_policy, stdout) {
                     Ok(decoder) => {

crates/store/re_data_loader/src/loader_rrd.rs

@@ -40,7 +40,7 @@ impl crate::DataLoader for RrdLoader {
             "Loading rrd data from filesystem…",
         );
 
-        let version_policy = re_log_encoding::decoder::VersionPolicy::Warn;
+        let version_policy = re_log_encoding::VersionPolicy::Warn;
 
         match extension.as_str() {
             "rbl" => {
@@ -118,7 +118,7 @@ impl crate::DataLoader for RrdLoader {
             return Err(crate::DataLoaderError::Incompatible(filepath));
         }
 
-        let version_policy = re_log_encoding::decoder::VersionPolicy::Warn;
+        let version_policy = re_log_encoding::VersionPolicy::Warn;
         let contents = std::io::Cursor::new(contents);
         let decoder = match re_log_encoding::decoder::Decoder::new(version_policy, contents) {
             Ok(decoder) => decoder,
@@ -308,7 +308,7 @@ impl RetryableFileReader {
 mod tests {
     use re_build_info::CrateVersion;
     use re_chunk::RowId;
-    use re_log_encoding::{decoder, encoder::DroppableEncoder};
+    use re_log_encoding::{encoder::DroppableEncoder, VersionPolicy};
     use re_log_types::{
         ApplicationId, LogMsg, SetStoreInfo, StoreId, StoreInfo, StoreKind, StoreSource, Time,
     };
@@ -372,7 +372,7 @@ mod tests {
         encoder.flush_blocking().expect("failed to flush messages");
 
         let reader = RetryableFileReader::new(&rrd_file_path).unwrap();
-        let mut decoder = Decoder::new(decoder::VersionPolicy::Warn, reader).unwrap();
+        let mut decoder = Decoder::new(VersionPolicy::Warn, reader).unwrap();
 
         // we should be able to read 5 messages that we wrote
         let decoded_messages = (0..5)

crates/store/re_data_source/src/load_stdin.rs

@@ -6,7 +6,7 @@ use re_smart_channel::Sender;
 /// This fails synchronously iff the standard input stream could not be opened, otherwise errors
 /// are handled asynchronously (as in: they're logged).
 pub fn load_stdin(tx: Sender<LogMsg>) -> anyhow::Result<()> {
-    let version_policy = re_log_encoding::decoder::VersionPolicy::Warn;
+    let version_policy = re_log_encoding::VersionPolicy::Warn;
 
     let stdin = std::io::BufReader::new(std::io::stdin());
     let decoder = re_log_encoding::decoder::Decoder::new_concatenated(version_policy, stdin)?;

crates/store/re_dataframe/Cargo.toml

@@ -24,16 +24,17 @@ all-features = true
 [features]
 default = []
 
-
 [dependencies]
 # Rerun dependencies:
 re_chunk.workspace = true
 re_chunk_store.workspace = true
 re_log.workspace = true
+re_log_encoding.workspace = true
 re_log_types.workspace = true
 re_query.workspace = true
 re_tracing.workspace = true
 re_types_core.workspace = true
+
 # External dependencies:
 anyhow.workspace = true
 arrow2.workspace = true

crates/store/re_dataframe/examples/query.rs

@@ -4,8 +4,9 @@ use itertools::Itertools;
 
 use re_dataframe::{
     ChunkStoreConfig, EntityPathFilter, QueryEngine, QueryExpression, ResolvedTimeRange,
-    SparseFillStrategy, StoreKind, TimeInt, Timeline, VersionPolicy,
+    SparseFillStrategy, StoreKind, TimeInt, Timeline,
 };
+use re_log_encoding::VersionPolicy;
 
 fn main() -> anyhow::Result<()> {
     let args = std::env::args().collect_vec();

crates/store/re_dataframe/src/engine.rs

@@ -3,7 +3,6 @@ use std::collections::BTreeMap;
 use re_chunk::{EntityPath, TransportChunk};
 use re_chunk_store::{
     ChunkStore, ChunkStoreConfig, ChunkStoreHandle, ColumnDescriptor, QueryExpression,
-    VersionPolicy,
 };
 use re_log_types::{EntityPathFilter, StoreId};
 use re_query::{QueryCache, QueryCacheHandle, StorageEngine, StorageEngineLike};
@@ -59,7 +58,7 @@ impl QueryEngine<StorageEngine> {
     pub fn from_rrd_filepath(
         store_config: &ChunkStoreConfig,
         path_to_rrd: impl AsRef<std::path::Path>,
-        version_policy: VersionPolicy,
+        version_policy: re_log_encoding::VersionPolicy,
     ) -> anyhow::Result<BTreeMap<StoreId, Self>> {
         Ok(
             ChunkStore::handle_from_rrd_filepath(store_config, path_to_rrd, version_policy)?

crates/store/re_dataframe/src/lib.rs

@@ -13,8 +13,7 @@ pub use self::external::re_chunk::{util::concatenate_record_batches, TransportCh
 #[doc(no_inline)]
 pub use self::external::re_chunk_store::{
     ChunkStoreConfig, ChunkStoreHandle, ColumnSelector, ComponentColumnSelector, Index, IndexRange,
-    IndexValue, QueryExpression, SparseFillStrategy, TimeColumnSelector, VersionPolicy,
-    ViewContentsSelector,
+    IndexValue, QueryExpression, SparseFillStrategy, TimeColumnSelector, ViewContentsSelector,
 };
 #[doc(no_inline)]
 pub use self::external::re_log_types::{

crates/store/re_entity_db/examples/memory_usage.rs

@@ -78,7 +78,7 @@ fn log_messages() {
     }
 
     fn decode_log_msg(mut bytes: &[u8]) -> LogMsg {
-        let version_policy = re_log_encoding::decoder::VersionPolicy::Error;
+        let version_policy = re_log_encoding::VersionPolicy::Error;
         let mut messages = re_log_encoding::decoder::Decoder::new(version_policy, &mut bytes)
             .unwrap()
             .collect::<Result<Vec<LogMsg>, _>>()

crates/store/re_entity_db/src/store_bundle.rs

@@ -1,7 +1,7 @@
 use itertools::Itertools as _;
 
 use crate::EntityDb;
-use re_log_encoding::decoder::VersionPolicy;
+use re_log_encoding::VersionPolicy;
 use re_log_types::{StoreId, StoreKind};
 
 #[derive(thiserror::Error, Debug)]

crates/store/re_grpc_client/src/lib.rs

@@ -151,7 +151,7 @@ async fn stream_recording_async(
             .connect()
             .await?;
 
-        StorageNodeClient::new(tonic_client)
+        StorageNodeClient::new(tonic_client).max_decoding_message_size(usize::MAX)
     };
 
     re_log::debug!("Fetching {recording_id}…");

crates/store/re_log_encoding/benches/msg_encode_benchmark.rs

@@ -46,7 +46,7 @@ fn encode_log_msgs(messages: &[LogMsg]) -> Vec<u8> {
 }
 
 fn decode_log_msgs(mut bytes: &[u8]) -> Vec<LogMsg> {
-    let version_policy = re_log_encoding::decoder::VersionPolicy::Error;
+    let version_policy = re_log_encoding::VersionPolicy::Error;
     let messages = re_log_encoding::decoder::Decoder::new(version_policy, &mut bytes)
         .unwrap()
         .collect::<Result<Vec<LogMsg>, _>>()

crates/store/re_log_encoding/src/codec/arrow.rs

@@ -0,0 +1,52 @@
+use super::CodecError;
+
+/// Helper function that serializes given arrow schema and record batch into bytes
+/// using Arrow IPC format.
+pub(crate) fn write_arrow_to_bytes<W: std::io::Write>(
+    writer: &mut W,
+    schema: &arrow2::datatypes::Schema,
+    data: &arrow2::chunk::Chunk<Box<dyn re_chunk::Arrow2Array>>,
+) -> Result<(), CodecError> {
+    use arrow2::io::ipc;
+
+    let options = ipc::write::WriteOptions { compression: None };
+    let mut sw = ipc::write::StreamWriter::new(writer, options);
+
+    sw.start(schema, None)
+        .map_err(CodecError::ArrowSerialization)?;
+    sw.write(data, None)
+        .map_err(CodecError::ArrowSerialization)?;
+    sw.finish().map_err(CodecError::ArrowSerialization)?;
+
+    Ok(())
+}
+
+/// Helper function that deserializes raw bytes into arrow schema and record batch
+/// using Arrow IPC format.
+pub(crate) fn read_arrow_from_bytes<R: std::io::Read>(
+    reader: &mut R,
+) -> Result<
+    (
+        arrow2::datatypes::Schema,
+        arrow2::chunk::Chunk<Box<dyn re_chunk::Arrow2Array>>,
+    ),
+    CodecError,
+> {
+    use arrow2::io::ipc;
+
+    let metadata =
+        ipc::read::read_stream_metadata(reader).map_err(CodecError::ArrowSerialization)?;
+    let mut stream = ipc::read::StreamReader::new(reader, metadata, None);
+
+    let schema = stream.schema().clone();
+    // there should be at least one record batch in the stream
+    let stream_state = stream
+        .next()
+        .ok_or(CodecError::MissingRecordBatch)?
+        .map_err(CodecError::ArrowSerialization)?;
+
+    match stream_state {
+        ipc::read::StreamState::Waiting => Err(CodecError::UnexpectedStreamState),
+        ipc::read::StreamState::Some(chunk) => Ok((schema, chunk)),
+    }
+}

crates/store/re_log_encoding/src/codec/file/decoder.rs

@@ -0,0 +1,107 @@
+use super::{MessageHeader, MessageKind};
+use crate::codec::arrow::read_arrow_from_bytes;
+use crate::decoder::DecodeError;
+use crate::{codec::CodecError, Compression};
+use re_log_types::LogMsg;
+use re_protos::TypeConversionError;
+
+impl MessageKind {
+    pub(crate) fn decode(data: &mut impl std::io::Read) -> Result<Self, DecodeError> {
+        let mut buf = [0; 4];
+        data.read_exact(&mut buf)?;
+
+        match u32::from_le_bytes(buf) {
+            1 => Ok(Self::SetStoreInfo),
+            2 => Ok(Self::ArrowMsg),
+            3 => Ok(Self::BlueprintActivationCommand),
+            255 => Ok(Self::End),
+            _ => Err(DecodeError::Codec(CodecError::UnknownMessageHeader)),
+        }
+    }
+}
+
+impl MessageHeader {
+    pub(crate) fn decode(data: &mut impl std::io::Read) -> Result<Self, DecodeError> {
+        let kind = MessageKind::decode(data)?;
+        let mut buf = [0; 4];
+        data.read_exact(&mut buf)?;
+        let len = u32::from_le_bytes(buf);
+
+        Ok(Self { kind, len })
+    }
+}
+
+pub(crate) fn decode(
+    data: &mut impl std::io::Read,
+    compression: Compression,
+) -> Result<(u64, Option<LogMsg>), DecodeError> {
+    use re_protos::external::prost::Message;
+    use re_protos::log_msg::v0::{ArrowMsg, BlueprintActivationCommand, Encoding, SetStoreInfo};
+
+    let mut read_bytes = 0u64;
+    let header = MessageHeader::decode(data)?;
+    read_bytes += std::mem::size_of::<MessageHeader>() as u64 + header.len as u64;
+
+    let mut buf = vec![0; header.len as usize];
+    data.read_exact(&mut buf[..])?;
+
+    let msg = match header.kind {
+        MessageKind::SetStoreInfo => {
+            let set_store_info = SetStoreInfo::decode(&buf[..])?;
+            Some(LogMsg::SetStoreInfo(set_store_info.try_into()?))
+        }
+        MessageKind::ArrowMsg => {
+            let arrow_msg = ArrowMsg::decode(&buf[..])?;
+            if arrow_msg.encoding() != Encoding::ArrowIpc {
+                return Err(DecodeError::Codec(CodecError::UnsupportedEncoding));
+            }
+
+            let (schema, chunk) = decode_arrow(&arrow_msg.payload, compression)?;
+
+            let store_id: re_log_types::StoreId = arrow_msg
+                .store_id
+                .ok_or_else(|| {
+                    TypeConversionError::missing_field("rerun.log_msg.v0.ArrowMsg", "store_id")
+                })?
+                .into();
+
+            let chunk = re_chunk::Chunk::from_transport(&re_chunk::TransportChunk {
+                schema,
+                data: chunk,
+            })?;
+
+            Some(LogMsg::ArrowMsg(store_id, chunk.to_arrow_msg()?))
+        }
+        MessageKind::BlueprintActivationCommand => {
+            let blueprint_activation_command = BlueprintActivationCommand::decode(&buf[..])?;
+            Some(LogMsg::BlueprintActivationCommand(
+                blueprint_activation_command.try_into()?,
+            ))
+        }
+        MessageKind::End => None,
+    };
+
+    Ok((read_bytes, msg))
+}
+
+fn decode_arrow(
+    data: &[u8],
+    compression: crate::Compression,
+) -> Result<
+    (
+        arrow2::datatypes::Schema,
+        arrow2::chunk::Chunk<Box<dyn re_chunk::Arrow2Array>>,
+    ),
+    DecodeError,
+> {
+    let mut uncompressed = Vec::new();
+    let data = match compression {
+        crate::Compression::Off => data,
+        crate::Compression::LZ4 => {
+            lz4_flex::block::decompress_into(data, &mut uncompressed)?;
+            uncompressed.as_slice()
+        }
+    };
+
+    Ok(read_arrow_from_bytes(&mut &data[..])?)
+}