rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-12 23:12:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: apply effective query options throughout planning and execution

Browse Source 
Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
This commit is contained in:
discord9
2026-10-09 16:35:10 +08:00
parent 9d9039e07b
commit 17af970ddd
12 changed files with 471 additions and 169 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/datanode/src/region_server.rs
+7 -6
View File
@@ -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)
src/datanode/src/tests.rs
+4 -1
View File
@@ -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<QueryEngineContext> {
unimplemented!()
}
fn engine_state(&self) -> &QueryEngineState {
src/frontend/src/instance.rs
+15
View File
@@ -292,6 +292,21 @@ impl Instance {
async fn query_statement(&self, stmt: Statement, query_ctx: QueryContextRef) -> Result<Output> {
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::<SqlQueryInterceptorRef<Error>>();
let query_interceptor = query_interceptor.as_ref();
src/frontend/src/instance/grpc.rs
+12 -4
View File
@@ -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
src/operator/src/statement/set.rs
+7 -9
View File
@@ -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<Expr>, 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<Expr>, 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<Expr>, 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)
src/query/src/datafusion.rs
+204 -100
View File
@@ -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<Output> {
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::<RemoteDynFilterReceiverInjectorRef>()
{
@@ -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::<u64>() {
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::<bool>().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<QueryEngineContext> {
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::<DistPlannerOptions>()
.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::<DistPlannerOptions>()
.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::<DistPlannerOptions>()
.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::<DatafusionQueryEngine>()
.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::<DatafusionQueryEngine>()
.unwrap();
let engine_ctx = engine.engine_context(QueryContext::arc());
let engine_ctx = engine.engine_context(QueryContext::arc()).unwrap();
let state = engine_ctx.state();
assert!(
src/query/src/optimizer/constant_term.rs
+1 -1
View File
@@ -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
src/query/src/planner.rs
+18 -42
View File
@@ -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<SessionState> {
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<DfExpr> {
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<LogicalPlan> {
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<LogicalPlan> {
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()
src/query/src/query_engine.rs
+1 -1
View File
@@ -97,7 +97,7 @@ pub trait QueryEngine: Send + Sync {
fn read_table(&self, table: TableRef) -> Result<DataFrame>;
/// Create a [`QueryEngineContext`].
fn engine_context(&self, query_ctx: QueryContextRef) -> QueryEngineContext;
fn engine_context(&self, query_ctx: QueryContextRef) -> Result<QueryEngineContext>;
/// Retrieve the query engine state [`QueryEngineState`]
fn engine_state(&self) -> &QueryEngineState;
src/query/src/query_engine/default_serializer.rs
+2
View File
@@ -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();
src/query/src/query_engine/state.rs
+133 -4
View File
@@ -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<Self>,
query_ctx: QueryContextRef,
) -> crate::error::Result<QueryEngineContext> {
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::<usize>().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<std::sync::Mutex<Option<(bool, usize)>>>);
impl ExtensionAnalyzerRule for RecordingAnalyzer {
fn analyze(
&self,
plan: DfLogicalPlan,
_ctx: &QueryEngineContext,
config: &datafusion_common::config::ConfigOptions,
) -> DfResult<DfLogicalPlan> {
*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,
src/query/src/sql.rs
+67 -1
View File
@@ -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<Option<Output>> {
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::<crate::dist_plan::DistPlannerOptions>()
.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<Output> {
let variable = stmt.variable.to_string().to_uppercase();
let value = match variable.as_str() {