chore(cubestore): Upgrade DF: Construct SessionConfig, update datafusion interfaces

Author: srh

rust/cubestore/Cargo.lock

@@ -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/config/mod.rs

@@ -2,8 +2,10 @@ use bytes::Bytes;
 use datafusion::datasource::physical_plan::parquet::DefaultParquetFileReaderFactory;
 use datafusion::datasource::physical_plan::{FileMeta, ParquetFileReaderFactory};
 use datafusion::parquet::arrow::async_reader::AsyncFileReader;
+use datafusion::parquet::file::encryption::ParquetEncryptionConfig;
 use datafusion::parquet::file::metadata::ParquetMetaData;
 use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
+use datafusion::prelude::SessionConfig;
 use futures_util::future::BoxFuture;
 use futures_util::FutureExt;
 use std::fmt;
@@ -22,6 +24,9 @@ pub trait MetadataCacheFactory: Sync + Send {
         max_capacity: u64,
         time_to_idle: Duration,
     ) -> Arc<dyn ParquetFileReaderFactory>;
+    fn make_session_config(&self) -> SessionConfig {
+        SessionConfig::new()
+    }
 }
 /// Default MetadataCache, does not cache anything
 #[derive(Debug)]
@@ -132,6 +137,16 @@ pub struct LruCachingFileReader {
     cache: Arc<moka::sync::Cache<object_store::path::Path, Arc<ParquetMetaData>>>,
 }
 
+impl LruCachingFileReader {
+    pub fn new(path: object_store::path::Path, reader: Box<dyn AsyncFileReader>, cache: Arc<moka::sync::Cache<object_store::path::Path, Arc<ParquetMetaData>>>) -> LruCachingFileReader {
+        LruCachingFileReader {
+            path,
+            reader,
+            cache,
+        }
+    }
+}
+
 impl AsyncFileReader for LruCachingFileReader {
     fn get_bytes(
         &mut self,
@@ -149,14 +164,16 @@ impl AsyncFileReader for LruCachingFileReader {
 
     fn get_metadata(
         &mut self,
+        encryption_config: &Option<ParquetEncryptionConfig>
     ) -> BoxFuture<'_, datafusion::parquet::errors::Result<Arc<ParquetMetaData>>> {
         let cache = self.cache.clone();
         let path = self.path.clone();
+        let encryption_config = encryption_config.clone();
         async move {
             match cache.get(&path) {
                 Some(metadata) => Ok(metadata),
                 None => {
-                    let metadata = self.reader.get_metadata().await?;
+                    let metadata = self.reader.get_metadata(&encryption_config).await?;
                     cache.insert(path, metadata.clone());
                     Ok(metadata)
                 }

rust/cubestore/cubestore/src/queryplanner/metadata_cache.rs

@@ -87,7 +87,7 @@ use datafusion::physical_plan::{
     collect, DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan, Partitioning,
     PlanProperties, SendableRecordBatchStream,
 };
-use datafusion::prelude::SessionContext;
+use datafusion::prelude::{SessionConfig, SessionContext};
 use datafusion::sql::parser::Statement;
 use datafusion::sql::planner::{ContextProvider, SqlToRel};
 use datafusion::{cube_ext, datasource::TableProvider};
@@ -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,);
@@ -245,8 +245,8 @@ impl QueryPlannerImpl {
 }
 
 impl QueryPlannerImpl {
-    pub fn make_execution_context() -> SessionContext {
-        let context = SessionContext::new();
+    pub fn execution_context_helper(config: SessionConfig) -> SessionContext {
+        let context = SessionContext::new_with_config(config);
         // TODO upgrade DF: build SessionContexts consistently -- that now means check all appropriate SessionContext constructors use this make_execution_context or execution_context function.
         for udaf in registerable_aggregate_udfs() {
             context.register_udaf(udaf);
@@ -266,8 +266,13 @@ impl QueryPlannerImpl {
         context
     }
 
+    pub fn make_execution_context() -> SessionContext {
+        Self::execution_context_helper(SessionConfig::new())
+    }
+
+    // TODO upgrade DF: Don't be async
     async fn execution_context(&self) -> Result<Arc<SessionContext>, CubeError> {
-        Ok(Arc::new(Self::make_execution_context()))
+        Ok(Arc::new(Self::execution_context_helper(self.metadata_cache_factory.make_session_config())))
     }
 }
 

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

@@ -19,13 +19,11 @@ impl PartitionFilter {
     const SIZE_LIMIT: usize = 50;
 
     pub fn extract(s: &Schema, filters: &[Expr]) -> PartitionFilter {
-        println!("Calling extract on filters {:?}", filters);
         let builder = Builder { schema: s };
 
         let mut r = vec![];
         for f in filters {
             r = builder.extract_filter(f, r);
-            println!("Extracted.  r = {:?}", r);
         }
 
         PartitionFilter { min_max: r }

rust/cubestore/cubestore/src/queryplanner/partition_filter.rs

@@ -97,6 +97,7 @@ use super::udfs::{
     aggregate_udf_by_kind, registerable_aggregate_udfs, registerable_arc_aggregate_udfs,
     registerable_arc_scalar_udfs, CubeAggregateUDFKind,
 };
+use super::QueryPlannerImpl;
 
 #[automock]
 #[async_trait]
@@ -139,14 +140,21 @@ pub trait QueryExecutor: DIService + Send + Sync {
 crate::di_service!(MockQueryExecutor, [QueryExecutor]);
 
 pub struct QueryExecutorImpl {
-    // TODO: Why do we need a MetadataCacheFactory when we have a ParquetMetadataCache?
+    // TODO: Why do we need a MetadataCacheFactory when we have a ParquetMetadataCache?  (We use its make_session_config() now, TODO rename stuff)
     metadata_cache_factory: Arc<dyn MetadataCacheFactory>,
     parquet_metadata_cache: Arc<dyn CubestoreParquetMetadataCache>,
     memory_handler: Arc<dyn MemoryHandler>,
 }
 
 crate::di_service!(QueryExecutorImpl, [QueryExecutor]);
 
+impl QueryExecutorImpl {
+    fn execution_context(&self) -> Result<Arc<SessionContext>, CubeError> {
+        // This is supposed to be identical to QueryImplImpl::execution_context.
+        Ok(Arc::new(QueryPlannerImpl::execution_context_helper(self.metadata_cache_factory.make_session_config())))
+    }
+}
+
 #[async_trait]
 impl QueryExecutor for QueryExecutorImpl {
     #[instrument(level = "trace", skip(self, plan, cluster))]
@@ -174,7 +182,8 @@ impl QueryExecutor for QueryExecutorImpl {
 
         let execution_time = SystemTime::now();
 
-        let results = collect(split_plan.clone(), Arc::new(TaskContext::default()))
+        let session_context = self.execution_context()?;
+        let results = collect(split_plan.clone(), session_context.task_ctx())
             .instrument(collect_span)
             .await;
         let execution_time = execution_time.elapsed()?;
@@ -241,8 +250,9 @@ impl QueryExecutor for QueryExecutorImpl {
         );
 
         let execution_time = SystemTime::now();
+        let session_context = self.execution_context()?;
         // TODO context
-        let results = collect(worker_plan.clone(), Arc::new(TaskContext::default()))
+        let results = collect(worker_plan.clone(), session_context.task_ctx())
             .instrument(tracing::span!(
                 tracing::Level::TRACE,
                 "collect_physical_plan"

rust/cubestore/cubestore/src/queryplanner/query_executor.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/compaction.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/store/mod.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();