diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 527ab0491c0..597fdd66d96 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -296,6 +296,12 @@ impl RegionServer { }; let region_id = RegionId::from_u64(request.region_id); + let engine_ctx = self + .inner + .query_engine + .engine_context(query_ctx) + .context(NewPlanDecoderSnafu)?; + let query_ctx = engine_ctx.query_ctx(); let catalog_list = Arc::new(NameAwareCatalogList::new( self.clone(), region_id, @@ -306,12 +312,7 @@ impl RegionServer { common_telemetry::info!("Handle remote read for region: {}", region_id); } - let decoder = self - .inner - .query_engine - .engine_context(query_ctx.clone()) - .new_plan_decoder() - .context(NewPlanDecoderSnafu)?; + let decoder = engine_ctx.new_plan_decoder().context(NewPlanDecoderSnafu)?; let plan = decoder .decode(Bytes::from(request.plan), catalog_list, false) diff --git a/src/datanode/src/tests.rs b/src/datanode/src/tests.rs index 843ac275220..9ca771d2186 100644 --- a/src/datanode/src/tests.rs +++ b/src/datanode/src/tests.rs @@ -90,7 +90,10 @@ impl QueryEngine for MockQueryEngine { unimplemented!() } - fn engine_context(&self, _query_ctx: QueryContextRef) -> QueryEngineContext { + fn engine_context( + &self, + _query_ctx: QueryContextRef, + ) -> query::error::Result { unimplemented!() } fn engine_state(&self) -> &QueryEngineState { diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 22d2afaa7bc..5a99ccbf584 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -292,6 +292,21 @@ impl Instance { async fn query_statement(&self, stmt: Statement, query_ctx: QueryContextRef) -> Result { check_permission(self.plugins.clone(), &stmt, &query_ctx)?; + let query_ctx = if matches!(&stmt, Statement::SetVariables(_)) { + query_ctx + } else { + Arc::new( + query_ctx + .query_option_snapshot() + .map_err(|error| { + query::error::InvalidQueryContextExtensionSnafu { + reason: error.to_string(), + } + .build() + }) + .context(PlanStatementSnafu)?, + ) + }; let query_interceptor = self.plugins.get::>(); let query_interceptor = query_interceptor.as_ref(); diff --git a/src/frontend/src/instance/grpc.rs b/src/frontend/src/instance/grpc.rs index e2435b65689..13898f71247 100644 --- a/src/frontend/src/instance/grpc.rs +++ b/src/frontend/src/instance/grpc.rs @@ -127,9 +127,12 @@ impl GrpcQueryHandler for Instance { let timer = GRPC_HANDLE_PLAN_ELAPSED.start_timer(); // use dummy catalog to provide table - let plan_decoder = self + let engine_ctx = self .query_engine() .engine_context(ctx.clone()) + .context(PlanStatementSnafu)?; + let ctx = engine_ctx.query_ctx(); + let plan_decoder = engine_ctx .new_plan_decoder() .context(PlanStatementSnafu)?; @@ -475,11 +478,12 @@ impl Instance { })?; // use dummy catalog to provide table - let plan_decoder = self + let engine_ctx = self .query_engine() .engine_context(ctx.clone()) - .new_plan_decoder() .context(PlanStatementSnafu)?; + let ctx = engine_ctx.query_ctx(); + let plan_decoder = engine_ctx.new_plan_decoder().context(PlanStatementSnafu)?; let dummy_catalog_list = Arc::new( catalog::table_source::dummy_catalog::DummyCatalogList::new_with_query_ctx( @@ -522,7 +526,11 @@ impl Instance { let insert_into = add_insert_to_logical_plan(table_name, table_source, logical_plan) .context(SubstraitDecodeLogicalPlanSnafu)?; - let engine_ctx = self.query_engine().engine_context(ctx.clone()); + let engine_ctx = self + .query_engine() + .engine_context(ctx.clone()) + .context(PlanStatementSnafu)?; + let ctx = engine_ctx.query_ctx(); let state = engine_ctx.state(); // Analyze the plan let analyzed_plan = state diff --git a/src/operator/src/statement/set.rs b/src/operator/src/statement/set.rs index c5246f065c2..194d068eb35 100644 --- a/src/operator/src/statement/set.rs +++ b/src/operator/src/statement/set.rs @@ -22,7 +22,7 @@ use session::ReadPreference; use session::context::Channel::Postgres; use session::context::QueryContextRef; use session::session_config::{PGByteaOutputValue, PGDateOrder, PGDateTimeStyle, PGIntervalStyle}; -use snafu::{OptionExt, ResultExt, ensure}; +use snafu::{IntoError, OptionExt, ResultExt, ensure}; use sql::ast::{Expr, Ident, Value}; use sql::statements::set_variables::SetVariables; use sqlparser::ast::ValueWithSpan; @@ -264,8 +264,8 @@ pub fn set_query_option(variable: &str, exprs: Vec, ctx: QueryContextRef) } let [Expr::Value(value)] = exprs.as_slice() else { - return InvalidConfigValueSnafu { - source: session::session_config::Error::InvalidConfigValue { + return Err(InvalidConfigValueSnafu.into_error( + session::session_config::Error::InvalidConfigValue { name: variable.to_string(), value: exprs .iter() @@ -275,8 +275,7 @@ pub fn set_query_option(variable: &str, exprs: Vec, ctx: QueryContextRef) hint: "Expected exactly one literal value".to_string(), location: snafu::Location::new(file!(), line!(), column!()), }, - } - .fail(); + )); }; let value = match &value.value { Value::Boolean(value) => value.to_string(), @@ -284,15 +283,14 @@ pub fn set_query_option(variable: &str, exprs: Vec, ctx: QueryContextRef) | Value::SingleQuotedString(value) | Value::DoubleQuotedString(value) => value.clone(), _ => { - return InvalidConfigValueSnafu { - source: session::session_config::Error::InvalidConfigValue { + return Err(InvalidConfigValueSnafu.into_error( + session::session_config::Error::InvalidConfigValue { name: variable.to_string(), value: value.value.to_string(), hint: "Expected a boolean, number, or quoted string literal".to_string(), location: snafu::Location::new(file!(), line!(), column!()), }, - } - .fail(); + )); } }; if let Some((key, value)) = session::query_options::parse_query_option(variable, &value) diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index 7df99f059ca..398f4eeed3f 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -27,7 +27,6 @@ use async_trait::async_trait; use common_base::Plugins; use common_catalog::consts::is_readonly_table; use common_error::ext::BoxedError; -use common_function::function::FunctionContext; use common_function::function_factory::ScalarFunctionFactory; use common_query::{Output, OutputData, OutputMeta}; use common_recordbatch::adapter::{RecordBatchStreamAdapter, RegionQueryStatCounters}; @@ -56,9 +55,7 @@ use tracing::Span; use crate::analyze::DistAnalyzeExec; pub use crate::datafusion::planner::DfContextProviderAdapter; -use crate::dist_plan::{ - DistPlannerOptions, MergeScanLogicalPlan, RemoteDynFilterReceiverInjectorRef, -}; +use crate::dist_plan::{MergeScanLogicalPlan, RemoteDynFilterReceiverInjectorRef}; use crate::error::{ CatalogSnafu, CreateRecordBatchSnafu, MissingTableMutationHandlerSnafu, MissingTimestampColumnSnafu, QueryExecutionSnafu, Result, TableMutationSnafu, @@ -69,7 +66,6 @@ use crate::metrics::{ OnDone, QUERY_STAGE_ELAPSED, maybe_attach_region_watermark_metrics, should_collect_region_watermark_from_query_ctx, }; -use crate::options::ScheduledTimeExtension; use crate::physical_wrapper::PhysicalPlanWrapperRef; use crate::planner::{DfLogicalPlanner, LogicalPlanner}; use crate::query_engine::{DescribeResult, QueryEngineContext, QueryEngineState}; @@ -179,7 +175,8 @@ impl DatafusionQueryEngine { plan: LogicalPlan, query_ctx: QueryContextRef, ) -> Result { - let mut ctx = self.engine_context(query_ctx.clone()); + let mut ctx = self.engine_context(query_ctx)?; + let query_ctx = ctx.query_ctx(); let plan = if let Some(receiver_injector) = self.plugins.get::() { @@ -617,94 +614,8 @@ impl QueryEngine for DatafusionQueryEngine { self.state.read_table(table).map_err(Into::into) } - fn engine_context(&self, query_ctx: QueryContextRef) -> QueryEngineContext { - let mut state = self.state.session_state(); - state.config_mut().set_extension(query_ctx.clone()); - state.config_mut().set_extension(self.state.clone()); - // note that hints in "x-greptime-hints" is automatically parsed - // and set to query context's extension, so we can get it from query context. - if let Some(parallelism) = query_ctx.extension(QUERY_PARALLELISM_HINT) { - if let Ok(n) = parallelism.parse::() { - if n > 0 { - let new_cfg = state.config().clone().with_target_partitions(n as usize); - *state.config_mut() = new_cfg; - } - } else { - common_telemetry::warn!( - "Failed to parse query_parallelism: {}, using default value", - parallelism - ); - } - } - - // configure execution options - state.config_mut().options_mut().execution.time_zone = - Some(query_ctx.timezone().to_string()); - - // usually it's impossible to have both `set variable` set by sql client and - // hint in header by grpc client, so only need to deal with them separately - if query_ctx.configuration_parameter().allow_query_fallback() { - state - .config_mut() - .options_mut() - .extensions - .insert(DistPlannerOptions { - allow_query_fallback: true, - }); - } else if let Some(fallback) = query_ctx.extension(QUERY_FALLBACK_HINT) { - // also check the query context for fallback hint - // if it is set, we will enable the fallback - if fallback.to_lowercase().parse::().unwrap_or(false) { - state - .config_mut() - .options_mut() - .extensions - .insert(DistPlannerOptions { - allow_query_fallback: true, - }); - } - } - - state - .config_mut() - .options_mut() - .extensions - .insert(FunctionContext { - query_ctx: query_ctx.clone(), - state: self.engine_state().function_state(), - }); - - // Carry scheduled Flow time through ConfigOptions.extensions so that - // the distributed plan analyzer can read it during expression - // simplification (preventing wall-clock constant-folding of `now()`). - state - .config_mut() - .options_mut() - .extensions - .insert(ScheduledTimeExtension { - scheduled_time: crate::options::scheduled_time_from_ctx(&query_ctx), - }); - - let config_options = state.config_options().clone(); - let _ = state - .execution_props_mut() - .config_options - .insert(config_options); - - // Apply scheduled time from query context if present, so that `now()` / - // `current_timestamp()` functions evaluate against the logical scheduled time - // rather than wall-clock. - match crate::options::parse_scheduled_time_datetime(&query_ctx.extensions()) { - Ok(Some(scheduled_rt)) => { - state.execution_props_mut().query_execution_start_time = Some(scheduled_rt); - } - Ok(None) => {} - Err(err) => { - common_telemetry::warn!(err; "Ignoring invalid scheduled time query extension"); - } - } - - QueryEngineContext::new(state, query_ctx) + fn engine_context(&self, query_ctx: QueryContextRef) -> Result { + self.state.query_engine_context(query_ctx) } fn engine_state(&self) -> &QueryEngineState { @@ -825,7 +736,7 @@ mod tests { use datafusion::physical_plan::expressions::PhysicalSortExpr; use datafusion::physical_plan::joins::{HashJoinExec, JoinOn, PartitionMode}; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; - use datafusion::physical_plan::{ExecutionPlan, PhysicalExpr}; + use datafusion::physical_plan::{ExecutionPlan, PhysicalExpr, displayable}; use datafusion::prelude::{col, lit}; use datafusion_common::{JoinType, NullEquality, ScalarValue}; use datafusion_physical_expr::expressions::Column; @@ -843,6 +754,7 @@ mod tests { use table::table::scan::RegionScanExec; use super::*; + use crate::dist_plan::DistPlannerOptions; use crate::options::QueryOptions; use crate::parser::{QueryLanguageParser, QueryStatement}; use crate::part_sort::PartSortExec; @@ -1084,7 +996,7 @@ mod tests { assert!(super::query_stat_counters(&plan).is_none()); } - async fn create_test_engine() -> QueryEngineRef { + async fn create_test_engine_with_fallback(allow_query_fallback: bool) -> QueryEngineRef { let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap(); let req = RegisterTableRequest { catalog: DEFAULT_CATALOG_NAME.to_string(), @@ -1094,7 +1006,6 @@ mod tests { table: NumbersTable::table(NUMBERS_TABLE_ID), }; catalog_manager.register_table_sync(req).unwrap(); - QueryEngineFactory::new( catalog_manager, None, @@ -1102,11 +1013,204 @@ mod tests { None, None, false, - QueryOptions::default(), + QueryOptions { + allow_query_fallback, + ..QueryOptions::default() + }, ) .query_engine() } + async fn create_test_engine() -> QueryEngineRef { + create_test_engine_with_fallback(false).await + } + + #[tokio::test] + async fn test_query_fallback_request_overrides_engine_default() { + for server_default in [false, true] { + let engine = create_test_engine_with_fallback(server_default).await; + for session_default in [None, Some(false), Some(true)] { + for request in [None, Some("false"), Some("true")] { + let mut ctx = QueryContextBuilder::default().build(); + if let Some(value) = session_default { + ctx.configuration_parameter() + .set_allow_query_fallback(value); + } + if let Some(value) = request { + ctx.set_extension(QUERY_FALLBACK_HINT, value); + } + let engine_ctx = engine.engine_context(Arc::new(ctx)).unwrap(); + let fallback = engine_ctx + .state() + .config_options() + .extensions + .get::() + .is_some_and(|options| options.allow_query_fallback); + assert_eq!( + fallback, + request + .map(|v| v == "true") + .or(session_default) + .unwrap_or(server_default), + "server={server_default}, session={session_default:?}, request={request:?}" + ); + } + } + } + } + + #[tokio::test] + async fn query_options_apply_to_engine_context_without_mutating_defaults() { + let engine = create_test_engine().await; + let mut ctx = QueryContextBuilder::default().build(); + ctx.configuration_parameter() + .set_query_option("datafusion.optimizer.repartition_joins", "true") + .unwrap(); + ctx.configuration_parameter() + .set_query_option("datafusion.optimizer.prefer_hash_join", "true") + .unwrap(); + ctx.configuration_parameter() + .set_query_option("query.parallelism", "8") + .unwrap(); + ctx.configuration_parameter().set_allow_query_fallback(true); + ctx.set_extension("query_parallelism", "3"); + ctx.set_extension("datafusion.optimizer.prefer_hash_join", "false"); + + let engine_ctx = engine.engine_context(Arc::new(ctx)).unwrap(); + let config = engine_ctx.state().config_options(); + assert!(config.optimizer.repartition_joins); + assert!(!config.optimizer.prefer_hash_join); + assert_eq!(engine_ctx.state().config().target_partitions(), 3); + assert!( + config + .extensions + .get::() + .unwrap() + .allow_query_fallback + ); + assert!( + engine + .engine_context(QueryContext::arc()) + .unwrap() + .state() + .config_options() + .optimizer + .prefer_hash_join + ); + } + + #[tokio::test] + async fn query_option_snapshot_survives_query_context_wire_round_trip() { + let engine = create_test_engine().await; + let mut ctx = QueryContextBuilder::default().build(); + ctx.configuration_parameter() + .set_query_option("datafusion.optimizer.repartition_joins", "true") + .unwrap(); + ctx.configuration_parameter() + .set_query_option("query.parallelism", "7") + .unwrap(); + ctx.configuration_parameter() + .set_allow_query_fallback(false); + ctx.set_extension("query.enable_remote_dynamic_filter_pushdown", "false"); + ctx.set_extension("flow.return_region_seq", "true"); + + let original_ctx = Arc::new(ctx); + let first = engine.engine_context(original_ctx.clone()).unwrap(); + original_ctx + .configuration_parameter() + .set_query_option("query.parallelism", "9") + .unwrap(); + let remote_ctx = QueryContext::from(api::v1::QueryContext::from( + first.query_ctx().as_ref().clone(), + )); + let remote = engine.engine_context(Arc::new(remote_ctx)).unwrap(); + let config = remote.state().config_options(); + assert!(config.optimizer.repartition_joins); + assert_eq!(remote.state().config().target_partitions(), 7); + assert!( + !config + .extensions + .get::() + .unwrap() + .allow_query_fallback + ); + assert_eq!( + remote + .query_ctx() + .extension("query.enable_remote_dynamic_filter_pushdown"), + Some("false") + ); + assert_eq!( + remote.query_ctx().extension("flow.return_region_seq"), + Some("true") + ); + } + + #[tokio::test] + async fn invalid_query_options_fail_engine_context_creation() { + let engine = create_test_engine().await; + let mut ctx = QueryContextBuilder::default().build(); + ctx.set_extension("datafusion.optimizer.repartition_joins", "not-a-bool"); + let error = engine.engine_context(Arc::new(ctx)).unwrap_err(); + use common_error::ext::ErrorExt; + assert_eq!( + error.status_code(), + common_error::status_code::StatusCode::InvalidArguments + ); + + let mut ctx = QueryContextBuilder::default().build(); + ctx.set_extension("datafusion.execution.batch_size", "12"); + let error = engine.engine_context(Arc::new(ctx)).unwrap_err(); + assert_eq!( + error.status_code(), + common_error::status_code::StatusCode::InvalidArguments + ); + + let mut ctx = QueryContextBuilder::default().build(); + ctx.set_extension("flow.scheduled_time_millis", "not-a-time"); + let context = Arc::new(ctx); + let stmt = QueryLanguageParser::parse_sql("SELECT 1", &QueryContext::arc()).unwrap(); + let error = engine.planner().plan(&stmt, context).await.unwrap_err(); + assert_eq!( + error.status_code(), + common_error::status_code::StatusCode::InvalidArguments + ); + } + + #[tokio::test] + async fn approved_optimizer_option_changes_join_plan_without_changing_results() { + let engine = create_test_engine().await; + let sql = "SELECT l.number FROM numbers l JOIN numbers r ON l.number = r.number WHERE l.number < 4 ORDER BY l.number"; + let mut plans = Vec::new(); + let mut results = Vec::new(); + for prefer_hash_join in [true, false] { + let mut ctx = QueryContextBuilder::default().build(); + ctx.set_extension( + "datafusion.optimizer.prefer_hash_join", + if prefer_hash_join { "true" } else { "false" }, + ); + ctx.set_extension(QUERY_PARALLELISM_HINT, "1"); + let context = Arc::new(ctx); + let stmt = QueryLanguageParser::parse_sql(sql, &context).unwrap(); + let logical_plan = engine.planner().plan(&stmt, context.clone()).await.unwrap(); + let output = engine.execute(logical_plan, context).await.unwrap(); + plans.push( + displayable(output.meta.plan.as_ref().unwrap().as_ref()) + .indent(true) + .to_string(), + ); + let OutputData::Stream(stream) = output.data else { + panic!("expected query result stream") + }; + results.push(util::collect(stream).await.unwrap()); + } + assert!(plans[0].contains("HashJoinExec"), "{}", plans[0]); + assert!(plans[1].contains("SortMergeJoinExec"), "{}", plans[1]); + assert_eq!(results[0].len(), results[1].len()); + assert_eq!(results[0][0].column(0), results[1][0].column(0)); + assert_eq!(results[0][0].num_rows(), 4); + } + #[tokio::test] async fn test_sql_to_plan() { let engine = create_test_engine().await; @@ -1259,7 +1363,7 @@ mod tests { .as_any() .downcast_ref::() .unwrap(); - let engine_ctx = engine.engine_context(QueryContext::arc()); + let engine_ctx = engine.engine_context(QueryContext::arc()).unwrap(); let state = engine_ctx.state(); let schema = Arc::new(datatypes::schema::Schema::new(vec![ColumnSchema::new( @@ -1319,7 +1423,7 @@ mod tests { .as_any() .downcast_ref::() .unwrap(); - let engine_ctx = engine.engine_context(QueryContext::arc()); + let engine_ctx = engine.engine_context(QueryContext::arc()).unwrap(); let state = engine_ctx.state(); assert!( diff --git a/src/query/src/optimizer/constant_term.rs b/src/query/src/optimizer/constant_term.rs index 70ddcc2f647..79afd005b5d 100644 --- a/src/query/src/optimizer/constant_term.rs +++ b/src/query/src/optimizer/constant_term.rs @@ -469,7 +469,7 @@ mod tests { .await .unwrap(); - let engine_ctx = engine.engine_context(query_ctx); + let engine_ctx = engine.engine_context(query_ctx).unwrap(); let state = engine_ctx.state(); let analyzed_plan = state diff --git a/src/query/src/planner.rs b/src/query/src/planner.rs index cc1c9996483..59f2e9b9e90 100644 --- a/src/query/src/planner.rs +++ b/src/query/src/planner.rs @@ -26,7 +26,6 @@ use common_error::ext::BoxedError; use common_query::promql_annotations::promql_annotation_collector; use common_telemetry::tracing; use datafusion::common::{DFSchema, plan_err}; -use datafusion::execution::SessionStateBuilder; use datafusion::execution::context::SessionState; use datafusion::sql::planner::PlannerContext; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; @@ -90,31 +89,6 @@ impl DfLogicalPlanner { } } - /// Derive a [`SessionState`] whose [`ExecutionProps`] includes - /// `query_execution_start_time` if a scheduled time extension is present - /// in the query context. - fn derive_session_state_with_scheduled_time( - &self, - query_ctx: &QueryContextRef, - ) -> Result { - let extensions = query_ctx.extensions(); - match crate::options::parse_scheduled_time_datetime(&extensions)? { - Some(dt) => { - let execution_props = self - .session_state - .execution_props() - .clone() - .with_query_execution_start_time(dt); - Ok( - SessionStateBuilder::new_from_existing(self.session_state.clone()) - .with_execution_props(execution_props) - .build(), - ) - } - None => Ok(self.session_state.clone()), - } - } - /// Basically the same with `explain_to_plan` in DataFusion, but adapted to Greptime's /// `plan_sql` to support Greptime Statements. async fn explain_to_plan( @@ -218,7 +192,9 @@ impl DfLogicalPlanner { .fail()?; } - let scheduled_state = self.derive_session_state_with_scheduled_time(&query_ctx)?; + let engine_context = self.engine_state.query_engine_context(query_ctx)?; + let query_ctx = engine_context.query_ctx(); + let scheduled_state = engine_context.state().clone(); let table_provider = DfTableSourceProvider::new( self.engine_state.catalog_manager().clone(), self.engine_state.disallow_cross_catalog_query(), @@ -241,7 +217,7 @@ impl DfLogicalPlanner { ) .await?; - let config_options = self.session_state.config().options(); + let config_options = scheduled_state.config().options(); let parser_options = &config_options.sql_parser; let parser_options = ParserOptions { map_string_types_to_utf8view: false, @@ -289,16 +265,17 @@ impl DfLogicalPlanner { normalize_ident: bool, query_ctx: QueryContextRef, ) -> Result { - let scheduled_state = self.derive_session_state_with_scheduled_time(&query_ctx)?; + let engine_context = self.engine_state.query_engine_context(query_ctx)?; + let query_ctx = engine_context.query_ctx(); let context_provider = DfContextProviderAdapter::try_new( self.engine_state.clone(), - scheduled_state, + engine_context.state().clone(), None, query_ctx, ) .await?; - let config_options = self.session_state.config().options(); + let config_options = engine_context.state().config().options(); let parser_options = &config_options.sql_parser; let parser_options: ParserOptions = ParserOptions { map_string_types_to_utf8view: false, @@ -313,7 +290,9 @@ impl DfLogicalPlanner { #[tracing::instrument(skip_all)] async fn plan_pql(&self, stmt: &EvalStmt, query_ctx: QueryContextRef) -> Result { - let mut scheduled_state = self.derive_session_state_with_scheduled_time(&query_ctx)?; + let engine_context = self.engine_state.query_engine_context(query_ctx)?; + let query_ctx = engine_context.query_ctx(); + let mut scheduled_state = engine_context.state().clone(); let promql_annotations = query_ctx.remote_query_id().map(promql_annotation_collector); if let Some(collector) = &promql_annotations { scheduled_state @@ -704,22 +683,19 @@ impl LogicalPlanner for DfLogicalPlanner { query: LogQuery, query_ctx: QueryContextRef, ) -> Result { - let plan_decoder = Arc::new(DefaultPlanDecoder::new( - self.session_state.clone(), - &query_ctx, - )?); + let engine_context = self.engine_state.query_engine_context(query_ctx)?; + let query_ctx = engine_context.query_ctx(); + let state = engine_context.state().clone(); + let plan_decoder = Arc::new(DefaultPlanDecoder::new(state.clone(), &query_ctx)?); let table_provider = DfTableSourceProvider::new( self.engine_state.catalog_manager().clone(), self.engine_state.disallow_cross_catalog_query(), query_ctx, plan_decoder, - self.session_state - .config_options() - .sql_parser - .enable_ident_normalization, + state.config_options().sql_parser.enable_ident_normalization, ); - let mut planner = LogQueryPlanner::new(table_provider, self.session_state.clone()); + let mut planner = LogQueryPlanner::new(table_provider, state); planner .query_to_plan(query) .await @@ -960,7 +936,7 @@ mod tests { .plan(&stmt, query_ctx.clone()) .await .unwrap(); - let context = engine.engine_context(query_ctx.clone()); + let context = engine.engine_context(query_ctx.clone()).unwrap(); let state = context.state(); state .analyzer() diff --git a/src/query/src/query_engine.rs b/src/query/src/query_engine.rs index c512b9c1e5e..0355dc2294d 100644 --- a/src/query/src/query_engine.rs +++ b/src/query/src/query_engine.rs @@ -97,7 +97,7 @@ pub trait QueryEngine: Send + Sync { fn read_table(&self, table: TableRef) -> Result; /// Create a [`QueryEngineContext`]. - fn engine_context(&self, query_ctx: QueryContextRef) -> QueryEngineContext; + fn engine_context(&self, query_ctx: QueryContextRef) -> Result; /// Retrieve the query engine state [`QueryEngineState`] fn engine_state(&self) -> &QueryEngineState; diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index f4d6b29fa2c..d6904c7ece0 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -297,6 +297,7 @@ mod tests { let plan_decoder = engine .engine_context(QueryContext::arc()) + .unwrap() .new_plan_decoder() .unwrap(); let catalog_list = Arc::new(DummyCatalogList::with_table_provider(table_provider)); @@ -389,6 +390,7 @@ mod tests { let table_provider = Arc::new(MemTable::try_new(schema, vec![vec![]]).unwrap()); let plan_decoder = engine .engine_context(QueryContext::arc()) + .unwrap() .new_plan_decoder() .unwrap(); diff --git a/src/query/src/query_engine/state.rs b/src/query/src/query_engine/state.rs index aa3b9417b14..2c6372eb1e2 100644 --- a/src/query/src/query_engine/state.rs +++ b/src/query/src/query_engine/state.rs @@ -21,6 +21,7 @@ use async_trait::async_trait; use catalog::CatalogManagerRef; use common_base::Plugins; use common_function::aggrs::aggr_wrapper::fix_order::FixStateUdafOrderingAnalyzer; +use common_function::function::FunctionContext; use common_function::function_factory::ScalarFunctionFactory; use common_function::function_registry::FUNCTION_REGISTRY; use common_function::handlers::{ @@ -59,6 +60,7 @@ use crate::dist_plan::{ MergeSortExtensionPlanner, RemoteDynFilterReceiverExtensionPlanner, RemoteDynFilterRegistryLease, }; +use crate::error::InvalidQueryContextExtensionSnafu; use crate::metrics::{QUERY_MEMORY_POOL_REJECTED_TOTAL, QUERY_MEMORY_POOL_USAGE_BYTES}; use crate::optimizer::ExtensionAnalyzerRule; use crate::optimizer::const_normalization::ConstNormalizationRule; @@ -78,7 +80,9 @@ use crate::optimizer::string_normalization::StringNormalizationRule; use crate::optimizer::transcribe_atat::TranscribeAtatRule; use crate::optimizer::type_conversion::TypeConversionRule; use crate::optimizer::windowed_sort::WindowedSortPhysicalRule; -use crate::options::{QueryMemoryPoolPolicy, QueryOptions as QueryOptionsNew}; +use crate::options::{ + QueryMemoryPoolPolicy, QueryOptions as QueryOptionsNew, ScheduledTimeExtension, +}; use crate::query_engine::DefaultSerializer; use crate::query_engine::options::QueryOptions; use crate::query_engine::runtime::{ @@ -310,7 +314,84 @@ impl QueryEngineState { rules.retain(|rule| rule.name() != name); } - /// Optimize the logical plan by the extension analyzer rules. + /// Builds a per-query DataFusion context with validated option overrides. + pub fn query_engine_context( + self: &Arc, + query_ctx: QueryContextRef, + ) -> crate::error::Result { + let query_ctx = Arc::new(query_ctx.query_option_snapshot().map_err(|error| { + InvalidQueryContextExtensionSnafu { + reason: error.to_string(), + } + .build() + })?); + let mut state = self.session_state(); + state.config_mut().set_extension(query_ctx.clone()); + state.config_mut().set_extension(self.clone()); + let options = query_ctx.extensions(); + for (key, value) in &options { + if key.starts_with("datafusion.") { + state + .config_mut() + .options_mut() + .set(key, value) + .map_err(|error| { + InvalidQueryContextExtensionSnafu { + reason: error.to_string(), + } + .build() + })?; + } + } + if let Some(partitions) = options.get("query.parallelism") { + let partitions = partitions.parse::().map_err(|error| { + InvalidQueryContextExtensionSnafu { + reason: error.to_string(), + } + .build() + })?; + *state.config_mut() = state.config().clone().with_target_partitions(partitions); + } + if let Some(fallback) = options.get("query.allow_query_fallback") { + state + .config_mut() + .options_mut() + .extensions + .insert(DistPlannerOptions { + allow_query_fallback: fallback == "true", + }); + } + state.config_mut().options_mut().execution.time_zone = + Some(query_ctx.timezone().to_string()); + state + .config_mut() + .options_mut() + .extensions + .insert(FunctionContext { + query_ctx: query_ctx.clone(), + state: self.function_state(), + }); + state + .config_mut() + .options_mut() + .extensions + .insert(ScheduledTimeExtension { + scheduled_time: crate::options::scheduled_time_from_ctx(&query_ctx), + }); + let config_options = state.config_options().clone(); + let _ = state + .execution_props_mut() + .config_options + .insert(config_options); + if let Some(scheduled_rt) = + crate::options::parse_scheduled_time_datetime(&query_ctx.extensions())? + { + state.execution_props_mut().query_execution_start_time = Some(scheduled_rt); + } + Ok(QueryEngineContext::new(state, query_ctx)) + } + + /// Optimize the logical plan using registered extension analyzer rules. pub fn optimize_by_extension_rules( &self, plan: DfLogicalPlan, @@ -319,7 +400,7 @@ impl QueryEngineState { self.extension_rules .iter() .try_fold(plan, |acc_plan, rule| { - rule.analyze(acc_plan, context, self.session_state().config_options()) + rule.analyze(acc_plan, context, context.state().config_options()) }) } @@ -695,10 +776,11 @@ mod tests { use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryLimit as DfMemoryLimit}; use datafusion::execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; use datafusion_common::config::SpillCompression; - use session::context::QueryContext; + use session::context::{QueryContext, QueryContextBuilder}; use super::*; use crate::options::{QueryOptions, QuerySpillCompression, QuerySpillMode}; + use crate::planner::LogicalPlanner; use crate::query_engine::runtime::{ DefaultQueryRuntimeProvider, QueryRuntimeContext, QueryRuntimeProvider, QueryRuntimeProviderRef, @@ -722,6 +804,53 @@ mod tests { ) } + struct RecordingAnalyzer(Arc>>); + + impl ExtensionAnalyzerRule for RecordingAnalyzer { + fn analyze( + &self, + plan: DfLogicalPlan, + _ctx: &QueryEngineContext, + config: &datafusion_common::config::ConfigOptions, + ) -> DfResult { + *self.0.lock().unwrap() = Some(( + config.optimizer.prefer_hash_join, + config.execution.target_partitions, + )); + Ok(plan) + } + } + + #[tokio::test] + async fn extension_analyzer_receives_query_specific_options() { + let observed = Arc::new(std::sync::Mutex::new(None)); + let mut state = new_query_engine_state(); + state + .extension_rules + .push(Arc::new(RecordingAnalyzer(observed.clone()))); + let state = Arc::new(state); + + let query_ctx = QueryContextBuilder::default().build(); + query_ctx + .configuration_parameter() + .set_query_option("datafusion.optimizer.prefer_hash_join", "false") + .unwrap(); + query_ctx + .configuration_parameter() + .set_query_option("query.parallelism", "2") + .unwrap(); + let query_ctx = Arc::new(query_ctx); + let statement = + crate::parser::QueryLanguageParser::parse_sql("SELECT 1", &query_ctx).unwrap(); + + crate::planner::DfLogicalPlanner::new(state) + .plan(&statement, query_ctx) + .await + .unwrap(); + + assert_eq!(*observed.lock().unwrap(), Some((false, 2))); + } + struct TestRuntimeProvider { build_called: AtomicBool, configure_called: AtomicBool, diff --git a/src/query/src/sql.rs b/src/query/src/sql.rs index 7527d44d821..0c6e278978b 100644 --- a/src/query/src/sql.rs +++ b/src/query/src/sql.rs @@ -376,7 +376,7 @@ async fn query_from_information_schema_dataframe( let view = dataframe.into_view(); let dataframe = SessionContext::new_with_state( query_engine - .engine_context(query_ctx.clone()) + .engine_context(query_ctx.clone())? .state() .clone(), ) @@ -844,6 +844,72 @@ pub async fn show_charsets_dataframe( .await } +/// Shows one supported query optimizer option, returning `None` for legacy SHOW handling. +pub fn show_query_option( + stmt: ShowVariables, + query_ctx: QueryContextRef, + engine: &dyn crate::QueryEngine, +) -> Result> { + let name = stmt.variable.to_string(); + let Some(canonical) = session::query_options::canonical_query_option_name(&name) else { + if name.to_ascii_lowercase().starts_with("query.") + || name.to_ascii_lowercase().starts_with("datafusion.") + { + session::query_options::parse_query_option(&name, "").map_err(|error| { + error::InvalidQueryContextExtensionSnafu { + reason: error.to_string(), + } + .build() + })?; + } + return Ok(None); + }; + let context = engine.engine_context(query_ctx.clone())?; + let value = match canonical { + "query.allow_query_fallback" => context + .state() + .config_options() + .extensions + .get::() + .map(|options| options.allow_query_fallback) + .unwrap_or(false) + .to_string(), + "query.enable_remote_dynamic_filter_pushdown" => { + crate::options::remote_dyn_filter_pushdown_enabled_from_extensions( + &context.query_ctx().extensions(), + )? + .to_string() + } + "query.parallelism" => context.state().config().target_partitions().to_string(), + _ => context + .state() + .config_options() + .entries() + .into_iter() + .find(|entry| entry.key == canonical) + .and_then(|entry| entry.value) + .ok_or_else(|| { + error::InvalidQueryContextExtensionSnafu { + reason: format!( + "No DataFusion configuration entry for supported option {canonical}" + ), + } + .build() + })?, + }; + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + name.clone(), + ConcreteDataType::string_datatype(), + false, + )])); + let records = RecordBatches::try_from_columns( + schema, + vec![Arc::new(StringVector::from(vec![value])) as _], + ) + .context(error::CreateRecordBatchSnafu)?; + Ok(Some(Output::new_with_record_batches(records))) +} + pub fn show_variable(stmt: ShowVariables, query_ctx: QueryContextRef) -> Result { let variable = stmt.variable.to_string().to_uppercase(); let value = match variable.as_str() {