Partition-to-segment rename (wave 1): redap layer (#12017)

Author: abey79

crates/store/re_data_source/src/data_source.rs

@@ -255,7 +255,7 @@ impl LogDataSource {
                         .client(uri_clone.origin.clone())
                         .await
                         .map_err(|err| ApiError::connection(err, "failed to connect to server"))?;
-                    re_redap_client::stream_blueprint_and_partition_from_server(
+                    re_redap_client::stream_blueprint_and_segment_from_server(
                         client, tx, uri_clone, on_msg,
                     )
                     .await

crates/store/re_datafusion/src/dataframe_query_common.rs

@@ -22,11 +22,11 @@ use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec;
 use re_dataframe::external::re_chunk_store::ChunkStore;
 use re_dataframe::{Index, QueryExpression};
 use re_log_types::EntryId;
-use re_protos::cloud::v1alpha1::FetchChunksRequest;
 use re_protos::{
     cloud::v1alpha1::{
-        GetDatasetSchemaRequest, QueryDatasetRequest, ScanPartitionTableResponse,
-        ext::{Query, QueryLatestAt, QueryRange},
+        FetchChunksRequest, GetDatasetSchemaRequest, QueryDatasetResponse,
+        ScanSegmentTableResponse,
+        ext::{Query, QueryDatasetRequest, QueryLatestAt, QueryRange},
     },
     common::v1alpha1::ext::ScanParameters,
     headers::RerunHeadersInjectorExt as _,
@@ -109,33 +109,27 @@ impl DataframeQueryTableProvider {
         let query = query_from_query_expression(query_expression);
 
         let dataset_query = QueryDatasetRequest {
-            partition_ids: partition_ids
+            segment_ids: partition_ids
                 .iter()
                 .map(|id| id.as_ref().to_owned().into())
                 .collect(),
             chunk_ids: vec![],
-            entity_paths: entity_paths
-                .into_iter()
-                .map(|p| (*p).clone().into())
-                .collect(),
+            entity_paths: entity_paths.into_iter().map(|p| (*p).clone()).collect(),
             select_all_entity_paths,
             fuzzy_descriptors,
             exclude_static_data: false,
             exclude_temporal_data: false,
-            query: Some(query.into()),
-            scan_parameters: Some(
-                ScanParameters {
-                    columns: FetchChunksRequest::required_column_names(),
-                    ..Default::default()
-                }
-                .into(),
-            ),
+            query: Some(query),
+            scan_parameters: Some(ScanParameters {
+                columns: FetchChunksRequest::required_column_names(),
+                ..Default::default()
+            }),
         };
 
         let response_stream = client
             .inner()
             .query_dataset(
-                tonic::Request::new(dataset_query)
+                tonic::Request::new(dataset_query.into())
                     .with_entry_id(dataset_id)
                     .map_err(|err| exec_datafusion_err!("{err}"))?,
             )
@@ -161,7 +155,7 @@ impl DataframeQueryTableProvider {
 
         let schema = Arc::new(prepend_string_column_schema(
             &schema,
-            ScanPartitionTableResponse::FIELD_PARTITION_ID,
+            ScanSegmentTableResponse::FIELD_SEGMENT_ID,
         ));
 
         Ok(Self {
@@ -415,21 +409,24 @@ pub(crate) fn group_chunk_infos_by_partition_id(
 
     for batch in chunk_info_batches.as_ref() {
         let partition_ids = batch
-            .column_by_name("chunk_partition_id")
+            .column_by_name(QueryDatasetResponse::FIELD_CHUNK_SEGMENT_ID)
             .ok_or(exec_datafusion_err!(
-                "Unable to find chunk_partition_id column"
+                "Unable to find {} column",
+                QueryDatasetResponse::FIELD_CHUNK_SEGMENT_ID
             ))?
             .as_any()
             .downcast_ref::<StringArray>()
             .ok_or(exec_datafusion_err!(
-                "chunk_partition_id must be string type"
+                "{} must be string type",
+                QueryDatasetResponse::FIELD_CHUNK_SEGMENT_ID
             ))?;
 
         // group rows by partition ID
         let mut partition_rows: BTreeMap<String, Vec<usize>> = BTreeMap::new();
         for (row_idx, partition_id) in partition_ids.iter().enumerate() {
             let pid = partition_id.ok_or(exec_datafusion_err!(
-                "Found null partition_id in chunk_partition_id column at row {row_idx}"
+                "Found null segment id in {} column at row {row_idx}",
+                QueryDatasetResponse::FIELD_CHUNK_SEGMENT_ID
             ))?;
             partition_rows
                 .entry(pid.to_owned())
@@ -532,19 +529,19 @@ mod tests {
     fn test_batches_grouping() {
         let schema = Arc::new(Schema::new_with_metadata(
             vec![
-                Field::new("chunk_partition_id", DataType::Utf8, false),
-                Field::new("chunk_id", DataType::FixedSizeBinary(32), false),
+                QueryDatasetResponse::field_chunk_segment_id(),
+                QueryDatasetResponse::field_chunk_id(),
             ],
             HashMap::default(),
         ));
 
         let capacity = 4;
-        let byte_width = 32;
+        let byte_width = 16;
         let mut chunk_id_builder = FixedSizeBinaryBuilder::with_capacity(capacity, byte_width);
-        chunk_id_builder.append_value([0u8; 32]).unwrap();
-        chunk_id_builder.append_value([1u8; 32]).unwrap();
-        chunk_id_builder.append_value([2u8; 32]).unwrap();
-        chunk_id_builder.append_value([3u8; 32]).unwrap();
+        chunk_id_builder.append_value([0u8; 16]).unwrap();
+        chunk_id_builder.append_value([1u8; 16]).unwrap();
+        chunk_id_builder.append_value([2u8; 16]).unwrap();
+        chunk_id_builder.append_value([3u8; 16]).unwrap();
         let chunk_id_array = Arc::new(chunk_id_builder.finish());
 
         let batch1 = RecordBatch::try_new_with_options(
@@ -563,9 +560,9 @@ mod tests {
         .unwrap();
 
         let mut chunk_id_builder = FixedSizeBinaryBuilder::with_capacity(capacity, byte_width);
-        chunk_id_builder.append_value([4u8; 32]).unwrap();
-        chunk_id_builder.append_value([5u8; 32]).unwrap();
-        chunk_id_builder.append_value([6u8; 32]).unwrap();
+        chunk_id_builder.append_value([4u8; 16]).unwrap();
+        chunk_id_builder.append_value([5u8; 16]).unwrap();
+        chunk_id_builder.append_value([6u8; 16]).unwrap();
         let chunk_id_array = Arc::new(chunk_id_builder.finish());
 
         let batch2 = RecordBatch::try_new_with_options(
@@ -593,8 +590,8 @@ mod tests {
             .downcast_ref::<FixedSizeBinaryArray>()
             .unwrap();
         assert_eq!(chunk_ids_a.len(), 2);
-        assert_eq!(chunk_ids_a.value(0), [0u8; 32]);
-        assert_eq!(chunk_ids_a.value(1), [2u8; 32]);
+        assert_eq!(chunk_ids_a.value(0), [0u8; 16]);
+        assert_eq!(chunk_ids_a.value(1), [2u8; 16]);
 
         let group_b = grouped.get("B").unwrap();
         assert_eq!(group_b.len(), 2);
@@ -605,15 +602,15 @@ mod tests {
             .downcast_ref::<FixedSizeBinaryArray>()
             .unwrap();
         assert_eq!(chunk_ids_b1.len(), 1);
-        assert_eq!(chunk_ids_b1.value(0), [1u8; 32]);
+        assert_eq!(chunk_ids_b1.value(0), [1u8; 16]);
         let chunk_ids_b2 = group_b[1]
             .column_by_name("chunk_id")
             .unwrap()
             .as_any()
             .downcast_ref::<FixedSizeBinaryArray>()
             .unwrap();
         assert_eq!(chunk_ids_b2.len(), 1);
-        assert_eq!(chunk_ids_b2.value(0), [4u8; 32]);
+        assert_eq!(chunk_ids_b2.value(0), [4u8; 16]);
 
         let group_c = grouped.get("C").unwrap();
         assert_eq!(group_c.len(), 2);
@@ -624,15 +621,15 @@ mod tests {
             .downcast_ref::<FixedSizeBinaryArray>()
             .unwrap();
         assert_eq!(chunk_ids_c1.len(), 1);
-        assert_eq!(chunk_ids_c1.value(0), [3u8; 32]);
+        assert_eq!(chunk_ids_c1.value(0), [3u8; 16]);
         let chunk_ids_c2 = group_c[1]
             .column_by_name("chunk_id")
             .unwrap()
             .as_any()
             .downcast_ref::<FixedSizeBinaryArray>()
             .unwrap();
         assert_eq!(chunk_ids_c2.len(), 1);
-        assert_eq!(chunk_ids_c2.value(0), [5u8; 32]);
+        assert_eq!(chunk_ids_c2.value(0), [5u8; 16]);
 
         let group_d = grouped.get("D").unwrap();
         assert_eq!(group_d.len(), 1);
@@ -643,6 +640,6 @@ mod tests {
             .downcast_ref::<FixedSizeBinaryArray>()
             .unwrap();
         assert_eq!(chunk_ids_d.len(), 1);
-        assert_eq!(chunk_ids_d.value(0), [6u8; 32]);
+        assert_eq!(chunk_ids_d.value(0), [6u8; 16]);
     }
 }

crates/store/re_datafusion/src/dataframe_query_provider.rs

@@ -32,7 +32,7 @@ use re_dataframe::{
     ChunkStoreHandle, Index, QueryCache, QueryEngine, QueryExpression, QueryHandle, StorageEngine,
 };
 use re_log_types::{ApplicationId, StoreId, StoreKind};
-use re_protos::cloud::v1alpha1::{FetchChunksRequest, ScanPartitionTableResponse};
+use re_protos::cloud::v1alpha1::{FetchChunksRequest, ScanSegmentTableResponse};
 use re_redap_client::ConnectionClient;
 use re_sorbet::{ColumnDescriptor, ColumnSelector};
 
@@ -199,12 +199,10 @@ impl PartitionStreamExec {
         let orderings = if projected_schema
             .fields()
             .iter()
-            .any(|f| f.name().as_str() == ScanPartitionTableResponse::FIELD_PARTITION_ID)
+            .any(|f| f.name().as_str() == ScanSegmentTableResponse::FIELD_SEGMENT_ID)
         {
-            let partition_col = Arc::new(Column::new(
-                ScanPartitionTableResponse::FIELD_PARTITION_ID,
-                0,
-            )) as Arc<dyn PhysicalExpr>;
+            let partition_col = Arc::new(Column::new(ScanSegmentTableResponse::FIELD_SEGMENT_ID, 0))
+                as Arc<dyn PhysicalExpr>;
             let order_col = sort_index
                 .and_then(|index| {
                     let index_name = index.as_str();
@@ -243,7 +241,7 @@ impl PartitionStreamExec {
         let output_partitioning = if partition_in_output_schema {
             Partitioning::Hash(
                 vec![Arc::new(Column::new(
-                    ScanPartitionTableResponse::FIELD_PARTITION_ID,
+                    ScanSegmentTableResponse::FIELD_SEGMENT_ID,
                     0,
                 ))],
                 num_partitions,
@@ -306,7 +304,7 @@ async fn send_next_row(
 
     let batch_schema = Arc::new(prepend_string_column_schema(
         &query_schema,
-        ScanPartitionTableResponse::FIELD_PARTITION_ID,
+        ScanSegmentTableResponse::FIELD_SEGMENT_ID,
     ));
 
     let batch = RecordBatch::try_new_with_options(
@@ -436,7 +434,7 @@ async fn chunk_stream_io_loop(
 
         // Then we need to fully decode these chunks, i.e. both the transport layer (Protobuf)
         // and the app layer (Arrow).
-        let mut chunk_stream = re_redap_client::fetch_chunks_response_to_chunk_and_partition_id(
+        let mut chunk_stream = re_redap_client::fetch_chunks_response_to_chunk_and_segment_id(
             fetch_chunks_response_stream,
         );
 

crates/store/re_datafusion/src/dataframe_query_provider_wasm.rs

@@ -27,7 +27,7 @@ use re_dataframe::{
     ChunkStoreHandle, Index, QueryCache, QueryEngine, QueryExpression, QueryHandle, StorageEngine,
 };
 use re_log_types::{StoreId, StoreKind};
-use re_protos::cloud::v1alpha1::{FetchChunksRequest, ScanPartitionTableResponse};
+use re_protos::cloud::v1alpha1::{FetchChunksRequest, ScanSegmentTableResponse};
 use re_redap_client::ConnectionClient;
 
 use crate::dataframe_query_common::{
@@ -78,7 +78,7 @@ impl DataframePartitionStream {
 
         // Then we need to fully decode these chunks, i.e. both the transport layer (Protobuf)
         // and the app layer (Arrow).
-        let mut chunk_stream = re_redap_client::fetch_chunks_response_to_chunk_and_partition_id(
+        let mut chunk_stream = re_redap_client::fetch_chunks_response_to_chunk_and_segment_id(
             fetch_chunks_response_stream,
         );
 
@@ -191,10 +191,8 @@ impl PartitionStreamExec {
             None => Arc::clone(table_schema),
         };
 
-        let partition_col = Arc::new(Column::new(
-            ScanPartitionTableResponse::FIELD_PARTITION_ID,
-            0,
-        )) as Arc<dyn PhysicalExpr>;
+        let partition_col = Arc::new(Column::new(ScanSegmentTableResponse::FIELD_SEGMENT_ID, 0))
+            as Arc<dyn PhysicalExpr>;
         let order_col = sort_index
             .and_then(|index| {
                 let index_name = index.as_str();
@@ -231,7 +229,7 @@ impl PartitionStreamExec {
         let output_partitioning = if partition_in_output_schema {
             Partitioning::Hash(
                 vec![Arc::new(Column::new(
-                    ScanPartitionTableResponse::FIELD_PARTITION_ID,
+                    ScanSegmentTableResponse::FIELD_SEGMENT_ID,
                     0,
                 ))],
                 num_partitions,
@@ -292,7 +290,7 @@ fn create_next_row(
 
     let batch_schema = Arc::new(prepend_string_column_schema(
         &query_schema,
-        ScanPartitionTableResponse::FIELD_PARTITION_ID,
+        ScanSegmentTableResponse::FIELD_SEGMENT_ID,
     ));
 
     let batch = RecordBatch::try_new_with_options(

crates/store/re_datafusion/src/dataset_manifest.rs

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

crates/store/re_datafusion/src/lib.rs

@@ -20,6 +20,6 @@ pub(crate) use dataframe_query_provider::PartitionStreamExec;
 #[cfg(target_arch = "wasm32")]
 pub(crate) use dataframe_query_provider_wasm::PartitionStreamExec;
 pub use dataset_manifest::DatasetManifestProvider;
-pub use partition_table::PartitionTableProvider;
+pub use partition_table::SegmentTableProvider;
 pub use search_provider::SearchResultsTableProvider;
 pub use table_entry_provider::TableEntryTableProvider;