Introduce APIs for list_recordings and update_metadata

Cargo.lock

@@ -5779,6 +5779,7 @@ dependencies = [
  "prost",
  "re_arrow2",
  "re_dataframe",
+ "re_log",
  "re_log_types",
  "thiserror",
  "tonic",
@@ -6971,6 +6972,7 @@ dependencies = [
  "re_web_viewer_server",
  "re_ws_comms",
  "tokio",
+ "tokio-stream",
  "tonic",
  "url",
  "uuid",

crates/store/re_protos/Cargo.toml

@@ -15,6 +15,7 @@ version.workspace = true
 [dependencies]
 re_log_types.workspace = true
 re_dataframe.workspace = true
+re_log.workspace = true
 
 # External
 arrow2 = { workspace = true, features = ["io_ipc"] }

examples/python/remote/metadata.py

@@ -0,0 +1,41 @@
+"""Script to show how to interact with a remote storage node via python APIs."""
+
+from __future__ import annotations
+
+import argparse
+
+import polars as pl
+import pyarrow as pa
+import rerun as rr
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    subparsers = parser.add_subparsers(dest="subcommand")
+
+    print_cmd = subparsers.add_parser("print", help="Print everything")
+    update_cmd = subparsers.add_parser("update", help="Update metadata for a recording")
+
+    update_cmd.add_argument("id", help="ID of the recording to update")
+    update_cmd.add_argument("key", help="Key of the metadata to update")
+    update_cmd.add_argument("value", help="Value of the metadata to update")
+
+    args = parser.parse_args()
+
+    # Register the new rrd
+    conn = rr.remote.connect("http://0.0.0.0:51234")
+
+    catalog = pl.from_arrow(conn.list_recordings())
+
+    if args.subcommand == "print":
+        print(catalog)
+
+    if args.subcommand == "update":
+        id = catalog.filter(catalog["id"].str.starts_with(args.id)).select(pl.first("id")).item()
+
+        if id is None:
+            print("ID not found")
+            exit(1)
+        print(f"Updating metadata for {id}")
+
+        conn.update_metadata(id, {args.key: pa.array([args.value])})

pixi.toml

@@ -605,6 +605,7 @@ python = "=3.11"
 [feature.examples-common.pypi-dependencies]
 # External deps
 jupyter = ">=1.0"
+polars = ">=0.12.0"
 
 segment-anything = { git = "https://github.com/facebookresearch/segment-anything.git" }
 mesh-to-sdf = { git = "https://github.com/marian42/mesh_to_sdf.git" }

rerun_py/Cargo.toml

@@ -15,7 +15,7 @@ name = "rerun_bindings" # name of the .so library that the Python module will im
 
 
 [features]
-default = ["extension-module"]
+default = ["extension-module", "remote"]
 
 ## Extra features that aren't ready to be included in release builds yet.
 extra = ["pypi", "remote"]
@@ -40,6 +40,7 @@ remote = [
   "dep:re_protos",
   "dep:re_ws_comms",
   "dep:tokio",
+  "dep:tokio-stream",
   "dep:tonic",
   "dep:url",
 ]
@@ -87,6 +88,7 @@ uuid.workspace = true
 object_store = { workspace = true, optional = true, features = ["aws"] }
 re_protos = { workspace = true, optional = true }
 tokio = { workspace = true, optional = true }
+tokio-stream = { workspace = true, optional = true }
 # Not used yet, but we will need it when we start streaming data
 #tokio-stream = { workspace = true, optional = true }
 tonic = { workspace = true, default-features = false, features = [

rerun_py/src/dataframe.rs

@@ -565,8 +565,8 @@ impl PySchema {
 /// to retrieve the data.
 #[pyclass(name = "Recording")]
 pub struct PyRecording {
-    store: ChunkStoreHandle,
-    cache: re_dataframe::QueryCacheHandle,
+    pub(crate) store: ChunkStoreHandle,
+    pub(crate) cache: re_dataframe::QueryCacheHandle,
 }
 
 /// A view of a recording restricted to a given index, containing a specific set of entities and components.

rerun_py/src/remote.rs

@@ -1,12 +1,26 @@
 #![allow(unsafe_op_in_unsafe_fn)]
-use arrow::{array::ArrayData, pyarrow::PyArrowType};
+use arrow::{
+    array::{ArrayData, RecordBatch, RecordBatchIterator, RecordBatchReader},
+    datatypes::Schema,
+    pyarrow::PyArrowType,
+};
 // False positive due to #[pyfunction] macro
 use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyDict, Bound, PyResult};
-use re_chunk::TransportChunk;
-use re_protos::v0::{
-    storage_node_client::StorageNodeClient, EncoderVersion, ListRecordingsRequest,
-    RecordingMetadata, RecordingType, RegisterRecordingRequest,
+use re_chunk::{Chunk, TransportChunk};
+use re_chunk_store::ChunkStore;
+use re_dataframe::ChunkStoreHandle;
+use re_log_types::{StoreInfo, StoreSource};
+use re_protos::{
+    codec::decode,
+    v0::{
+        storage_node_client::StorageNodeClient, EncoderVersion, FetchRecordingRequest,
+        ListRecordingsRequest, RecordingId, RecordingMetadata, RecordingType,
+        RegisterRecordingRequest, UpdateRecordingMetadataRequest,
+    },
 };
+use re_sdk::{ApplicationId, StoreId, StoreKind, Time};
+
+use crate::dataframe::PyRecording;
 
 /// Register the `rerun.remote` module.
 pub(crate) fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
@@ -52,8 +66,9 @@ pub struct PyConnection {
 #[pymethods]
 impl PyConnection {
     /// List all recordings registered with the node.
-    fn list_recordings(&mut self) -> PyResult<Vec<PyRecordingMetadata>> {
-        self.runtime.block_on(async {
+    fn list_recordings(&mut self) -> PyResult<PyArrowType<Box<dyn RecordBatchReader + Send>>> {
+        let reader = self.runtime.block_on(async {
+            // TODO(jleibs): Support column projection
             let request = ListRecordingsRequest {
                 column_projection: None,
             };
@@ -64,13 +79,33 @@ impl PyConnection {
                 .await
                 .map_err(|err| PyRuntimeError::new_err(err.to_string()))?;
 
-            Ok(resp
+            let transport_chunks = resp
                 .into_inner()
                 .recordings
                 .into_iter()
-                .map(|recording| PyRecordingMetadata { info: recording })
-                .collect())
-        })
+                .map(|recording| recording.data())
+                .collect::<Result<Vec<_>, _>>()
+                .map_err(|err| PyRuntimeError::new_err(err.to_string()))?;
+
+            let record_batches: Vec<Result<RecordBatch, arrow::error::ArrowError>> =
+                transport_chunks
+                    .into_iter()
+                    .map(|tc| tc.try_to_arrow_record_batch())
+                    .collect();
+
+            // TODO(jleibs): surfacing this schema is awkward. This should be more explicit in
+            // the gRPC APIs somehow.
+            let schema = record_batches
+                .first()
+                .and_then(|batch| batch.as_ref().ok().map(|batch| batch.schema()))
+                .unwrap_or(std::sync::Arc::new(Schema::empty()));
+
+            let reader = RecordBatchIterator::new(record_batches, schema);
+
+            Ok::<_, PyErr>(reader)
+        })?;
+
+        Ok(PyArrowType(Box::new(reader)))
     }
 
     /// Register a recording along with some metadata
@@ -146,6 +181,119 @@ impl PyConnection {
             Ok(recording_id)
         })
     }
+
+    /// Updates the metadata for a recording.
+    #[pyo3(signature = (
+        id,
+        metadata
+    ))]
+    fn update_metadata(&mut self, id: &str, metadata: &Bound<'_, PyDict>) -> PyResult<()> {
+        self.runtime.block_on(async {
+            let (schema, data): (
+                Vec<arrow2::datatypes::Field>,
+                Vec<Box<dyn arrow2::array::Array>>,
+            ) = metadata
+                .iter()
+                .map(|(key, value)| {
+                    let key = key.to_string();
+                    let value = value.extract::<MetadataLike>()?;
+                    let value_array = value.to_arrow2()?;
+                    let field =
+                        arrow2::datatypes::Field::new(key, value_array.data_type().clone(), true);
+                    Ok((field, value_array))
+                })
+                .collect::<PyResult<Vec<_>>>()?
+                .into_iter()
+                .unzip();
+
+            let schema = arrow2::datatypes::Schema::from(schema);
+
+            let data = arrow2::chunk::Chunk::new(data);
+
+            let metadata_tc = TransportChunk {
+                schema: schema.clone(),
+                data,
+            };
+
+            let metadata = RecordingMetadata::try_from(EncoderVersion::V0, &metadata_tc)
+                .map_err(|err| PyRuntimeError::new_err(err.to_string()))?;
+
+            let request = UpdateRecordingMetadataRequest {
+                recording_id: Some(RecordingId { id: id.to_owned() }),
+                metadata: Some(metadata),
+            };
+
+            self.client
+                .update_recording_metadata(request)
+                .await
+                .map_err(|err| PyRuntimeError::new_err(err.to_string()))?;
+
+            Ok(())
+        })
+    }
+
+    /// Opens a recording for query and analysis.
+    #[pyo3(signature = (
+        id,
+    ))]
+    fn open_recording(&mut self, id: &str) -> PyResult<PyRecording> {
+        use tokio_stream::StreamExt as _;
+        let store = self.runtime.block_on(async {
+            let mut resp = self
+                .client
+                .fetch_recording(FetchRecordingRequest {
+                    recording_id: Some(RecordingId { id: id.to_owned() }),
+                })
+                .await
+                .map_err(|err| PyRuntimeError::new_err(err.to_string()))?
+                .into_inner();
+
+            // TODO(jleibs): Does this come from RDP?
+            let store_id = StoreId::from_string(StoreKind::Recording, id.to_owned());
+
+            let store_info = StoreInfo {
+                application_id: ApplicationId::from("rerun_data_platform"),
+                store_id: store_id.clone(),
+                cloned_from: None,
+                is_official_example: false,
+                started: Time::now(),
+                store_source: StoreSource::Unknown,
+                store_version: None,
+            };
+
+            let mut store = ChunkStore::new(store_id, Default::default());
+            store.set_info(store_info);
+
+            while let Some(result) = resp.next().await {
+                let response = result.map_err(|err| PyRuntimeError::new_err(err.to_string()))?;
+                let tc = decode(EncoderVersion::V0, &response.payload)
+                    .map_err(|err| PyRuntimeError::new_err(err.to_string()))?;
+
+                let Some(tc) = tc else {
+                    return Err(PyRuntimeError::new_err("Stream error"));
+                };
+
+                let chunk = Chunk::from_transport(&tc)
+                    .map_err(|err| PyRuntimeError::new_err(err.to_string()))?;
+
+                store
+                    .insert_chunk(&std::sync::Arc::new(chunk))
+                    .map_err(|err| PyRuntimeError::new_err(err.to_string()))?;
+            }
+
+            Ok(store)
+        })?;
+
+        let handle = ChunkStoreHandle::new(store);
+
+        let cache =
+            re_dataframe::QueryCacheHandle::new(re_dataframe::QueryCache::new(handle.clone()));
+
+        Ok(PyRecording {
+            store: handle,
+            cache,
+        })
+    }
 }
 
 /// A type alias for metadata.