Clean-up of the partition table datafusion table provider (#11434)

Author: abey79

Cargo.lock

@@ -8657,7 +8657,6 @@ name = "re_datafusion"
 version = "0.26.0-alpha.1+dev"
 dependencies = [
  "ahash",
- "anyhow",
  "arrow",
  "async-stream",
  "async-trait",
@@ -8666,7 +8665,6 @@ dependencies = [
  "futures",
  "futures-util",
  "getrandom 0.3.3",
- "itertools 0.14.0",
  "log",
  "parking_lot",
  "re_arrow_util",
@@ -8676,7 +8674,6 @@ dependencies = [
  "re_protos",
  "re_redap_client",
  "re_sorbet",
- "re_tuid",
  "re_uri",
  "tokio",
  "tokio-stream",

crates/store/re_datafusion/Cargo.toml

@@ -37,19 +37,16 @@ re_log_encoding.workspace = true
 re_log_types.workspace = true
 re_protos.workspace = true
 re_sorbet.workspace = true
-re_tuid.workspace = true
 re_uri.workspace = true
 
 # External dependencies:
 ahash.workspace = true
-anyhow.workspace = true
 arrow.workspace = true
 async-stream.workspace = true
 async-trait.workspace = true
 datafusion.workspace = true
 futures.workspace = true
 futures-util.workspace = true
-itertools.workspace = true
 log.workspace = true
 parking_lot.workspace = true
 tokio.workspace = true
@@ -67,5 +64,4 @@ getrandom = { workspace = true, features = ["wasm_js"] }
 
 
 [dev-dependencies]
-re_tuid.workspace = true
 tokio = { workspace = true, features = ["macros", "rt-multi-thread", "signal"] }

crates/store/re_datafusion/examples/catalog.rs

@@ -1,15 +1,18 @@
-use crate::TableEntryTableProvider;
+use std::any::Any;
+use std::iter;
+use std::sync::Arc;
+
 use ahash::{HashMap, HashSet};
 use async_trait::async_trait;
 use datafusion::catalog::{CatalogProvider, SchemaProvider, TableProvider};
 use datafusion::common::{DataFusionError, Result as DataFusionResult, TableReference, exec_err};
 use parking_lot::Mutex;
-use re_redap_client::ConnectionClient;
-use std::any::Any;
-use std::iter;
-use std::sync::Arc;
 use tokio::runtime::Handle as RuntimeHandle;
 
+use re_redap_client::ConnectionClient;
+
+use crate::TableEntryTableProvider;
+
 // These are to match the defaults in datafusion.
 pub const DEFAULT_CATALOG_NAME: &str = "datafusion";
 const DEFAULT_SCHEMA_NAME: &str = "public";

crates/store/re_datafusion/src/catalog_provider.rs

@@ -6,7 +6,6 @@ mod dataframe_query_common;
 mod dataframe_query_provider;
 #[cfg(target_arch = "wasm32")]
 mod dataframe_query_provider_wasm;
-mod datafusion_connector;
 mod grpc_streaming_provider;
 mod partition_table;
 mod search_provider;
@@ -19,7 +18,6 @@ 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 datafusion_connector::DataFusionConnector;
 pub use partition_table::PartitionTableProvider;
 pub use search_provider::SearchResultsTableProvider;
 pub use table_entry_provider::TableEntryTableProvider;

crates/store/re_datafusion/src/datafusion_connector.rs

@@ -10,10 +10,9 @@ use tracing::instrument;
 
 use re_log_encoding::codec::wire::decoder::Decode as _;
 use re_log_types::EntryId;
-use re_protos::cloud::v1alpha1::GetPartitionTableSchemaRequest;
-use re_protos::headers::RerunHeadersInjectorExt as _;
 use re_protos::{
-    cloud::v1alpha1::ScanPartitionTableRequest, cloud::v1alpha1::ScanPartitionTableResponse,
+    cloud::v1alpha1::{ScanPartitionTableRequest, ScanPartitionTableResponse},
+    headers::RerunHeadersInjectorExt as _,
 };
 use re_redap_client::ConnectionClient;
 
@@ -22,7 +21,6 @@ use crate::wasm_compat::make_future_send;
 
 #[derive(Clone)]
 pub struct PartitionTableProvider {
-    //TODO(#10191): this should use a `ConnectionRegistryHandle` instead
     client: ConnectionClient,
     dataset_id: EntryId,
 }
@@ -52,27 +50,27 @@ impl GrpcStreamToTable for PartitionTableProvider {
 
     #[instrument(skip(self), err)]
     async fn fetch_schema(&mut self) -> DataFusionResult<SchemaRef> {
-        let req = tonic::Request::new(GetPartitionTableSchemaRequest {})
-            .with_entry_id(self.dataset_id)
-            .map_err(|err| DataFusionError::External(Box::new(err)))?;
-
         let mut client = self.client.clone();
 
+        let dataset_id = self.dataset_id;
+
         Ok(Arc::new(
-            make_future_send(
-                async move { Ok(client.inner().get_partition_table_schema(req).await) },
-            )
-            .await?
-            .map_err(|err| DataFusionError::External(Box::new(err)))?
-            .into_inner()
-            .schema
-            .ok_or(DataFusionError::External(
-                "Schema missing from GetPartitionTableSchema response".into(),
-            ))?
-            .try_into()?,
+            make_future_send(async move {
+                client
+                    .get_partition_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,

crates/store/re_datafusion/src/lib.rs

@@ -6,11 +6,11 @@ use datafusion::{
     catalog::TableProvider,
     error::{DataFusionError, Result as DataFusionResult},
 };
-use re_log_types::EntryId;
 use tokio_stream::StreamExt as _;
 use tracing::instrument;
 
 use re_log_encoding::codec::wire::decoder::Decode as _;
+use re_log_types::EntryId;
 use re_protos::{
     cloud::v1alpha1::{SearchDatasetRequest, SearchDatasetResponse},
     common::v1alpha1::ScanParameters,
@@ -23,7 +23,6 @@ use crate::wasm_compat::make_future_send;
 
 #[derive(Clone)]
 pub struct SearchResultsTableProvider {
-    //TODO(#10191): this should use a `ConnectionRegistryHandle` instead
     client: ConnectionClient,
     dataset_id: EntryId,
     request: SearchDatasetRequest,

crates/store/re_datafusion/src/partition_table.rs

@@ -21,7 +21,6 @@ use crate::wasm_compat::make_future_send;
 
 #[derive(Clone)]
 pub struct TableEntryTableProvider {
-    //TODO(#10191): this should use a `ConnectionRegistryHandle` instead
     client: ConnectionClient,
     table: EntryIdOrName,
 
@@ -47,6 +46,10 @@ impl TableEntryTableProvider {
         }
     }
 
+    pub fn new_entry_list(client: ConnectionClient) -> Self {
+        Self::new(client, "__entries")
+    }
+
     /// This is a convenience function
     pub async fn into_provider(self) -> Result<Arc<dyn TableProvider>, DataFusionError> {
         Ok(GrpcStreamProvider::prepare(self).await?)