From 340ea2d71e1d787fc5aabf4c690c4432aadc38ee Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 16:26:19 +0800 Subject: [PATCH] fix(query): retain metric table restriction in native scans Signed-off-by: discord9 --- src/query/src/promql/planner/native.rs | 111 ++++++++++++++++++------- src/query/src/promql/planner/test.rs | 91 ++++++++++++++++---- 2 files changed, 155 insertions(+), 47 deletions(-) diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs index 14fa4d99734..fbca18f39d4 100644 --- a/src/query/src/promql/planner/native.rs +++ b/src/query/src/promql/planner/native.rs @@ -30,11 +30,12 @@ use std::sync::Arc; use std::time::{Duration, UNIX_EPOCH}; -use common_catalog::consts::{MITO_ENGINE, MITO2_ENGINE}; +use common_catalog::consts::{METRIC_ENGINE, MITO_ENGINE, MITO2_ENGINE}; use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, TimeUnit as ArrowTimeUnit}; use datafusion::common::{DFSchema, DFSchemaRef}; use datafusion::error::DataFusionError; use datafusion::logical_expr::{Extension, LogicalPlan}; +use datafusion_expr::Expr; use datatypes::series::SeriesOrdering; use datatypes::series::samples::{FieldSpec, SeriesSchema}; use promql::extension_plan::{FusedSeries, FusedSeriesBinding}; @@ -45,6 +46,9 @@ use promql_parser::parser::{ }; use session::context::QueryContextRef; use snafu::ResultExt; +use store_api::metric_engine_consts::{ + DATA_SCHEMA_TABLE_ID_COLUMN_NAME, LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, +}; use crate::promql::error::{DataFusionPlanningSnafu, Result}; use crate::promql::native_compile::{compile, statement_grid}; @@ -62,7 +66,7 @@ use crate::promql::planner::PromPlanner; /// [`QueryOptions::experimental_native_promql`]: crate::options::QueryOptions::experimental_native_promql /// [`remote_dyn_filter_pushdown_enabled_from_extensions`]: /// crate::options::remote_dyn_filter_pushdown_enabled_from_extensions -pub(crate) const NATIVE_PROMQL_EXTENSION_KEY: &str = "query.experimental_native_promql"; +pub const NATIVE_PROMQL_EXTENSION_KEY: &str = "query.experimental_native_promql"; /// Whether `query_ctx` asks for the native fused-series route. /// @@ -160,12 +164,41 @@ impl PromPlanner { // instead would let the two paths disagree about the columns they read. let table_ref = self.table_ref()?; let mut input = self.create_table_scan_plan(table_ref.clone()).await?; - // The node reads the region scan itself, so the plan below it has to *be* that scan: a - // filter, a projection or a `__tsid` rewrite in between produces a different physical - // child, and the physical layer cannot reach through it. Declining here, while the logical - // plan is still this method's to choose, is what keeps such a query on the existing path. - let LogicalPlan::TableScan(scan) = &input else { - return Ok(None); + let logical_scan = metric_scan_and_table_id(&input); + let has_verified_metric_rewrite = logical_scan.is_some(); + let scan = match logical_scan { + Some((mut scan, predicate, table_id)) => { + let logical_provider = self + .table_provider + .resolve_table(table_ref.clone()) + .await + .context(crate::promql::error::CatalogSnafu)?; + let logical_table = self.table_from_source(&logical_provider)?; + let meta = &logical_table.table_info().meta; + if meta.engine != METRIC_ENGINE_NAME + || !meta + .options + .extra_options + .contains_key(LOGICAL_TABLE_METADATA_KEY) + || logical_table.table_info().ident.table_id != table_id + { + return Ok(None); + } + // The physical scan bypasses logical-region request transformation, so preserve + // the validated logical-table restriction as a pushed-down scan predicate. + scan.filters.push(predicate); + input = LogicalPlan::TableScan(scan.clone()); + scan + } + None => { + if self.ctx.use_tsid { + return Ok(None); + } + match &input { + LogicalPlan::TableScan(scan) => scan.clone(), + _ => return Ok(None), + } + } }; let table = self.table_from_source(&scan.source)?; let table_meta = &table.table_info().meta; @@ -174,7 +207,9 @@ impl PromPlanner { // planning error, not a fallback. So the engine is checked here, where the existing plan is // still the answer. A file-engine table is the concrete case: its scanner is a // `SinglePartitionScanner`, which leaves `raw_series_schema` at the default `None`. - if !supports_native_raw_series(&table_meta.engine) { + if !supports_native_raw_series(&table_meta.engine) + || (table_meta.engine == METRIC_ENGINE_NAME && !has_verified_metric_rewrite) + { return Ok(None); } // The region's raw-series preflight refuses these, and the physical layer turns that refusal into an error, so decline this route here instead. @@ -379,29 +414,45 @@ fn selector_range_millis(expr: &PromExpr) -> Option { } } -/// Whether `engine` can serve raw series frames to the node this route builds. +/// Whether `engine` can serve raw series frames to the fused-series node. /// -/// A whitelist, not a probe: the logical planner cannot ask a region what its scanner supports -/// without reading the table, and a false negative only costs the query its native route. -/// -/// Mito's series scanner implements the raw path. -/// -/// The metric engine is deliberately absent. A logical metric table does read through its physical -/// region, which is a mito scanner, but the planner reaches this node with a bare scan only when it -/// could not resolve that physical table: the ordinary rewrite puts a `__table_id` filter and an -/// alias between the scan and this node, which the bare-scan check above already refuses. A logical -/// table whose physical table is missing is not a table this route can serve - its `TableInfo` -/// carries no `append_mode` either, so the preflight's own check could not be made here - and -/// naming the engine would admit exactly that unverifiable case. -/// -/// Every other engine in the workspace leaves `raw_series_schema` at its default `None`: a -/// file-engine table, for one, is read by a `SinglePartitionScanner` that cannot frame a series. -/// -/// This is deliberately not [`supports_pk_dictionary_encoding`]: that answers a different question. -/// -/// [`supports_pk_dictionary_encoding`]: table::table::scan +/// A whitelist avoids probing a region from logical planning; declining only costs the native route. fn supports_native_raw_series(engine: &str) -> bool { - matches!(engine, MITO_ENGINE | MITO2_ENGINE) + matches!(engine, MITO_ENGINE | MITO2_ENGINE | METRIC_ENGINE) +} + +/// Recognizes only the metric logical-table scan emitted by `create_table_scan_plan`. +fn metric_scan_and_table_id( + plan: &LogicalPlan, +) -> Option<(datafusion::logical_expr::TableScan, Expr, u32)> { + let plan = match plan { + LogicalPlan::Projection(projection) => projection.input.as_ref(), + plan => plan, + }; + let plan = match plan { + LogicalPlan::SubqueryAlias(alias) => alias.input.as_ref(), + plan => plan, + }; + let LogicalPlan::Filter(filter) = plan else { + return None; + }; + let Expr::BinaryExpr(binary) = &filter.predicate else { + return None; + }; + if binary.op != datafusion::logical_expr::Operator::Eq + || !matches!(binary.left.as_ref(), Expr::Column(column) if column.name == DATA_SCHEMA_TABLE_ID_COLUMN_NAME) + { + return None; + } + let Expr::Literal(datafusion::scalar::ScalarValue::UInt32(Some(table_id)), _) = + binary.right.as_ref() + else { + return None; + }; + let LogicalPlan::TableScan(scan) = filter.input.as_ref() else { + return None; + }; + Some((scan.clone(), filter.predicate.clone(), *table_id)) } fn single_selector(expr: &PromExpr) -> Option<&VectorSelector> { diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 15732c987bb..b085e3b4f2f 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -38,7 +38,7 @@ use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; use datafusion::datasource::{MemTable, provider_as_source}; use datafusion::execution::context::SessionContext; -use datafusion::logical_expr::{Extension, col}; +use datafusion::logical_expr::{Extension, UserDefinedLogicalNodeCore, col}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; use promql::extension_plan::{FusedSeries, HistogramFold}; @@ -1748,6 +1748,22 @@ async fn build_native_fused_plan_with_provider( .unwrap() } +fn find_metric_physical_scan(plan: &LogicalPlan) -> Option { + match plan { + LogicalPlan::Extension(Extension { node }) => node + .as_any() + .downcast_ref::() + .and_then(|fused| match fused.inputs().first() { + Some(LogicalPlan::TableScan(scan)) => Some(scan.clone()), + _ => None, + }), + _ => plan + .inputs() + .into_iter() + .find_map(find_metric_physical_scan), + } +} + fn find_fused_series(plan: &LogicalPlan) -> Option<&FusedSeries> { if let LogicalPlan::Extension(Extension { node }) = plan && let Some(fused) = node.as_any().downcast_ref::() @@ -1881,16 +1897,61 @@ async fn native_fused_tag_partitioned_table_still_plans_fused_series() { assert!(plan.display_indent().to_string().contains("FusedSeries")); } -/// The flag on and a shape the compiler refuses: the existing planner plans it unchanged. -/// A metric-engine logical table keeps the existing plan. -/// -/// A logical table reads through its physical region, which is a mito scanner, but the planner -/// only reaches this node with a bare scan when that physical table could not be resolved: the -/// ordinary rewrite puts a `__table_id` filter and an alias in between. Such a table's `TableInfo` -/// also carries no `append_mode`, so the preflight's own check could not be made here. Both are -/// reasons the engine is not on the route's whitelist, and this test pins that it falls back. +/// Metric logical scans retain their table-id restriction on the direct fused scan input. #[tokio::test] -async fn native_fused_metric_engine_table_falls_back() { +async fn native_fused_metric_engine_table_is_accepted() { + let native = build_native_fused_plan_with_provider( + "rate(some_metric[5m])", + true, + build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await, + ) + .await; + let legacy = build_native_fused_plan_with_provider( + "rate(some_metric[5m])", + false, + build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await, + ) + .await; + assert!( + find_fused_series(&native).is_some(), + "{}", + native.display_indent() + ); + assert!( + legacy + .display_indent() + .to_string() + .contains("PromRangeManipulate") + ); + let scan = find_metric_physical_scan(&native).expect("fused route keeps physical scan"); + assert!(scan.filters.iter().any(|filter| { + *filter + == datafusion::logical_expr::Expr::Column(datafusion_common::Column::new( + Some(datafusion_common::TableReference::bare("phy")), + DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(), + )) + .eq(datafusion::logical_expr::lit(1024u32)) + })); + assert!( + scan.filters + .iter() + .any(|filter| format!("{filter:?}").contains("timestamp")) + ); +} + +/// A metric table without logical-table metadata declines to the legacy plan. +#[tokio::test] +async fn native_fused_metric_table_without_logical_metadata_falls_back() { let native = build_native_fused_plan_with_provider( "rate(some_metric[5m])", true, @@ -1907,7 +1968,7 @@ async fn native_fused_metric_engine_table_falls_back() { .await, ) .await; - let without_flag = build_native_fused_plan_with_provider( + let legacy = build_native_fused_plan_with_provider( "rate(some_metric[5m])", false, build_test_table_provider_with_options( @@ -1923,14 +1984,10 @@ async fn native_fused_metric_engine_table_falls_back() { .await, ) .await; - assert!( - find_fused_series(&native).is_none(), - "a metric-engine table must not reach the native route" - ); + assert!(find_fused_series(&native).is_none()); assert_eq!( native.display_indent().to_string(), - without_flag.display_indent().to_string(), - "a metric-engine table must keep the existing plan" + legacy.display_indent().to_string() ); }