Implement Dataset.download_partition() to download a partition as a Recording (#9451)

Title. This enable to tie the new Python API to the existing dataframe
API:

<img width="1010" alt="image"
src="https://github.com/user-attachments/assets/fb623700-f5f3-431c-8ab0-9d17f35bb5d6"
/>

Author: abey79

crates/store/re_grpc_client/src/lib.rs

@@ -70,6 +70,9 @@ pub enum StreamError {
 
     #[error(transparent)]
     TypeConversionError(#[from] re_protos::TypeConversionError),
+
+    #[error("Chunk data missing in response")]
+    MissingChunkData,
 }
 
 impl From<tonic::Status> for StreamError {

crates/store/re_grpc_client/src/redap/mod.rs

@@ -1,4 +1,4 @@
-use tokio_stream::StreamExt as _;
+use tokio_stream::{Stream, StreamExt as _};
 
 use re_chunk::Chunk;
 use re_log_encoding::codec::wire::decoder::Decode as _;
@@ -7,7 +7,7 @@ use re_protos::catalog::v1alpha1::ext::ReadDatasetEntryResponse;
 use re_protos::catalog::v1alpha1::ReadDatasetEntryRequest;
 use re_protos::frontend::v1alpha1::frontend_service_client::FrontendServiceClient;
 use re_protos::frontend::v1alpha1::FetchPartitionRequest;
-
+use re_protos::manifest_registry::v1alpha1::FetchPartitionResponse;
 use re_uri::{DatasetDataEndpoint, Origin};
 
 use crate::{spawn_future, StreamError, MAX_DECODING_MESSAGE_SIZE};
@@ -154,6 +154,20 @@ pub async fn client_with_interceptor<I: tonic::service::Interceptor>(
     )
 }
 
+/// Converts a `FetchPartitionResponse` stream into a stream of `Chunk`s.
+//TODO(#9430): ideally this should be factored as a nice helper in `re_proto`
+pub fn fetch_partition_response_to_chunk(
+    response: tonic::Streaming<FetchPartitionResponse>,
+) -> impl Stream<Item = Result<Chunk, StreamError>> {
+    response.map(|resp| {
+        resp.map_err(Into::into).and_then(|r| {
+            let batch = r.chunk.ok_or(StreamError::MissingChunkData)?.decode()?;
+
+            Chunk::from_record_batch(&batch).map_err(Into::into)
+        })
+    })
+}
+
 pub async fn stream_partition_async(
     tx: re_smart_channel::Sender<LogMsg>,
     endpoint: re_uri::DatasetDataEndpoint,
@@ -188,15 +202,6 @@ pub async fn stream_partition_async(
         .await?
         .into_inner();
 
-    let mut chunk_stream = catalog_chunk_stream.map(|resp| {
-        resp.and_then(|r| {
-            r.chunk
-                .ok_or_else(|| tonic::Status::internal("missing chunk in FetchPartitionResponse"))?
-                .decode()
-                .map_err(|err| tonic::Status::internal(err.to_string()))
-        })
-    });
-
     drop(client);
 
     let store_id = StoreId::from_string(StoreKind::Recording, endpoint.partition_id.clone());
@@ -227,9 +232,10 @@ pub async fn stream_partition_async(
         return Ok(());
     }
 
-    while let Some(result) = chunk_stream.next().await {
-        let batch = result?;
-        let chunk = Chunk::from_record_batch(&batch)?;
+    let mut chunk_stream = fetch_partition_response_to_chunk(catalog_chunk_stream);
+
+    while let Some(chunk) = chunk_stream.next().await {
+        let chunk = chunk?;
 
         if tx
             .send(LogMsg::ArrowMsg(store_id.clone(), chunk.to_arrow_msg()?))

rerun_py/src/catalog/dataset.rs

@@ -1,8 +1,15 @@
-use pyo3::{pyclass, pymethods};
+use pyo3::{pyclass, pymethods, PyRef, PyResult};
+use tokio_stream::StreamExt as _;
 
+use re_chunk_store::{ChunkStore, ChunkStoreHandle};
+use re_grpc_client::redap::fetch_partition_response_to_chunk;
+use re_log_types::{StoreId, StoreInfo, StoreKind, StoreSource};
 use re_protos::common::v1alpha1::ext::DatasetHandle;
+use re_protos::frontend::v1alpha1::FetchPartitionRequest;
 
-use crate::catalog::PyEntry;
+use crate::catalog::{to_py_err, PyEntry};
+use crate::dataframe::PyRecording;
+use crate::utils::wait_for_future;
 
 #[pyclass(name = "Dataset", extends=PyEntry)]
 pub struct PyDataset {
@@ -15,4 +22,56 @@ impl PyDataset {
     fn manifest_url(&self) -> String {
         self.dataset_handle.url.to_string()
     }
+
+    fn download_partition(self_: PyRef<'_, Self>, partition_id: String) -> PyResult<PyRecording> {
+        let super_ = self_.as_super();
+        let mut client = super_.client.borrow(self_.py()).connection().client();
+
+        let dataset_id = super_.details.id;
+        let dataset_name = super_.details.name.clone();
+
+        let store: PyResult<ChunkStore> = wait_for_future(self_.py(), async move {
+            let catalog_chunk_stream = client
+                .fetch_partition(FetchPartitionRequest {
+                    dataset_id: Some(dataset_id.into()),
+                    partition_id: Some(partition_id.clone().into()),
+                })
+                .await
+                .map_err(to_py_err)?
+                .into_inner();
+
+            let store_id = StoreId::from_string(StoreKind::Recording, partition_id);
+            let store_info = StoreInfo {
+                application_id: dataset_name.into(),
+                store_id: store_id.clone(),
+                cloned_from: None,
+                store_source: StoreSource::Unknown,
+                store_version: None,
+            };
+
+            let mut store = ChunkStore::new(store_id, Default::default());
+            store.set_info(store_info);
+
+            let mut chunk_stream = fetch_partition_response_to_chunk(catalog_chunk_stream);
+
+            while let Some(chunk) = chunk_stream.next().await {
+                let chunk = chunk.map_err(to_py_err)?;
+                store
+                    .insert_chunk(&std::sync::Arc::new(chunk))
+                    .map_err(to_py_err)?;
+            }
+
+            Ok(store)
+        });
+
+        let handle = ChunkStoreHandle::new(store?);
+
+        let cache =
+            re_dataframe::QueryCacheHandle::new(re_dataframe::QueryCache::new(handle.clone()));
+
+        Ok(PyRecording {
+            store: handle,
+            cache,
+        })
+    }
 }

rerun_py/src/catalog/errors.rs

@@ -29,6 +29,9 @@ enum ExternalError {
     ConnectionError(ConnectionError),
     TonicStatusError(tonic::Status),
     UriError(re_uri::Error),
+    ChunkError(re_chunk::ChunkError),
+    ChunkStoreError(re_chunk_store::ChunkStoreError),
+    StreamError(re_grpc_client::StreamError),
 }
 
 impl From<ConnectionError> for ExternalError {
@@ -49,6 +52,24 @@ impl From<re_uri::Error> for ExternalError {
     }
 }
 
+impl From<re_chunk::ChunkError> for ExternalError {
+    fn from(value: re_chunk::ChunkError) -> Self {
+        Self::ChunkError(value)
+    }
+}
+
+impl From<re_chunk_store::ChunkStoreError> for ExternalError {
+    fn from(value: re_chunk_store::ChunkStoreError) -> Self {
+        Self::ChunkStoreError(value)
+    }
+}
+
+impl From<re_grpc_client::StreamError> for ExternalError {
+    fn from(value: re_grpc_client::StreamError) -> Self {
+        Self::StreamError(value)
+    }
+}
+
 impl From<ExternalError> for PyErr {
     fn from(err: ExternalError) -> Self {
         match err {
@@ -69,6 +90,16 @@ impl From<ExternalError> for PyErr {
             }
 
             ExternalError::UriError(err) => PyValueError::new_err(format!("Invalid URI: {err}")),
+
+            ExternalError::ChunkError(err) => PyValueError::new_err(format!("Chunk error: {err}")),
+
+            ExternalError::ChunkStoreError(err) => {
+                PyValueError::new_err(format!("Chunk store error: {err}"))
+            }
+
+            ExternalError::StreamError(err) => {
+                PyValueError::new_err(format!("Data streaming error: {err}"))
+            }
         }
     }
 }