WIP: Avoid TaskContext::default, which is necessary as metadata cache factory is not used

Author: srh

rust/cubestore/cubestore/src/config/mod.rs

@@ -2103,6 +2103,10 @@ impl Config {
                     i.get_service_typed().await,
                     i.get_service_typed().await,
                     i.get_service_typed().await,
+                    i.get_service_typed::<dyn CubestoreMetadataCacheFactory>()
+                        .await
+                        .cache_factory()
+                        .clone()
                 )
             })
             .await;

rust/cubestore/cubestore/src/queryplanner/mod.rs

@@ -217,7 +217,7 @@ impl QueryPlanner for QueryPlannerImpl {
         let physical_plan = plan_ctx.state().create_physical_plan(&plan_to_move).await?;
 
         let execution_time = SystemTime::now();
-        let results = collect(physical_plan, Arc::new(TaskContext::default())).await?;
+        let results = collect(physical_plan, ctx.task_ctx()).await?;
         let execution_time = execution_time.elapsed()?;
         app_metrics::META_QUERY_TIME_MS.report(execution_time.as_millis() as i64);
         debug!("Meta query data processing time: {:?}", execution_time,);

rust/cubestore/cubestore/src/store/compaction.rs

@@ -12,6 +12,7 @@ use crate::metastore::{
 use crate::queryplanner::merge_sort::LastRowByUniqueKeyExec;
 use crate::queryplanner::metadata_cache::MetadataCacheFactory;
 use crate::queryplanner::trace_data_loaded::{DataLoadedSize, TraceDataLoadedExec};
+use crate::queryplanner::QueryPlannerImpl;
 use crate::remotefs::{ensure_temp_file_is_dropped, RemoteFs};
 use crate::store::{min_max_values_from_data, ChunkDataStore, ChunkStore, ROW_GROUP_SIZE};
 use crate::table::data::{cmp_min_rows, cmp_partition_key};
@@ -190,11 +191,14 @@ impl CompactionServiceImpl {
         let deactivate_res = self
             .deactivate_and_mark_failed_chunks_for_replay(failed)
             .await;
+
+        let task_context = QueryPlannerImpl::execution_context_helper(self.metadata_cache_factory.cache_factory().make_session_config()).task_ctx();
+
         let in_memory_res = self
-            .compact_chunks_to_memory(mem_chunks, &partition, &index, &table)
+            .compact_chunks_to_memory(mem_chunks, &partition, &index, &table, task_context.clone())
             .await;
         let persistent_res = self
-            .compact_chunks_to_persistent(persistent_chunks, &partition, &index, &table)
+            .compact_chunks_to_persistent(persistent_chunks, &partition, &index, &table, task_context)
             .await;
         deactivate_res?;
         in_memory_res?;
@@ -209,6 +213,7 @@ impl CompactionServiceImpl {
         partition: &IdRow<Partition>,
         index: &IdRow<Index>,
         table: &IdRow<Table>,
+        task_context: Arc<TaskContext>,
     ) -> Result<(), CubeError> {
         if chunks.is_empty() {
             return Ok(());
@@ -290,6 +295,7 @@ impl CompactionServiceImpl {
                 in_memory_columns,
                 unique_key.clone(),
                 aggregate_columns.clone(),
+                task_context.clone(),
             )
             .await?;
             let batches = collect(batches_stream).await?;
@@ -337,6 +343,7 @@ impl CompactionServiceImpl {
         partition: &IdRow<Partition>,
         index: &IdRow<Index>,
         table: &IdRow<Table>,
+        task_context: Arc<TaskContext>,
     ) -> Result<(), CubeError> {
         if chunks.is_empty() {
             return Ok(());
@@ -381,6 +388,7 @@ impl CompactionServiceImpl {
             in_memory_columns,
             unique_key.clone(),
             aggregate_columns.clone(),
+            task_context,
         )
         .await?;
 
@@ -687,8 +695,9 @@ impl CompactionService for CompactionServiceImpl {
             IndexType::Regular => None,
             IndexType::Aggregate => Some(table.get_row().aggregate_columns()),
         };
+        let task_context = QueryPlannerImpl::execution_context_helper(self.metadata_cache_factory.cache_factory().make_session_config()).task_ctx();
         let records =
-            merge_chunks(key_size, main_table, new, unique_key, aggregate_columns).await?;
+            merge_chunks(key_size, main_table, new, unique_key, aggregate_columns, task_context).await?;
         let count_and_min = write_to_files(
             records,
             total_rows as usize,
@@ -890,6 +899,7 @@ impl CompactionService for CompactionServiceImpl {
             key_len,
             // TODO should it respect table partition_split_threshold?
             self.config.partition_split_threshold() as usize,
+            QueryPlannerImpl::execution_context_helper(self.metadata_cache_factory.cache_factory().make_session_config()).task_ctx(),
         )
         .await?;
         // There is no point if we cannot split the partition.
@@ -988,8 +998,9 @@ async fn find_partition_keys(
     p: AggregateExec,
     key_len: usize,
     rows_per_partition: usize,
+    context: Arc<TaskContext>,
 ) -> Result<Vec<Row>, CubeError> {
-    let mut s = p.execute(0, Arc::new(TaskContext::default()))?;
+    let mut s = p.execute(0, context)?;
     let mut points = Vec::new();
     let mut row_count = 0;
     while let Some(b) = s.next().await.transpose()? {
@@ -1364,6 +1375,7 @@ pub async fn merge_chunks(
     r: Vec<ArrayRef>,
     unique_key_columns: Option<Vec<&crate::metastore::Column>>,
     aggregate_columns: Option<Vec<AggregateColumn>>,
+    task_context: Arc<TaskContext>,
 ) -> Result<SendableRecordBatchStream, CubeError> {
     let schema = l.schema();
     let r = RecordBatch::try_new(schema.clone(), r)?;
@@ -1421,7 +1433,7 @@ pub async fn merge_chunks(
         )?);
     }
 
-    Ok(res.execute(0, Arc::new(TaskContext::default()))?)
+    Ok(res.execute(0, task_context)?)
 }
 
 pub async fn merge_replay_handles(
@@ -2331,6 +2343,7 @@ impl MultiSplit {
             ROW_GROUP_SIZE,
             self.metadata_cache_factory.clone(),
         );
+        let task_context = QueryPlannerImpl::execution_context_helper(self.metadata_cache_factory.cache_factory().make_session_config()).task_ctx();
         let records = if !in_files.is_empty() {
             read_files(
                 &in_files.into_iter().map(|(f, _)| f).collect::<Vec<_>>(),
@@ -2340,10 +2353,10 @@ impl MultiSplit {
                 Arc::new(store.arrow_schema()),
             )
             .await?
-            .execute(0, Arc::new(TaskContext::default()))?
+            .execute(0, task_context)?
         } else {
             EmptyExec::new(Arc::new(store.arrow_schema()))
-                .execute(0, Arc::new(TaskContext::default()))?
+                .execute(0, task_context)?
         };
         let row_counts = write_to_files_by_keys(
             records,

rust/cubestore/cubestore/src/store/mod.rs

@@ -18,6 +18,7 @@ use crate::metastore::{
     deactivate_table_due_to_corrupt_data, deactivate_table_on_corrupt_data, table::Table, Chunk,
     Column, ColumnType, IdRow, Index, IndexType, MetaStore, Partition, WAL,
 };
+use crate::queryplanner::QueryPlannerImpl;
 use crate::remotefs::{ensure_temp_file_is_dropped, RemoteFs};
 use crate::table::{Row, TableValue};
 use crate::util::batch_memory::columns_vec_buffer_size;
@@ -432,12 +433,15 @@ impl ChunkDataStore for ChunkStore {
         if old_chunk_ids.is_empty() {
             return Ok(());
         }
+        let task_context = QueryPlannerImpl::execution_context_helper(self.metadata_cache_factory.cache_factory().make_session_config()).task_ctx();
+
         let batches_stream = merge_chunks(
             key_size,
             main_table.clone(),
             in_memory_columns,
             unique_key.clone(),
             aggregate_columns.clone(),
+            task_context,
         )
         .await?;
         let batches = common_collect(batches_stream).await?;
@@ -1342,7 +1346,9 @@ impl ChunkStore {
 
                 assert!(aggregate.properties().output_ordering().is_some_and(|ordering| ordering.len() == key_size));
 
-                let batches = collect(aggregate, Arc::new(TaskContext::default())).await?;
+                let task_context = QueryPlannerImpl::execution_context_helper(self.metadata_cache_factory.cache_factory().make_session_config()).task_ctx();
+
+                let batches = collect(aggregate, task_context).await?;
                 if batches.is_empty() {
                     Ok(vec![])
                 } else if batches.len() == 1 {

rust/cubestore/cubestore/src/streaming/kafka.rs

@@ -2,6 +2,7 @@ use crate::config::injection::DIService;
 use crate::config::ConfigObj;
 use crate::metastore::table::StreamOffset;
 use crate::metastore::Column;
+use crate::queryplanner::metadata_cache::MetadataCacheFactory;
 use crate::streaming::kafka_post_processing::{KafkaPostProcessPlan, KafkaPostProcessPlanner};
 use crate::streaming::traffic_sender::TrafficSender;
 use crate::streaming::{parse_json_payload_and_key, StreamingSource};
@@ -59,6 +60,7 @@ impl KafkaStreamingSource {
         kafka_client: Arc<dyn KafkaClientService>,
         use_ssl: bool,
         trace_obj: Option<String>,
+        metadata_cache_factory: Arc<dyn MetadataCacheFactory>,
     ) -> Result<Self, CubeError> {
         let (post_processing_plan, columns, unique_key_columns, seq_column_index) =
             if let Some(select_statement) = select_statement {
@@ -70,7 +72,7 @@ impl KafkaStreamingSource {
                     source_columns,
                 );
                 let plan = planner
-                    .build(select_statement.clone())
+                    .build(select_statement.clone(), metadata_cache_factory)
                     .await?;
                 let columns = plan.source_columns().clone();
                 let seq_column_index = plan.source_seq_column_index();
@@ -448,7 +450,7 @@ mod tests {
             .await
             .unwrap();
 
-        let batches = collect(phys_plan, Arc::new(TaskContext::default()))
+        let batches = collect(phys_plan, plan_ctx.task_ctx())
             .await
             .unwrap();
         let res = batches_to_dataframe(batches).unwrap();
@@ -487,7 +489,7 @@ mod tests {
             .unwrap();
         let phys_plan = phys_plan.with_new_children(vec![inp]).unwrap();
 
-        let batches = collect(phys_plan, Arc::new(TaskContext::default()))
+        let batches = collect(phys_plan, plan_ctx.task_ctx())
             .await
             .unwrap();
         let res = batches_to_dataframe(batches).unwrap();

rust/cubestore/cubestore/src/streaming/kafka_post_processing.rs

@@ -1,4 +1,5 @@
 use crate::metastore::Column;
+use crate::queryplanner::metadata_cache::MetadataCacheFactory;
 use crate::queryplanner::{QueryPlan, QueryPlannerImpl};
 use crate::sql::MySqlDialectWithBackTicks;
 use crate::streaming::topic_table_provider::TopicTableProvider;
@@ -29,6 +30,7 @@ use std::sync::Arc;
 
 #[derive(Clone)]
 pub struct KafkaPostProcessPlan {
+    metadata_cache_factory: Arc<dyn MetadataCacheFactory>,
     projection_plan: Arc<dyn ExecutionPlan>,
     filter_plan: Option<Arc<dyn ExecutionPlan>>,
     source_columns: Vec<Column>,
@@ -44,6 +46,7 @@ impl KafkaPostProcessPlan {
         source_columns: Vec<Column>,
         source_unique_columns: Vec<Column>,
         source_seq_column_index: usize,
+        metadata_cache_factory: Arc<dyn MetadataCacheFactory>,
     ) -> Self {
         let source_schema = Arc::new(Schema::new(
             source_columns
@@ -58,6 +61,7 @@ impl KafkaPostProcessPlan {
             source_unique_columns,
             source_seq_column_index,
             source_schema,
+            metadata_cache_factory,
         }
     }
 
@@ -91,7 +95,9 @@ impl KafkaPostProcessPlan {
             .clone()
             .with_new_children(vec![filter_input])?;
 
-        let mut out_batches = collect(projection, Arc::new(TaskContext::default())).await?;
+        let task_context = QueryPlannerImpl::execution_context_helper(self.metadata_cache_factory.make_session_config()).task_ctx();
+
+        let mut out_batches = collect(projection, task_context).await?;
         let res = if out_batches.len() == 1 {
             out_batches.pop().unwrap()
         } else {
@@ -139,6 +145,7 @@ impl KafkaPostProcessPlanner {
     pub async fn build(
         &self,
         select_statement: String,
+        metadata_cache_factory: Arc<dyn MetadataCacheFactory>,
     ) -> Result<KafkaPostProcessPlan, CubeError> {
         let target_schema = Arc::new(Schema::new(
             self.columns
@@ -176,6 +183,7 @@ impl KafkaPostProcessPlanner {
             self.source_columns.clone(),
             source_unique_columns,
             source_seq_column_index,
+            metadata_cache_factory,
         ))
     }
 

rust/cubestore/cubestore/src/streaming/mod.rs

@@ -11,6 +11,7 @@ use crate::metastore::replay_handle::{ReplayHandle, SeqPointer, SeqPointerForLoc
 use crate::metastore::source::SourceCredentials;
 use crate::metastore::table::{StreamOffset, Table};
 use crate::metastore::{Column, ColumnType, IdRow, MetaStore};
+use crate::queryplanner::metadata_cache::MetadataCacheFactory;
 use crate::sql::timestamp_from_string;
 use crate::store::ChunkDataStore;
 use crate::streaming::kafka::{KafkaClientService, KafkaStreamingSource};
@@ -57,6 +58,7 @@ pub struct StreamingServiceImpl {
     chunk_store: Arc<dyn ChunkDataStore>,
     ksql_client: Arc<dyn KsqlClient>,
     kafka_client: Arc<dyn KafkaClientService>,
+    metadata_cache_factory: Arc<dyn MetadataCacheFactory>,
 }
 
 crate::di_service!(StreamingServiceImpl, [StreamingService]);
@@ -68,13 +70,15 @@ impl StreamingServiceImpl {
         chunk_store: Arc<dyn ChunkDataStore>,
         ksql_client: Arc<dyn KsqlClient>,
         kafka_client: Arc<dyn KafkaClientService>,
+        metadata_cache_factory: Arc<dyn MetadataCacheFactory>,
     ) -> Arc<Self> {
         Arc::new(Self {
             config_obj,
             meta_store,
             chunk_store,
             ksql_client,
             kafka_client,
+            metadata_cache_factory,
         })
     }
 
@@ -165,6 +169,7 @@ impl StreamingServiceImpl {
                 self.kafka_client.clone(),
                 *use_ssl,
                 trace_obj,
+                self.metadata_cache_factory.clone(),
             ).await?)),
         }
     }

rust/cubestore/cubestore/src/table/data.rs

@@ -241,6 +241,7 @@ pub fn rows_to_columns(cols: &[Column], rows: &[Row]) -> Vec<ArrayRef> {
 
 pub fn to_stream(r: RecordBatch) -> SendableRecordBatchStream {
     let schema = r.schema();
+    // TaskContext::default is OK here because it's a plain memory exec.
     MemoryExec::try_new(&[vec![r]], schema, None)
         .unwrap()
         .execute(0, Arc::new(TaskContext::default()))