add table provider for layer table

Author: abey79

crates/store/re_datafusion/src/layer_table.rs

@@ -0,0 +1,104 @@
+use std::sync::Arc;
+
+use arrow::{array::RecordBatch, datatypes::SchemaRef};
+use async_trait::async_trait;
+use datafusion::{
+    catalog::TableProvider,
+    error::{DataFusionError, Result as DataFusionResult},
+};
+use tracing::instrument;
+
+use re_log_encoding::codec::wire::decoder::Decode as _;
+use re_log_types::EntryId;
+use re_protos::{
+    cloud::v1alpha1::{ScanLayerTableRequest, ScanLayerTableResponse},
+    headers::RerunHeadersInjectorExt as _,
+};
+use re_redap_client::ConnectionClient;
+
+use crate::grpc_streaming_provider::{GrpcStreamProvider, GrpcStreamToTable};
+use crate::wasm_compat::make_future_send;
+
+//TODO(ab): deduplicate from PartitionTableProvider
+#[derive(Clone)]
+pub struct LayerTableProvider {
+    client: ConnectionClient,
+    dataset_id: EntryId,
+}
+
+impl std::fmt::Debug for LayerTableProvider {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("LayerTableProvider")
+            .field("dataset_id", &self.dataset_id)
+            .finish()
+    }
+}
+
+impl LayerTableProvider {
+    pub fn new(client: ConnectionClient, dataset_id: EntryId) -> Self {
+        Self { client, dataset_id }
+    }
+
+    /// This is a convenience function
+    pub async fn into_provider(self) -> DataFusionResult<Arc<dyn TableProvider>> {
+        Ok(GrpcStreamProvider::prepare(self).await?)
+    }
+}
+
+#[async_trait]
+impl GrpcStreamToTable for LayerTableProvider {
+    type GrpcStreamData = ScanLayerTableResponse;
+
+    #[instrument(skip(self), err)]
+    async fn fetch_schema(&mut self) -> DataFusionResult<SchemaRef> {
+        let mut client = self.client.clone();
+
+        let dataset_id = self.dataset_id;
+
+        Ok(Arc::new(
+            make_future_send(async move {
+                client
+                    .get_layer_table_schema(dataset_id)
+                    .await
+                    .map_err(|err| {
+                        DataFusionError::External(
+                            format!("Couldn't get partition table schema: {err}").into(),
+                        )
+                    })
+            })
+            .await?,
+        ))
+    }
+
+    // TODO(ab): what `GrpcStreamToTable` attempts to simplify should probably be handled by
+    // `ConnectionClient`
+    #[instrument(skip(self), err)]
+    async fn send_streaming_request(
+        &mut self,
+    ) -> DataFusionResult<tonic::Response<tonic::Streaming<Self::GrpcStreamData>>> {
+        let request = tonic::Request::new(ScanLayerTableRequest {
+            columns: vec![], // all of them
+        })
+        .with_entry_id(self.dataset_id)
+        .map_err(|err| DataFusionError::External(Box::new(err)))?;
+
+        let mut client = self.client.clone();
+
+        make_future_send(async move { Ok(client.inner().scan_layer_table(request).await) })
+            .await?
+            .map_err(|err| DataFusionError::External(Box::new(err)))
+    }
+
+    fn process_response(
+        &mut self,
+        response: Self::GrpcStreamData,
+    ) -> DataFusionResult<RecordBatch> {
+        response
+            .data
+            .ok_or(DataFusionError::Execution(
+                "DataFrame missing from PartitionList response".to_owned(),
+            ))?
+            .decode()
+            .map_err(|err| DataFusionError::External(Box::new(err)))
+    }
+}

crates/store/re_datafusion/src/lib.rs

@@ -7,6 +7,7 @@ mod dataframe_query_provider;
 #[cfg(target_arch = "wasm32")]
 mod dataframe_query_provider_wasm;
 mod grpc_streaming_provider;
+mod layer_table;
 mod partition_table;
 mod search_provider;
 mod table_entry_provider;
@@ -18,6 +19,7 @@ pub use dataframe_query_common::{DataframeQueryTableProvider, query_from_query_e
 pub(crate) use dataframe_query_provider::PartitionStreamExec;
 #[cfg(target_arch = "wasm32")]
 pub(crate) use dataframe_query_provider_wasm::PartitionStreamExec;
+pub use layer_table::LayerTableProvider;
 pub use partition_table::PartitionTableProvider;
 pub use search_provider::SearchResultsTableProvider;
 pub use table_entry_provider::TableEntryTableProvider;

crates/store/re_datafusion/src/partition_table.rs

@@ -19,6 +19,7 @@ use re_redap_client::ConnectionClient;
 use crate::grpc_streaming_provider::{GrpcStreamProvider, GrpcStreamToTable};
 use crate::wasm_compat::make_future_send;
 
+//TODO(ab): deduplicate from LayerTableProvider
 #[derive(Clone)]
 pub struct PartitionTableProvider {
     client: ConnectionClient,

crates/store/re_redap_client/src/connection_client.rs

@@ -9,9 +9,9 @@ use re_protos::{
     TypeConversionError,
     cloud::v1alpha1::{
         CreateDatasetEntryRequest, DeleteEntryRequest, EntryFilter, EntryKind, FindEntriesRequest,
-        GetPartitionTableSchemaRequest, GetPartitionTableSchemaResponse, ReadDatasetEntryRequest,
-        ReadTableEntryRequest, RegisterWithDatasetResponse, ScanPartitionTableRequest,
-        ScanPartitionTableResponse,
+        GetLayerTableSchemaRequest, GetLayerTableSchemaResponse, GetPartitionTableSchemaRequest,
+        GetPartitionTableSchemaResponse, ReadDatasetEntryRequest, ReadTableEntryRequest,
+        RegisterWithDatasetResponse, ScanPartitionTableRequest, ScanPartitionTableResponse,
         ext::{
             CreateDatasetEntryResponse, DataSource, DataSourceKind, DatasetDetails, DatasetEntry,
             EntryDetails, EntryDetailsUpdate, LanceTable, ProviderDetails as _,
@@ -269,6 +269,26 @@ where
         Ok(partition_ids)
     }
 
+    //TODO(ab): accept entry name
+    pub async fn get_layer_table_schema(
+        &mut self,
+        entry_id: EntryId,
+    ) -> Result<ArrowSchema, StreamError> {
+        Ok(self
+            .inner()
+            .get_layer_table_schema(
+                tonic::Request::new(GetLayerTableSchemaRequest {})
+                    .with_entry_id(entry_id)
+                    .map_err(|err| StreamEntryError::InvalidId(err.into()))?,
+            )
+            .await
+            .map_err(|err| StreamEntryError::GetLayerTableSchema(err.into()))?
+            .into_inner()
+            .schema
+            .ok_or_else(|| missing_field!(GetLayerTableSchemaResponse, "schema"))?
+            .try_into()?)
+    }
+
     /// Initiate registration of the provided recording URIs with a dataset and return the
     /// corresponding task descriptors.
     ///

crates/store/re_redap_client/src/lib.rs

@@ -96,6 +96,9 @@ pub enum StreamEntryError {
     #[error("Failed reading partition table scheme\nDetails:{0}")]
     GetPartitionTableSchema(TonicStatusError),
 
+    #[error("Failed reading layer table scheme\nDetails:{0}")]
+    GetLayerTableSchema(TonicStatusError),
+
     #[error("Failed scanning the partition table \nDetails:{0}")]
     ScanPartitionTable(TonicStatusError),