From 00aa65de7a9a3e72bc8674cf3a54b7f668809edf Mon Sep 17 00:00:00 2001 From: Yingwen Date: Wed, 14 Oct 2026 06:49:35 +0000 Subject: [PATCH] fix(query): guard deep expressions during query planning (#9564) * fix(query): avoid speculative clones in constant normalization Signed-off-by: evenyag * fix(query): reject deeply nested expressions before optimization Signed-off-by: evenyag --------- Signed-off-by: evenyag --- src/promql/src/extension_plan/empty_metric.rs | 5 + src/query/src/datafusion.rs | 5 + src/query/src/dist_plan.rs | 2 +- src/query/src/dist_plan/analyzer.rs | 3 + src/query/src/dist_plan/planner.rs | 2 + src/query/src/optimizer.rs | 1 + .../src/optimizer/const_normalization.rs | 189 +++++++-------- src/query/src/optimizer/expr_depth.rs | 223 ++++++++++++++++++ src/query/src/planner.rs | 2 + src/query/src/query_engine/state.rs | 9 +- tests-integration/tests/sql.rs | 51 ++++ .../common/filter/expression_depth.result | 24 ++ .../common/filter/expression_depth.sql | 9 + .../common/tql-explain-analyze/explain.result | 3 + 14 files changed, 433 insertions(+), 95 deletions(-) create mode 100644 src/query/src/optimizer/expr_depth.rs create mode 100644 tests/cases/standalone/common/filter/expression_depth.result create mode 100644 tests/cases/standalone/common/filter/expression_depth.sql diff --git a/src/promql/src/extension_plan/empty_metric.rs b/src/promql/src/extension_plan/empty_metric.rs index 4d2e029417a..a534bf20493 100644 --- a/src/promql/src/extension_plan/empty_metric.rs +++ b/src/promql/src/extension_plan/empty_metric.rs @@ -109,6 +109,11 @@ impl EmptyMetric { }) } + /// Returns the expression used to generate metric values. + pub fn expr(&self) -> Option<&Expr> { + self.expr.as_ref() + } + pub const fn name() -> &'static str { "EmptyMetric" } diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index 970f3109804..72788461836 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -68,6 +68,7 @@ use crate::metrics::{ OnDone, QUERY_STAGE_ELAPSED, maybe_attach_region_watermark_metrics, should_collect_region_watermark_from_query_ctx, }; +use crate::optimizer::expr_depth; use crate::options::ScheduledTimeExtension; use crate::physical_wrapper::PhysicalPlanWrapperRef; use crate::planner::{DfLogicalPlanner, LogicalPlanner}; @@ -381,6 +382,8 @@ impl DatafusionQueryEngine { } } + expr_depth::validate_plan(logical_plan)?; + let mut logger = PanicLogger { input_logical_plan: logical_plan, after_analyze: None, @@ -408,6 +411,7 @@ impl DatafusionQueryEngine { |_, _| {}, )?; + expr_depth::validate_plan(&analyzed_plan)?; logger.after_analyze = Some(analyzed_plan.clone()); common_telemetry::debug!("Create physical plan, analyzed plan: {analyzed_plan}"); @@ -423,6 +427,7 @@ impl DatafusionQueryEngine { .optimize(analyzed_plan, state, |_, _| {})? }; + expr_depth::validate_plan(&optimized_plan)?; common_telemetry::debug!("Create physical plan, optimized plan: {optimized_plan}"); logger.after_optimize = Some(optimized_plan.clone()); diff --git a/src/query/src/dist_plan.rs b/src/query/src/dist_plan.rs index 969f00e7c9c..a982eb1d8bf 100644 --- a/src/query/src/dist_plan.rs +++ b/src/query/src/dist_plan.rs @@ -27,7 +27,7 @@ mod remote_dyn_filter_registry; pub use analyzer::{DistPlannerAnalyzer, DistPlannerOptions}; pub use filter_id::{FilterFingerprint, FilterId, ParseFilterIdError, RemoteDynFilterProducerId}; pub use merge_scan::{MergeScanExec, MergeScanLogicalPlan}; -pub(crate) use merge_sort::MergeSortExec; +pub(crate) use merge_sort::{MergeSortExec, MergeSortLogicalPlan}; pub use planner::{DistExtensionPlanner, MergeSortExtensionPlanner}; pub use predicate_extractor::PredicateExtractor; pub use region_pruner::ConstraintPruner; diff --git a/src/query/src/dist_plan/analyzer.rs b/src/query/src/dist_plan/analyzer.rs index 38d386cc818..a39514a0e43 100644 --- a/src/query/src/dist_plan/analyzer.rs +++ b/src/query/src/dist_plan/analyzer.rs @@ -53,6 +53,7 @@ use crate::dist_plan::commutativity::{ use crate::dist_plan::merge_scan::MergeScanLogicalPlan; use crate::dist_plan::merge_sort::MergeSortLogicalPlan; use crate::metrics::PUSH_DOWN_FALLBACK_ERRORS_TOTAL; +use crate::optimizer::expr_depth; use crate::options::ScheduledTimeExtension; use crate::plan::ExtractExpr; use crate::query_engine::DefaultSerializer; @@ -118,6 +119,7 @@ impl AnalyzerRule for DistPlannerAnalyzer { plan: LogicalPlan, config: &ConfigOptions, ) -> datafusion_common::Result { + expr_depth::validate_plan(&plan)?; let mut config = config.clone(); // Aligned with the behavior in `datafusion_optimizer::OptimizerContext::new()`. config.optimizer.filter_null_join_keys = true; @@ -142,6 +144,7 @@ impl AnalyzerRule for DistPlannerAnalyzer { let plan = plan .rewrite_with_subqueries(&mut PlanTreeExpressionSimplifier::new(optimizer_context))? .data; + expr_depth::validate_plan(&plan)?; let fallback_plan = plan.clone(); // Run a filter-focused optimizer subset before MergeScan wraps remote diff --git a/src/query/src/dist_plan/planner.rs b/src/query/src/dist_plan/planner.rs index 88b13caf388..34c0994f9e0 100644 --- a/src/query/src/dist_plan/planner.rs +++ b/src/query/src/dist_plan/planner.rs @@ -50,6 +50,7 @@ use crate::dist_plan::merge_scan::{MergeScanExec, MergeScanLogicalPlan}; use crate::dist_plan::merge_sort::{MergeSortExec, MergeSortLogicalPlan}; use crate::dist_plan::region_pruner::ConstraintPruner; use crate::error::{CatalogSnafu, PartitionRuleManagerSnafu, TableNotFoundSnafu}; +use crate::optimizer::expr_depth; use crate::region_query::RegionQueryHandlerRef; /// Planner for converting merge sort logical plan to physical plan. @@ -447,6 +448,7 @@ impl DistExtensionPlanner { session: &dyn Session, plan: &LogicalPlan, ) -> Result { + expr_depth::validate_plan(plan)?; let session_state = session .as_any() .downcast_ref::() diff --git a/src/query/src/optimizer.rs b/src/query/src/optimizer.rs index 03d6ce3aa46..35a3cac2639 100644 --- a/src/query/src/optimizer.rs +++ b/src/query/src/optimizer.rs @@ -17,6 +17,7 @@ pub mod constant_term; pub mod count_nest_aggr; pub mod count_wildcard; pub mod enforce_sorting; +pub(crate) mod expr_depth; pub mod global_limit; pub(crate) mod insert_assignment; pub(crate) mod json_schema_concretize; diff --git a/src/query/src/optimizer/const_normalization.rs b/src/query/src/optimizer/const_normalization.rs index d29f0aa4f39..64d39a5f668 100644 --- a/src/query/src/optimizer/const_normalization.rs +++ b/src/query/src/optimizer/const_normalization.rs @@ -23,8 +23,6 @@ use datafusion_expr::{Between, BinaryExpr, Expr, ExprSchemable, LogicalPlan, Ope use datafusion_expr_common::casts::try_cast_literal_to_type; use datafusion_optimizer::analyzer::AnalyzerRule; -use crate::plan::ExtractExpr; - /// ConstNormalizationRule rewrites castable constants against their /// non-constant comparison operand ahead of filter pushdown. #[derive(Debug)] @@ -52,26 +50,12 @@ impl AnalyzerRule for ConstNormalizationRule { } fn rewrite_plan_exprs(plan: LogicalPlan, schema: DFSchemaRef) -> Result> { - let mut rewriter = ConstNormalizationRewriter { - schema, - transformed: false, - }; - let exprs = plan - .expressions_consider_join() - .into_iter() - .map(|expr| expr.rewrite(&mut rewriter).map(|rewritten| rewritten.data)) - .collect::>>()?; - if !rewriter.transformed { - return Ok(Transformed::no(plan)); - } - - let inputs = plan.inputs().into_iter().cloned().collect::>(); - plan.with_new_exprs(exprs, inputs).map(Transformed::yes) + let mut rewriter = ConstNormalizationRewriter { schema }; + plan.map_expressions(|expr| expr.rewrite(&mut rewriter)) } struct ConstNormalizationRewriter { schema: DFSchemaRef, - transformed: bool, } impl TreeNodeRewriter for ConstNormalizationRewriter { @@ -91,44 +75,36 @@ impl TreeNodeRewriter for ConstNormalizationRewriter { } fn f_up(&mut self, expr: Expr) -> Result> { - let rewritten = rewrite_expr_node(expr, &self.schema)?; - self.transformed |= rewritten.transformed; - Ok(rewritten) + rewrite_expr_node(expr, &self.schema) } } fn rewrite_expr_node(expr: Expr, schema: &DFSchemaRef) -> Result> { - match expr { - Expr::BinaryExpr(binary) => match rewrite_binary_expr(binary.clone(), schema)? { - Some(expr) => Ok(Transformed::yes(expr)), - None => Ok(Transformed::no(Expr::BinaryExpr(binary))), - }, - Expr::Between(between) => match rewrite_between_expr(between.clone(), schema)? { - Some(expr) => Ok(Transformed::yes(expr)), - None => Ok(Transformed::no(Expr::Between(between))), - }, - Expr::InList(in_list) => match rewrite_in_list_expr(in_list.clone(), schema)? { - Some(expr) => Ok(Transformed::yes(expr)), - None => Ok(Transformed::no(Expr::InList(in_list))), - }, - Expr::Like(like) => rewrite_like_expr(like, PatternMatchKind::Like, schema), - Expr::SimilarTo(like) => rewrite_like_expr(like, PatternMatchKind::SimilarTo, schema), - expr => Ok(Transformed::no(expr)), - } + // Borrow while checking eligibility: cloning unchanged subtrees here can overflow + // the stack even though DataFusion protects the surrounding rewrite traversal. + let rewritten = match &expr { + Expr::BinaryExpr(binary) => rewrite_binary_expr(binary, schema)?, + Expr::Between(between) => rewrite_between_expr(between, schema)?, + Expr::InList(in_list) => rewrite_in_list_expr(in_list, schema)?, + Expr::Like(like) => rewrite_like_expr(like, PatternMatchKind::Like, schema)?, + Expr::SimilarTo(like) => rewrite_like_expr(like, PatternMatchKind::SimilarTo, schema)?, + _ => None, + }; + Ok(match rewritten { + Some(rewritten) => Transformed::yes(rewritten), + None => Transformed::no(expr), + }) } -fn rewrite_between_expr(between: Between, schema: &DFSchemaRef) -> Result> { +fn rewrite_between_expr(between: &Between, schema: &DFSchemaRef) -> Result> { let Between { expr, negated, low, high, } = between; - let expr = *expr; - let low_expr = *low; - let high_expr = *high; let Some((target, constants)) = - extract_rewrite_operands(&expr, &[low_expr.clone(), high_expr.clone()], schema)? + extract_rewrite_operands(expr, [low.as_ref(), high.as_ref()], schema)? else { return Ok(None); }; @@ -142,7 +118,7 @@ fn rewrite_between_expr(between: Between, schema: &DFSchemaRef) -> Result