alamb commented on code in PR #11403:
URL: https://github.com/apache/datafusion/pull/11403#discussion_r1677188254
##########
datafusion/core/src/execution/session_state.rs:
##########
@@ -976,6 +872,713 @@ impl SessionState {
}
}
+/// A builder to be used for building [`SessionState`]'s. Defaults will
+/// be used for all values unless explicitly provided.
+pub struct SessionStateBuilder {
+ session_id: Option<String>,
+ analyzer: Option<Analyzer>,
+ expr_planners: Option<Vec<Arc<dyn ExprPlanner>>>,
+ optimizer: Option<Optimizer>,
+ physical_optimizers: Option<PhysicalOptimizer>,
+ query_planner: Option<Arc<dyn QueryPlanner + Send + Sync>>,
+ catalog_list: Option<Arc<dyn CatalogProviderList>>,
+ table_functions: Option<HashMap<String, Arc<TableFunction>>>,
+ scalar_functions: Option<Vec<Arc<ScalarUDF>>>,
+ aggregate_functions: Option<Vec<Arc<AggregateUDF>>>,
+ window_functions: Option<Vec<Arc<WindowUDF>>>,
+ serializer_registry: Option<Arc<dyn SerializerRegistry>>,
+ file_formats: Option<Vec<Arc<dyn FileFormatFactory>>>,
+ config: Option<SessionConfig>,
+ table_options: Option<TableOptions>,
+ execution_props: Option<ExecutionProps>,
+ table_factories: Option<HashMap<String, Arc<dyn TableProviderFactory>>>,
+ runtime_env: Option<Arc<RuntimeEnv>>,
+ function_factory: Option<Arc<dyn FunctionFactory>>,
+ // fields to support convenience functions
+ analyzer_rules: Option<Vec<Arc<dyn AnalyzerRule + Send + Sync>>>,
+ optimizer_rules: Option<Vec<Arc<dyn OptimizerRule + Send + Sync>>>,
+ physical_optimizer_rules: Option<Vec<Arc<dyn PhysicalOptimizerRule + Send
+ Sync>>>,
+}
+
+impl SessionStateBuilder {
+ /// Returns a new [`SessionStateBuilder`] with no options set.
+ pub fn new() -> Self {
+ Self {
+ session_id: None,
+ analyzer: None,
+ expr_planners: None,
+ optimizer: None,
+ physical_optimizers: None,
+ query_planner: None,
+ catalog_list: None,
+ table_functions: None,
+ scalar_functions: None,
+ aggregate_functions: None,
+ window_functions: None,
+ serializer_registry: None,
+ file_formats: None,
+ table_options: None,
+ config: None,
+ execution_props: None,
+ table_factories: None,
+ runtime_env: None,
+ function_factory: None,
+ // fields to support convenience functions
+ analyzer_rules: None,
+ optimizer_rules: None,
+ physical_optimizer_rules: None,
+ }
+ }
+
+ /// Returns a new [SessionStateBuilder] based on an existing [SessionState]
+ /// The session id for the new builder will be unset; all other fields will
+ /// be cloned from what is set in the provided session state
+ pub fn new_from_existing(existing: SessionState) -> Self {
+ let cloned = existing.clone();
+
+ Self {
+ session_id: None,
+ analyzer: Some(cloned.analyzer),
+ expr_planners: Some(cloned.expr_planners),
+ optimizer: Some(cloned.optimizer),
+ physical_optimizers: Some(cloned.physical_optimizers),
+ query_planner: Some(cloned.query_planner),
+ catalog_list: Some(cloned.catalog_list),
+ table_functions: Some(cloned.table_functions),
+ scalar_functions:
Some(cloned.scalar_functions.into_values().collect_vec()),
+ aggregate_functions: Some(
+ cloned.aggregate_functions.into_values().collect_vec(),
+ ),
+ window_functions:
Some(cloned.window_functions.into_values().collect_vec()),
+ serializer_registry: Some(cloned.serializer_registry),
+ file_formats:
Some(cloned.file_formats.into_values().collect_vec()),
+ config: Some(cloned.config),
+ table_options: Some(cloned.table_options),
+ execution_props: Some(cloned.execution_props),
+ table_factories: Some(cloned.table_factories),
+ runtime_env: Some(cloned.runtime_env),
+ function_factory: cloned.function_factory,
+
+ // fields to support convenience functions
+ analyzer_rules: None,
+ optimizer_rules: None,
+ physical_optimizer_rules: None,
+ }
+ }
+
+ /// Set defaults for table_factories, file formats, expr_planners and
builtin
+ /// scalar and aggregate functions.
+ pub fn with_default_features(mut self) -> Self {
+ self.table_factories =
Some(SessionStateDefaults::default_table_factories());
+ self.file_formats = Some(SessionStateDefaults::default_file_formats());
+ self.expr_planners =
Some(SessionStateDefaults::default_expr_planners());
+ self.scalar_functions =
Some(SessionStateDefaults::default_scalar_functions());
+ self.aggregate_functions =
+ Some(SessionStateDefaults::default_aggregate_functions());
+ self
+ }
+
+ /// Set the session id.
+ pub fn with_session_id(mut self, session_id: String) -> Self {
+ self.session_id = Some(session_id);
+ self
+ }
+
+ /// Set the [`AnalyzerRule`]s optimizer plan rules.
+ pub fn with_analyzer_rules(
+ mut self,
+ rules: Vec<Arc<dyn AnalyzerRule + Send + Sync>>,
+ ) -> Self {
+ self.analyzer = Some(Analyzer::with_rules(rules));
+ self
+ }
+
+ /// Add `analyzer_rule` to the end of the list of
+ /// [`AnalyzerRule`]s used to rewrite queries.
+ pub fn add_analyzer_rule(
+ mut self,
+ analyzer_rule: Arc<dyn AnalyzerRule + Send + Sync>,
+ ) -> Self {
+ let mut rules = self.analyzer_rules.unwrap_or_default();
+ rules.push(analyzer_rule);
+ self.analyzer_rules = Some(rules);
+ self
+ }
+
+ /// Set the [`OptimizerRule`]s used to optimize plans.
+ pub fn with_optimizer_rules(
+ mut self,
+ rules: Vec<Arc<dyn OptimizerRule + Send + Sync>>,
+ ) -> Self {
+ self.optimizer = Some(Optimizer::with_rules(rules));
+ self
+ }
+
+ /// Add `optimizer_rule` to the end of the list of
+ /// [`OptimizerRule`]s used to rewrite queries.
+ pub fn add_optimizer_rule(
+ mut self,
+ optimizer_rule: Arc<dyn OptimizerRule + Send + Sync>,
+ ) -> Self {
+ let mut rules = self.optimizer_rules.unwrap_or_default();
+ rules.push(optimizer_rule);
+ self.optimizer_rules = Some(rules);
+ self
+ }
+
+ /// Set the [`ExprPlanner`]s used to customize the behavior of the SQL
planner.
+ pub fn with_expr_planners(
+ mut self,
+ expr_planners: Vec<Arc<dyn ExprPlanner>>,
+ ) -> Self {
+ self.expr_planners = Some(expr_planners);
+ self
+ }
+
+ /// Set tje [`PhysicalOptimizerRule`]s used to optimize plans.
+ pub fn with_physical_optimizer_rules(
+ mut self,
+ physical_optimizers: Vec<Arc<dyn PhysicalOptimizerRule + Send + Sync>>,
+ ) -> Self {
+ self.physical_optimizers =
+ Some(PhysicalOptimizer::with_rules(physical_optimizers));
+ self
+ }
+
+ /// Add `physical_optimizer_rule` to the end of the list of
+ /// [`PhysicalOptimizerRule`]s used to rewrite queries.
+ pub fn add_physical_optimizer_rule(
+ mut self,
+ physical_optimizer_rule: Arc<dyn PhysicalOptimizerRule + Send + Sync>,
+ ) -> Self {
+ let mut rules = self.physical_optimizer_rules.unwrap_or_default();
+ rules.push(physical_optimizer_rule);
+ self.physical_optimizer_rules = Some(rules);
+ self
+ }
+
+ /// Set the [`QueryPlanner`]
+ pub fn with_query_planner(
+ mut self,
+ query_planner: Arc<dyn QueryPlanner + Send + Sync>,
+ ) -> Self {
+ self.query_planner = Some(query_planner);
+ self
+ }
+
+ /// Set the [`CatalogProviderList`]
+ pub fn with_catalog_list(
+ mut self,
+ catalog_list: Arc<dyn CatalogProviderList>,
+ ) -> Self {
+ self.catalog_list = Some(catalog_list);
+ self
+ }
+
+ /// Set the map of [`TableFunction`]s
+ pub fn with_table_functions(
+ mut self,
+ table_functions: HashMap<String, Arc<TableFunction>>,
+ ) -> Self {
+ self.table_functions = Some(table_functions);
+ self
+ }
+
+ /// Set the map of [`ScalarUDF`]s
+ pub fn with_scalar_functions(
+ mut self,
+ scalar_functions: Vec<Arc<ScalarUDF>>,
+ ) -> Self {
+ self.scalar_functions = Some(scalar_functions);
+ self
+ }
+
+ /// Set the map of [`AggregateUDF`]s
+ pub fn with_aggregate_functions(
+ mut self,
+ aggregate_functions: Vec<Arc<AggregateUDF>>,
+ ) -> Self {
+ self.aggregate_functions = Some(aggregate_functions);
+ self
+ }
+
+ /// Set the map of [`WindowUDF`]s
+ pub fn with_window_functions(
+ mut self,
+ window_functions: Vec<Arc<WindowUDF>>,
+ ) -> Self {
+ self.window_functions = Some(window_functions);
+ self
+ }
+
+ /// Set the [`SerializerRegistry`]
+ pub fn with_serializer_registry(
+ mut self,
+ serializer_registry: Arc<dyn SerializerRegistry>,
+ ) -> Self {
+ self.serializer_registry = Some(serializer_registry);
+ self
+ }
+
+ /// Set the map of [`FileFormatFactory`]s
+ pub fn with_file_formats(
+ mut self,
+ file_formats: Vec<Arc<dyn FileFormatFactory>>,
+ ) -> Self {
+ self.file_formats = Some(file_formats);
+ self
+ }
+
+ /// Set the [`SessionConfig`]
+ pub fn with_config(mut self, config: SessionConfig) -> Self {
+ self.config = Some(config);
+ self
+ }
+
+ /// Set the [`TableOptions`]
+ pub fn with_table_options(mut self, table_options: TableOptions) -> Self {
+ self.table_options = Some(table_options);
+ self
+ }
+
+ /// Set the [`ExecutionProps`]
+ pub fn with_execution_props(mut self, execution_props: ExecutionProps) ->
Self {
+ self.execution_props = Some(execution_props);
+ self
+ }
+
+ /// Set the map of [`TableProviderFactory`]s
+ pub fn with_table_factories(
+ mut self,
+ table_factories: HashMap<String, Arc<dyn TableProviderFactory>>,
+ ) -> Self {
+ self.table_factories = Some(table_factories);
+ self
+ }
+
+ /// Set the [`RuntimeEnv`]
+ pub fn with_runtime_env(mut self, runtime_env: Arc<RuntimeEnv>) -> Self {
+ self.runtime_env = Some(runtime_env);
+ self
+ }
+
+ /// Set a [`FunctionFactory`] to handle `CREATE FUNCTION` statements
+ pub fn with_function_factory(
+ mut self,
+ function_factory: Option<Arc<dyn FunctionFactory>>,
+ ) -> Self {
+ self.function_factory = function_factory;
+ self
+ }
+
+ /// Builds a [`SessionState`] with the current configuration.
+ ///
+ /// Note that there is an explicit option for enabling catalog and schema
defaults
+ /// in [SessionConfig::create_default_catalog_and_schema] which if enabled
+ /// will be built here.
+ pub fn build(self) -> SessionState {
+ let config = self.config.unwrap_or_default();
+ let runtime_env =
self.runtime_env.unwrap_or(Arc::new(RuntimeEnv::default()));
+
+ let mut state = SessionState {
Review Comment:
I switched a few places in the sql planner to use it after we got a few bugs
due to by newly fields being (silently) ignored.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
