chore(cubestore): Upgrade DF: Deduplicate and simplify SessionContext creation

Author: srh

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

@@ -3,6 +3,7 @@ pub mod optimizations;
 pub mod panic;
 mod partition_filter;
 mod planning;
+use datafusion::execution::runtime_env::RuntimeEnv;
 use datafusion::logical_expr::planner::ExprPlanner;
 use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
 use datafusion_datasource::memory::MemorySourceConfig;
@@ -19,7 +20,6 @@ mod topk;
 pub mod trace_data_loaded;
 use serialized_plan::PreSerializedPlan;
 pub use topk::MIN_TOPK_STREAM_ROWS;
-use udfs::{registerable_aggregate_udfs, registerable_scalar_udfs};
 mod filter_by_key_range;
 pub mod info_schema;
 pub mod merge_sort;
@@ -56,6 +56,7 @@ use crate::queryplanner::topk::ClusterAggregateTopKLower;
 use crate::queryplanner::metadata_cache::MetadataCacheFactory;
 use crate::queryplanner::optimizations::rolling_optimizer::RollingOptimizerRule;
 use crate::queryplanner::pretty_printers::{pp_plan_ext, PPOptions};
+use crate::queryplanner::udfs::{registerable_aggregate_udfs_iter, registerable_scalar_udfs_iter};
 use crate::sql::cache::SqlResultCache;
 use crate::sql::InlineTables;
 use crate::store::DataFrame;
@@ -72,7 +73,7 @@ use datafusion::common::{plan_datafusion_err, TableReference};
 use datafusion::config::ConfigOptions;
 use datafusion::datasource::{provider_as_source, TableType};
 use datafusion::error::DataFusionError;
-use datafusion::execution::{SessionState, TaskContext};
+use datafusion::execution::{SessionState, SessionStateBuilder, TaskContext};
 use datafusion::logical_expr::{
     AggregateUDF, Expr, Extension, LogicalPlan, ScalarUDF, TableProviderFilterPushDown,
     TableSource, WindowUDF,
@@ -280,22 +281,39 @@ impl QueryPlannerImpl {
 }
 
 impl QueryPlannerImpl {
+    /// Has the user defined functions to define query language behavior, but might exclude Cube
+    /// optimizer rules or other parameters affecting execution performance.  This is used by
+    /// `QueryPlannerImpl::make_execution_context`.
+    pub fn minimal_session_state_from_final_config(config: SessionConfig) -> SessionStateBuilder {
+        let mut state_builder = SessionStateBuilder::new()
+            .with_config(config)
+            .with_runtime_env(Arc::new(RuntimeEnv::default()))
+            .with_default_features();
+        state_builder
+            .aggregate_functions()
+            .get_or_insert_default()
+            .extend(registerable_aggregate_udfs_iter().map(Arc::new));
+        state_builder
+            .scalar_functions()
+            .get_or_insert_default()
+            .extend(registerable_scalar_udfs_iter().map(Arc::new));
+        state_builder
+    }
+
     pub fn make_execution_context(mut config: SessionConfig) -> SessionContext {
         // The config parameter is from metadata_cache_factory (which we need to rename) but doesn't
         // include all necessary configs.
         config
             .options_mut()
             .execution
             .dont_parallelize_sort_preserving_merge_exec_inputs = true;
-        let context = SessionContext::new_with_config(config);
+
         // TODO upgrade DF: build SessionContexts consistently
-        for udaf in registerable_aggregate_udfs() {
-            context.register_udaf(udaf);
-        }
-        for udf in registerable_scalar_udfs() {
-            context.register_udf(udf);
-        }
-        context.add_optimizer_rule(Arc::new(RollingOptimizerRule {}));
+        let state = Self::minimal_session_state_from_final_config(config)
+            .with_optimizer_rule(Arc::new(RollingOptimizerRule {}))
+            .build();
+
+        let context = SessionContext::new_with_state(state);
 
         // TODO upgrade DF
         // context

rust/cubestore/cubestore/src/queryplanner/query_executor.rs

@@ -46,8 +46,7 @@ use datafusion::datasource::physical_plan::{
 use datafusion::datasource::{TableProvider, TableType};
 use datafusion::error::DataFusionError;
 use datafusion::error::Result as DFResult;
-use datafusion::execution::runtime_env::RuntimeEnv;
-use datafusion::execution::{SessionStateBuilder, TaskContext};
+use datafusion::execution::TaskContext;
 use datafusion::logical_expr::{Expr, LogicalPlan};
 use datafusion::physical_expr;
 use datafusion::physical_expr::LexOrdering;
@@ -96,7 +95,6 @@ use std::time::SystemTime;
 use tracing::{instrument, Instrument};
 
 use super::serialized_plan::PreSerializedPlan;
-use super::udfs::{registerable_arc_aggregate_udfs, registerable_arc_scalar_udfs};
 use super::{try_make_memory_data_source, QueryPlannerImpl};
 
 #[automock]
@@ -444,15 +442,9 @@ impl QueryExecutorImpl {
         &self,
         query_planner: CubeQueryPlanner,
     ) -> Result<Arc<SessionContext>, CubeError> {
-        let runtime = Arc::new(RuntimeEnv::default());
         let config = self.session_config();
-        let session_state = SessionStateBuilder::new()
-            .with_config(config)
-            .with_runtime_env(runtime)
-            .with_default_features()
+        let session_state = QueryPlannerImpl::minimal_session_state_from_final_config(config)
             .with_query_planner(Arc::new(query_planner))
-            .with_aggregate_functions(registerable_arc_aggregate_udfs())
-            .with_scalar_functions(registerable_arc_scalar_udfs())
             .with_physical_optimizer_rules(self.physical_optimizer_rules())
             .build();
         let ctx = SessionContext::new_with_state(session_state);

rust/cubestore/cubestore/src/queryplanner/serialized_plan.rs

@@ -1,4 +1,3 @@
-use super::udfs::{registerable_aggregate_udfs, registerable_scalar_udfs};
 use crate::metastore::table::{Table, TablePath};
 use crate::metastore::{Chunk, IdRow, Index, Partition};
 use crate::queryplanner::panic::PanicWorkerNode;
@@ -7,7 +6,9 @@ use crate::queryplanner::providers::InfoSchemaQueryCacheTableProvider;
 use crate::queryplanner::query_executor::{CubeTable, InlineTableId, InlineTableProvider};
 use crate::queryplanner::rolling::RollingWindowAggregate;
 use crate::queryplanner::topk::{ClusterAggregateTopKLower, ClusterAggregateTopKUpper};
-use crate::queryplanner::{pretty_printers, CubeTableLogical, InfoSchemaTableProvider};
+use crate::queryplanner::{
+    pretty_printers, CubeTableLogical, InfoSchemaTableProvider, QueryPlannerImpl,
+};
 use crate::table::Row;
 use crate::CubeError;
 use datafusion::arrow::datatypes::SchemaRef;
@@ -27,7 +28,7 @@ use datafusion::logical_expr::{
     Projection, RecursiveQuery, Repartition, Sort, Subquery, SubqueryAlias, TableScan, Union,
     Unnest, Values, Window,
 };
-use datafusion::prelude::SessionContext;
+use datafusion::prelude::{SessionConfig, SessionContext};
 use datafusion_proto::bytes::logical_plan_from_bytes_with_extension_codec;
 use datafusion_proto::logical_plan::LogicalExtensionCodec;
 use std::collections::HashMap;
@@ -1037,16 +1038,13 @@ impl SerializedPlan {
         chunk_id_to_record_batches: HashMap<u64, Vec<RecordBatch>>,
         parquet_metadata_cache: Arc<dyn ParquetFileReaderFactory>,
     ) -> Result<LogicalPlan, CubeError> {
-        // TODO DF upgrade SessionContext::new()
-        // After this comment was made, we now register_udaf... what else?
-        let session_context = SessionContext::new();
-        // TODO DF upgrade: consistently build SessionContexts/register udafs/udfs.
-        for udaf in registerable_aggregate_udfs() {
-            session_context.register_udaf(udaf);
-        }
-        for udf in registerable_scalar_udfs() {
-            session_context.register_udf(udf);
-        }
+        // TODO upgrade DF: We might avoid constructing so many one-time-use SessionContexts.
+
+        // We need registered Cube UDFs and UDAFs (and there are no UDWFs) to deserialize the plan,
+        // but not much else.
+        let session_context = SessionContext::new_with_state(
+            QueryPlannerImpl::minimal_session_state_from_final_config(SessionConfig::new()).build(),
+        );
 
         let logical_plan = logical_plan_from_bytes_with_extension_codec(
             self.logical_plan.as_slice(),

rust/cubestore/cubestore/src/queryplanner/topk/plan.rs

@@ -3,7 +3,7 @@ use crate::queryplanner::topk::execute::{AggregateTopKExec, TopKAggregateFunctio
 use crate::queryplanner::topk::{
     ClusterAggregateTopKLower, ClusterAggregateTopKUpper, SortColumn, MIN_TOPK_STREAM_ROWS,
 };
-use crate::queryplanner::udfs::{scalar_udf_by_kind, CubeScalarUDFKind, HllCardinality};
+use crate::queryplanner::udfs::HllCardinality;
 use datafusion::arrow::compute::SortOptions;
 use datafusion::arrow::datatypes::{DataType, Schema};
 use datafusion::common::tree_node::{Transformed, TreeNode};
@@ -21,7 +21,8 @@ use datafusion::physical_plan::{ExecutionPlan, PhysicalExpr};
 
 use datafusion::common::{DFSchema, DFSchemaRef, Spans};
 use datafusion::logical_expr::{
-    Aggregate, Extension, FetchType, Filter, Limit, LogicalPlan, Projection, SkipType, SortExpr,
+    Aggregate, Extension, FetchType, Filter, Limit, LogicalPlan, Projection, ScalarUDF, SkipType,
+    SortExpr,
 };
 use datafusion::physical_planner::{create_aggregate_expr_and_maybe_filter, PhysicalPlanner};
 use datafusion::prelude::Expr;
@@ -678,7 +679,7 @@ pub fn make_sort_expr(
     // schema in create_physical_expr.
     match fun {
         TopKAggregateFunction::Merge => {
-            let udf = scalar_udf_by_kind(CubeScalarUDFKind::HllCardinality);
+            let udf = Arc::new(ScalarUDF::new_from_impl(HllCardinality::new()));
             Arc::new(ScalarFunctionExpr::new(
                 HllCardinality::static_name(),
                 udf,

rust/cubestore/cubestore/src/queryplanner/udfs.rs

@@ -21,45 +21,16 @@ use serde_derive::{Deserialize, Serialize};
 use std::any::Any;
 use std::sync::Arc;
 
-#[derive(Copy, Clone, Debug, Serialize, Deserialize)]
-pub enum CubeScalarUDFKind {
-    HllCardinality, // cardinality(), accepting the HyperLogLog sketches.
-    UnixTimestamp,
-    DateAdd,
-    DateSub,
-    DateBin,
-    ConvertTz,
-}
-
-pub fn scalar_udf_by_kind(k: CubeScalarUDFKind) -> Arc<ScalarUDF> {
-    match k {
-        CubeScalarUDFKind::HllCardinality => Arc::new(HllCardinality::descriptor()),
-        CubeScalarUDFKind::UnixTimestamp => {
-            Arc::new(ScalarUDF::new_from_impl(UnixTimestamp::new()))
-        }
-        CubeScalarUDFKind::DateAdd => Arc::new(ScalarUDF::new_from_impl(DateAddSub::new_add())),
-        CubeScalarUDFKind::DateSub => Arc::new(ScalarUDF::new_from_impl(DateAddSub::new_sub())),
-        CubeScalarUDFKind::DateBin => Arc::new(ScalarUDF::new_from_impl(DateBin::new())),
-        CubeScalarUDFKind::ConvertTz => Arc::new(ScalarUDF::new_from_impl(ConvertTz::new())),
-    }
-}
-
-pub fn registerable_scalar_udfs() -> Vec<ScalarUDF> {
-    vec![
-        HllCardinality::descriptor(),
+pub fn registerable_scalar_udfs_iter() -> impl Iterator<Item = ScalarUDF> {
+    [
+        ScalarUDF::new_from_impl(HllCardinality::new()),
         ScalarUDF::new_from_impl(DateBin::new()),
         ScalarUDF::new_from_impl(DateAddSub::new_add()),
         ScalarUDF::new_from_impl(DateAddSub::new_sub()),
         ScalarUDF::new_from_impl(UnixTimestamp::new()),
         ScalarUDF::new_from_impl(ConvertTz::new()),
     ]
-}
-
-pub fn registerable_arc_scalar_udfs() -> Vec<Arc<ScalarUDF>> {
-    registerable_scalar_udfs()
-        .into_iter()
-        .map(Arc::new)
-        .collect()
+    .into_iter()
 }
 
 #[derive(Copy, Clone, Debug, Serialize, Deserialize)]
@@ -68,18 +39,12 @@ pub enum CubeAggregateUDFKind {
     Xirr,
 }
 
-pub fn registerable_aggregate_udfs() -> Vec<AggregateUDF> {
-    vec![
+pub fn registerable_aggregate_udfs_iter() -> impl Iterator<Item = AggregateUDF> {
+    [
         AggregateUDF::new_from_impl(HllMergeUDF::new()),
         AggregateUDF::new_from_impl(XirrUDF::new()),
     ]
-}
-
-pub fn registerable_arc_aggregate_udfs() -> Vec<Arc<AggregateUDF>> {
-    registerable_aggregate_udfs()
-        .into_iter()
-        .map(Arc::new)
-        .collect()
+    .into_iter()
 }
 
 pub fn aggregate_udf_by_kind(k: CubeAggregateUDFKind) -> AggregateUDF {
@@ -586,9 +551,6 @@ impl HllCardinality {
 
         HllCardinality { signature }
     }
-    fn descriptor() -> ScalarUDF {
-        return ScalarUDF::new_from_impl(HllCardinality::new());
-    }
 
     /// Lets us call [`ScalarFunctionExpr::new`] in some cases without elaborately computing return
     /// type or using [`ScalarFunctionExpr::try_new`].

rust/cubestore/cubestore/src/streaming/topic_table_provider.rs

@@ -1,5 +1,5 @@
 use crate::metastore::Column;
-use crate::queryplanner::udfs::{registerable_arc_aggregate_udfs, registerable_arc_scalar_udfs};
+use crate::queryplanner::udfs::{registerable_aggregate_udfs_iter, registerable_scalar_udfs_iter};
 use crate::CubeError;
 use async_trait::async_trait;
 use chrono::{TimeZone, Utc};
@@ -47,7 +47,7 @@ impl TopicTableProvider {
                 .collect::<Vec<Field>>(),
         ));
         let mut udfs = SessionStateDefaults::default_scalar_functions();
-        udfs.append(&mut registerable_arc_scalar_udfs());
+        udfs.extend(registerable_scalar_udfs_iter().map(Arc::new));
         udfs.push(Arc::new(
             ScalarUDF::new_from_impl(ParseTimestampFunc::new()),
         ));
@@ -62,7 +62,7 @@ impl TopicTableProvider {
             .collect();
 
         let mut udafs = SessionStateDefaults::default_aggregate_functions();
-        udafs.append(&mut registerable_arc_aggregate_udfs());
+        udafs.extend(registerable_aggregate_udfs_iter().map(Arc::new));
 
         let udafs = udafs
             .into_iter()