implement 'rerun merge'

Author: teh-cmc

crates/build/re_build_info/src/crate_version.rs

@@ -49,6 +49,35 @@ pub struct CrateVersion {
     pub meta: Option<Meta>,
 }
 
+impl Ord for CrateVersion {
+    #[inline]
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        let Self {
+            major,
+            minor,
+            patch,
+            meta: _,
+        } = self;
+
+        match major.cmp(&other.major) {
+            core::cmp::Ordering::Equal => {}
+            ord => return ord,
+        }
+        match minor.cmp(&other.minor) {
+            core::cmp::Ordering::Equal => {}
+            ord => return ord,
+        }
+        patch.cmp(&other.patch)
+    }
+}
+
+impl PartialOrd for CrateVersion {
+    #[inline]
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
 impl CrateVersion {
     pub const LOCAL: Self = Self::parse(env!("CARGO_PKG_VERSION"));
 }

crates/top/rerun/src/run.rs

@@ -325,7 +325,7 @@ enum RrdCommands {
 
     /// Merges the contents of multiple .rrd and/or .rbl files, and writes the result to a new file.
     ///
-    /// Example: rerun merge -i input1.rrd -i input2.rbl -input3.rrd -o output.rrd`
+    /// Example: rerun merge -i input1.rrd -i input2.rbl -i input3.rrd -o output.rrd`
     Merge {
         #[arg(
             short = 'i',
@@ -497,10 +497,9 @@ fn run_rrd_commands(cmd: &RrdCommands) -> anyhow::Result<()> {
             path_to_input_rrds,
             path_to_output_rrd,
         } => {
-            // path_to_input_rrds.into_iter()_
-            // let path_to_input_rrd = PathBuf::from(path_to_input_rrd);
-            // let path_to_output_rrd = PathBuf::from(path_to_output_rrd);
-            // run_compact(&path_to_input_rrd, &path_to_output_rrd)
+            let path_to_input_rrds = path_to_input_rrds.iter().map(PathBuf::from).collect_vec();
+            let path_to_output_rrd = PathBuf::from(path_to_output_rrd);
+            run_merge(&path_to_input_rrds, &path_to_output_rrd)
         }
     }
 }
@@ -619,7 +618,7 @@ fn run_compact(path_to_input_rrd: &Path, path_to_output_rrd: &Path) -> anyhow::R
 
     use re_viewer::external::re_chunk_store::ChunkStoreConfig;
     let mut store_config = ChunkStoreConfig::from_env().unwrap_or_default();
-    // NOTE: We're doing headless, there's no point in running subscribers, it will just
+    // NOTE: We're doing headless processing, there's no point in running subscribers, it will just
     // (massively) slow us down.
     store_config.enable_changelog = false;
 
@@ -658,7 +657,7 @@ fn run_compact(path_to_input_rrd: &Path, path_to_output_rrd: &Path) -> anyhow::R
     );
 
     let mut rrd_out = std::fs::File::create(path_to_output_rrd)
-        .with_context(|| format!("{path_to_input_rrd:?}"))?;
+        .with_context(|| format!("{path_to_output_rrd:?}"))?;
 
     let messages: Result<Vec<Vec<LogMsg>>, _> = entity_dbs
         .into_values()
@@ -696,6 +695,101 @@ fn run_compact(path_to_input_rrd: &Path, path_to_output_rrd: &Path) -> anyhow::R
     Ok(())
 }
 
+fn run_merge(path_to_input_rrds: &[PathBuf], path_to_output_rrd: &Path) -> anyhow::Result<()> {
+    use re_entity_db::EntityDb;
+    use re_log_types::StoreId;
+
+    let rrds_in: Result<Vec<_>, _> = path_to_input_rrds
+        .iter()
+        .map(|path_to_input_rrd| {
+            std::fs::File::open(path_to_input_rrd).with_context(|| format!("{path_to_input_rrd:?}"))
+        })
+        .collect();
+    let rrds_in = rrds_in?;
+
+    let rrds_in_size = rrds_in
+        .iter()
+        .map(|rrd_in| rrd_in.metadata().ok().map(|md| md.len()))
+        .sum::<Option<u64>>();
+
+    let file_size_to_string = |size: Option<u64>| {
+        size.map_or_else(
+            || "<unknown>".to_owned(),
+            |size| re_format::format_bytes(size as _),
+        )
+    };
+
+    use re_viewer::external::re_chunk_store::ChunkStoreConfig;
+    let mut store_config = ChunkStoreConfig::from_env().unwrap_or_default();
+    // NOTE: We're doing headless processing, there's no point in running subscribers, it will just
+    // (massively) slow us down.
+    store_config.enable_changelog = false;
+
+    re_log::info!(
+        srcs = ?path_to_input_rrds,
+        dst = ?path_to_output_rrd,
+        max_num_rows = %re_format::format_uint(store_config.chunk_max_rows),
+        max_num_bytes = %re_format::format_bytes(store_config.chunk_max_bytes as _),
+        "merge started"
+    );
+
+    let now = std::time::Instant::now();
+
+    let mut entity_dbs: std::collections::HashMap<StoreId, EntityDb> = Default::default();
+    let mut version = None;
+    for rrd_in in rrds_in {
+        let version_policy = re_log_encoding::decoder::VersionPolicy::Warn;
+        let decoder = re_log_encoding::decoder::Decoder::new(version_policy, rrd_in)?;
+        version = version.max(Some(decoder.version()));
+        for msg in decoder {
+            let msg = msg.context("decode rrd message")?;
+            entity_dbs
+                .entry(msg.store_id().clone())
+                .or_insert_with(|| {
+                    re_entity_db::EntityDb::with_store_config(
+                        msg.store_id().clone(),
+                        store_config.clone(),
+                    )
+                })
+                .add(&msg)
+                .context("decode rrd file contents")?;
+        }
+    }
+
+    anyhow::ensure!(
+        !entity_dbs.is_empty(),
+        "no recordings found in rrd/rbl files"
+    );
+
+    let mut rrd_out = std::fs::File::create(path_to_output_rrd)
+        .with_context(|| format!("{path_to_output_rrd:?}"))?;
+
+    let messages: Result<Vec<Vec<LogMsg>>, _> = entity_dbs
+        .into_values()
+        .map(|entity_db| entity_db.to_messages(None /* time selection */))
+        .collect();
+    let messages = messages?;
+    let messages = messages.iter().flatten();
+
+    let encoding_options = re_log_encoding::EncodingOptions::COMPRESSED;
+    let version = version.unwrap_or(re_build_info::CrateVersion::LOCAL);
+    re_log_encoding::encoder::encode(version, encoding_options, messages, &mut rrd_out)
+        .context("Message encode")?;
+
+    let rrd_out_size = rrd_out.metadata().ok().map(|md| md.len());
+
+    re_log::info!(
+        srcs = ?path_to_input_rrds,
+        srcs_size_bytes = %file_size_to_string(rrds_in_size),
+        dst = ?path_to_output_rrd,
+        dst_size_bytes = %file_size_to_string(rrd_out_size),
+        time = ?now.elapsed(),
+        "merge finished"
+    );
+
+    Ok(())
+}
+
 impl PrintCommand {
     fn run(&self) -> anyhow::Result<()> {
         let rrd_path = PathBuf::from(&self.rrd_path);