diff --git a/src/query/src/dist_plan/commutativity.rs b/src/query/src/dist_plan/commutativity.rs index d02b746a8bf..4d6033729e0 100644 --- a/src/query/src/dist_plan/commutativity.rs +++ b/src/query/src/dist_plan/commutativity.rs @@ -21,8 +21,9 @@ use datafusion::error::Result as DfResult; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_expr::{Expr, LogicalPlan, UserDefinedLogicalNode}; use promql::extension_plan::{ - EmptyMetric, InstantManipulate, RangeManipulate, SeriesDivide, SeriesNormalize, + EmptyMetric, FusedSeries, InstantManipulate, RangeManipulate, SeriesDivide, SeriesNormalize, }; +use promql::native::plan::{NativeExpr, NativeNode}; use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME; use crate::dist_plan::MergeScanLogicalPlan; @@ -240,6 +241,29 @@ impl Categorizer { Commutativity::Commutative } + name if name == FusedSeries::name() => { + let fused = plan.as_any().downcast_ref::().unwrap(); + // A region builds this node from the region scan itself: the conversion takes that + // scan's raw-series entrypoint and cannot look through anything in between, so a + // node whose input is not that scan has nothing to read where it would be placed + // and stays where it was planned. + if !matches!(fused.inputs().as_slice(), [LogicalPlan::TableScan(_)]) { + return Commutativity::NonCommutative; + } + // The region asks the scan for the series of one named metric, so a node that binds + // no name cannot be built there at all. + if fused.binding().metric_name.is_none() { + return Commutativity::NonCommutative; + } + // A region holds a subset of the metric's series and no series is split across + // regions, so the node may be evaluated there only when every node of its tree + // reads one series at a time. + if reads_across_series(fused.tree().root()) { + return Commutativity::NonCommutative; + } + + Commutativity::Commutative + } name if name == SeriesNormalize::name() || name == InstantManipulate::name() || name == RangeManipulate::name() => @@ -346,19 +370,60 @@ impl Categorizer { } } +/// Whether any node of `tree` reads more than one series at a time. +/// +/// Everything the native compiler emits today reads one series: an instant selection, a window and +/// a point-wise map each select, fold or map the samples of the series they were handed. An +/// aggregate folds a label set's series together and a binary operation pairs them, so both consume +/// several series at once. Neither has a compiler producer yet, which is why refusing them costs +/// nothing today: a tree built some other way is exactly what this check is for, because evaluating +/// it one region at a time would fold a partial input and answer a wrong query silently. +fn reads_across_series(node: &NativeNode) -> bool { + match &node.expr { + NativeExpr::Input(_) | NativeExpr::Scalar(_) => false, + NativeExpr::Instant { input, .. } + | NativeExpr::Window { input, .. } + | NativeExpr::Map { input, .. } => reads_across_series(input), + NativeExpr::Aggregate { .. } | NativeExpr::Binary { .. } => true, + } +} + #[cfg(test)] mod tests { use std::collections::{BTreeMap, BTreeSet}; + use std::time::{Duration, UNIX_EPOCH}; + use datafusion::arrow::datatypes::{ + DataType as ArrowDataType, Field, TimeUnit as ArrowTimeUnit, + }; + use datafusion::common::{DFSchema, DFSchemaRef}; + use datafusion::datasource::DefaultTableSource; use datafusion_common::Column; + use datafusion_common::config::ConfigOptions; use datafusion_expr::LogicalPlanBuilder; use datafusion_expr::expr::ScalarFunction; + use datafusion_expr::{Extension, LogicalPlan}; use datafusion_functions::core::coalesce; + use datafusion_optimizer::analyzer::AnalyzerRule; + use datatypes::series::SeriesOrdering; + use datatypes::series::samples::{FieldSpec, SeriesSchema}; + use promql::extension_plan::{FusedSeries, FusedSeriesBinding}; use promql::functions::{ NativeHistogramDrop, NativeHistogramFraction, NativeHistogramQuantile, }; + use promql::native::plan::{ + AggregateKind, AggregateStage, AggregateStateSchema, BoundAggregate, BoundInput, Grouping, + NativeType, + }; + use promql_parser::parser; + use store_api::storage::RegionId; + use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; use super::*; + use crate::dist_plan::DistPlannerAnalyzer; + use crate::optimizer::test_util::mock_table_provider; + use crate::promql::native_compile::{compile, statement_grid}; + use crate::query_engine::DefaultSerializer; #[test] fn series_divide_by_tsid_is_commutative() { @@ -393,6 +458,237 @@ mod tests { )); } } + + /// The statement one `FusedSeries` test node is bound to. + fn native_statement() -> promql_parser::parser::EvalStmt { + promql_parser::parser::EvalStmt { + expr: parser::parse("some_metric").unwrap(), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + } + } + + /// A `FusedSeries` node over `input`, bound to the tree a bare selector compiles to. + /// + /// The tree is the simplest shape the compiler accepts, because what these tests change is + /// where the node may be placed, not what it computes. + fn fused_series_plan(input: LogicalPlan) -> LogicalPlan { + fused_series_plan_named(input, Some("some_metric")) + } + + /// The same node, with the metric name it binds chosen by the caller. + fn fused_series_plan_named(input: LogicalPlan, metric_name: Option<&str>) -> LogicalPlan { + let stmt = native_statement(); + let series_schema = Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { name: "v0".into() }]), + eval_grid: None, + }); + let tree = compile( + &stmt, + Arc::clone(&series_schema), + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + let binding = FusedSeriesBinding { + series_schema, + grid: statement_grid(&stmt).unwrap(), + metric_name: metric_name.map(Arc::from), + time_index: "ts".to_string(), + tags: Vec::new(), + }; + let output_schema: DFSchemaRef = Arc::new( + DFSchema::new_with_metadata( + vec![ + ( + None, + Arc::new(Field::new("v0", ArrowDataType::Float64, true)), + ), + ( + None, + Arc::new(Field::new( + "ts", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + )), + ), + ], + Default::default(), + ) + .unwrap(), + ); + + LogicalPlan::Extension(Extension { + node: Arc::new( + FusedSeries::try_new( + Arc::from("some_metric"), + tree, + binding, + input, + output_schema, + ) + .unwrap(), + ), + }) + } + + /// The region's own view of the metric's table: a scan whose provider is the region provider. + fn table_scan_plan() -> LogicalPlan { + let provider = mock_table_provider(RegionId::new(1, 1)); + LogicalPlanBuilder::scan( + "some_metric", + Arc::new(DefaultTableSource::new(Arc::new(provider))), + None, + ) + .unwrap() + .build() + .unwrap() + } + + fn holds_fused_series(plan: &LogicalPlan) -> bool { + if let LogicalPlan::Extension(extension) = plan + && extension.node.as_any().is::() + { + return true; + } + plan.inputs().into_iter().any(holds_fused_series) + } + + #[test] + fn fused_series_over_its_scan_is_commutative() { + let plan = fused_series_plan(table_scan_plan()); + + assert!(matches!( + Categorizer::check_plan(&plan, None).unwrap(), + Commutativity::Commutative + )); + } + + /// The region builds the node from the scan's own raw-series entrypoint, so an input that is + /// not that scan leaves the node with nothing to read where it would be placed. + #[test] + fn fused_series_over_anything_but_a_scan_is_not_commutative() { + let scan = table_scan_plan(); + let input = LogicalPlanBuilder::from(scan) + .filter(datafusion_expr::col("k0").eq(datafusion_expr::lit("a"))) + .unwrap() + .build() + .unwrap(); + let plan = fused_series_plan(input); + + assert!(matches!( + Categorizer::check_plan(&plan, None).unwrap(), + Commutativity::NonCommutative + )); + } + + /// The region asks the scan for one named metric's series, so a node binding no name cannot + /// be built there. + #[test] + fn fused_series_that_binds_no_metric_name_is_not_commutative() { + let plan = fused_series_plan_named(table_scan_plan(), None); + + assert!(matches!( + Categorizer::check_plan(&plan, None).unwrap(), + Commutativity::NonCommutative + )); + } + + /// A tree built by something other than the compiler can hold a node that folds several + /// series. The check is a property of the tree, so it is exercised on the node itself: the + /// compiler emits no such node today, and nothing else would build a tree containing one. + #[test] + fn a_tree_that_folds_across_series_is_not_per_series() { + let raw = NativeType::Samples(Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { name: "v0".into() }]), + eval_grid: None, + })); + let input = NativeNode { + expr: NativeExpr::Input(BoundInput { + output: raw.clone(), + ordering: SeriesOrdering::KeyThenTime, + }), + output: raw, + }; + let group = Grouping::By(Arc::from(["k0".into()])); + let aggregate = NativeNode { + expr: NativeExpr::Aggregate { + input: Box::new(input), + op: BoundAggregate { + kind: AggregateKind::Sum, + group: group.clone(), + stage: AggregateStage::Partial, + }, + }, + output: NativeType::AggregateState(Arc::new(AggregateStateSchema { + group, + accumulators: Arc::from([]), + grid: statement_grid(&native_statement()).unwrap(), + })), + }; + assert!(reads_across_series(&aggregate)); + + // The same walk over what the compiler emits stays per-series. + let single_selector = compile( + &native_statement(), + Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { name: "v0".into() }]), + eval_grid: None, + }), + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + assert!(!reads_across_series(single_selector.root())); + } + + /// Blocks placement, and pins the blocker rather than the intent. + /// + /// `PlanRewriter::should_expand` asks two questions about a node before it lets the descent + /// continue, and the Substrait one comes first: unless the *whole* root plan encodes, every + /// subtree that fails to encode is expanded on, whichever way `check_plan` answers. A + /// `FusedSeries` node has no serializer branch — `ExtensionSerializer` answers + /// `NotImplemented` for every name it does not list — so its subtree does not encode, the + /// boundary is drawn at the node, and the region receives the bare scan instead of the node. + /// + /// The commutativity arm above is necessary and not sufficient. This test fails the day the + /// node gains a serializer branch (or the analyzer stops asking), which is when the placement + /// assertion this one replaces can be written. + #[test] + fn fused_series_does_not_encode_so_the_analyzer_keeps_it_on_the_frontend() { + let plan = fused_series_plan(table_scan_plan()); + + assert!( + DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .is_err(), + "the fused node encodes now, so it can be placed inside the merge scan" + ); + + let analyzed = DistPlannerAnalyzer {} + .analyze(plan, &ConfigOptions::default()) + .unwrap(); + let mut remote_input_holds_fused = false; + analyzed + .apply(|node| { + if let LogicalPlan::Extension(extension) = node + && let Some(merge_scan) = extension + .node + .as_any() + .downcast_ref::() + { + remote_input_holds_fused = holds_fused_series(merge_scan.input()); + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + assert!( + !remote_input_holds_fused, + "the analyzer placed the node despite the encoding gate: {analyzed}" + ); + } } pub type Transformer = Arc Option>; diff --git a/src/query/src/optimizer/scan_hint.rs b/src/query/src/optimizer/scan_hint.rs index 4685889c998..0a1334dd9c0 100644 --- a/src/query/src/optimizer/scan_hint.rs +++ b/src/query/src/optimizer/scan_hint.rs @@ -26,7 +26,7 @@ use datafusion_expr::expr::Sort; use datafusion_expr::{Expr, LogicalPlan, utils}; use datafusion_optimizer::{OptimizerConfig, OptimizerRule}; use datatypes::arrow::datatypes::{DataType, TimeUnit as ArrowTimeUnit}; -use promql::extension_plan::{InstantManipulate, SeriesDivide, SeriesNormalize}; +use promql::extension_plan::{FusedSeries, InstantManipulate, SeriesDivide, SeriesNormalize}; use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME; use store_api::storage::{TimeSeriesDistribution, TimeSeriesRowSelector}; @@ -37,6 +37,9 @@ use crate::dummy_catalog::DummyTableProvider; /// - the nearest order requirement to the leaf table scan node as ordering hint. /// - the group by columns when all aggregate functions are `last_value` as /// time series row selector hint. +/// - the per-series distribution when the leaf table scan node is the input of a +/// [`FusedSeries`] node, which reads the raw series frames a region only packs +/// for that distribution. /// /// [`ScanRequest`]: store_api::storage::ScanRequest #[derive(Debug)] @@ -68,6 +71,7 @@ impl ScanHintRule { fn set_hints( plan: LogicalPlan, rewriter: &mut ScanHintRewriter, + fused_series_input: bool, ) -> Result> { let LogicalPlan::TableScan(mut table_scan) = plan else { return Ok(Transformed::no(plan)); @@ -89,8 +93,10 @@ impl ScanHintRule { }; let use_last_row = rewriter.inside_single_evaluation && filters_preserve_last_row; - let has_hint = - rewriter.order_expr.is_some() || rewriter.ts_row_selector.is_some() || use_last_row; + let has_hint = rewriter.order_expr.is_some() + || rewriter.ts_row_selector.is_some() + || use_last_row + || fused_series_input; if !has_hint { return Ok(Transformed::no(LogicalPlan::TableScan(table_scan))); } @@ -101,6 +107,14 @@ impl ScanHintRule { // local as well as the new LastRow hint. let adapter = original.clone_for_scan(); Self::apply_hints(&adapter, rewriter); + if fused_series_input { + // The scan below a `FusedSeries` node reads packs of raw series, which a region only + // produces for a per-series request; the existing PromQL pipeline asks for the same + // distribution through a sort on `(series key, time)`, which the native node does not + // have below it. Applying it after the order hints makes the plan shape the deciding + // one: the node's own input is what reads the frames. + adapter.with_distribution(TimeSeriesDistribution::PerSeries); + } if use_last_row { // Apply the instant-derived hint after the aggregate hint. Both // select LastRow today, and this ordering preserves the existing @@ -275,6 +289,12 @@ struct ScanHintRewriter { ts_stack: Vec, Column)>>, inside_single_evaluation: bool, single_evaluation_stack: Vec, + /// Whether the node being visited is the input of a [`FusedSeries`] node. + /// + /// The flag only describes that one edge: it is set while a `FusedSeries` node is visited, so + /// the node visited next — that node's own input — sees it, and a node further down does not. + fused_series_input: bool, + fused_series_input_stack: Vec, } impl TreeNodeRewriter for ScanHintRewriter { @@ -285,6 +305,13 @@ impl TreeNodeRewriter for ScanHintRewriter { self.ts_stack.push(self.ts_row_selector.clone()); self.single_evaluation_stack .push(self.inside_single_evaluation); + self.fused_series_input_stack.push(self.fused_series_input); + // Read the edge this node sits on before recording what this node is, so a `FusedSeries` + // node hands the flag to its own input and to nothing below it. + let fused_series_input = std::mem::replace( + &mut self.fused_series_input, + matches!(&node, LogicalPlan::Extension(extension) if extension.node.as_any().is::()), + ); if let LogicalPlan::Sort(sort) = &node { self.order_expr = Some(sort.expr.clone()); @@ -321,7 +348,7 @@ impl TreeNodeRewriter for ScanHintRewriter { } } - ScanHintRule::set_hints(node, self) + ScanHintRule::set_hints(node, self, fused_series_input) } fn f_up(&mut self, node: LogicalPlan) -> Result> { @@ -334,6 +361,9 @@ impl TreeNodeRewriter for ScanHintRewriter { if let Some(previous) = self.single_evaluation_stack.pop() { self.inside_single_evaluation = previous; } + if let Some(previous) = self.fused_series_input_stack.pop() { + self.fused_series_input = previous; + } Ok(Transformed::no(node)) } } @@ -446,6 +476,7 @@ impl ScanHintRewriter { mod test { use std::collections::HashMap; use std::sync::Arc; + use std::time::{Duration, UNIX_EPOCH}; use datafusion::functions_aggregate::first_last::last_value_udaf; use datafusion::functions_aggregate::min_max::max_udaf; @@ -454,21 +485,26 @@ mod test { use datafusion::logical_expr::{WindowFrame, WindowFunctionDefinition}; use datafusion::prelude::JoinType; use datafusion_common::tree_node::TreeNodeRecursion; + use datafusion_common::{DFSchema, DFSchemaRef}; use datafusion_expr::expr::{ AggregateFunction, AggregateFunctionParams, Cast, WindowFunctionParams, }; use datafusion_expr::expr_fn::scalar_subquery; use datafusion_expr::{Extension, LogicalPlan, LogicalPlanBuilder, col, lit}; use datafusion_optimizer::OptimizerContext; - use datatypes::arrow::datatypes::DataType; + use datatypes::arrow::datatypes::{DataType, Field, TimeUnit}; use datatypes::data_type::ConcreteDataType; use datatypes::schema::ColumnSchema; - use promql::extension_plan::RangeManipulate; + use datatypes::series::SeriesOrdering; + use datatypes::series::samples::{FieldSpec, SeriesSchema}; + use promql::extension_plan::{FusedSeries, FusedSeriesBinding, RangeManipulate}; + use promql_parser::parser; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME; use store_api::storage::{RegionId, TimeSeriesRowSelector}; use super::*; + use crate::promql::native_compile::{compile, statement_grid}; fn scan_requests(plan: &LogicalPlan) -> Vec { scan_requests_with_names(plan) @@ -1520,4 +1556,123 @@ mod test { Some(TimeSeriesDistribution::PerSeries) ); } + + /// A `FusedSeries` node over `input`, bound to the tree a bare selector compiles to. + /// + /// The rule reads the plan shape, not what the node computes, so the simplest shape the + /// compiler accepts is enough for these tests. + fn fused_series(input: LogicalPlan) -> LogicalPlan { + let stmt = promql_parser::parser::EvalStmt { + expr: parser::parse("some_metric").unwrap(), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + let series_schema = Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { name: "v0".into() }]), + eval_grid: None, + }); + let tree = compile( + &stmt, + Arc::clone(&series_schema), + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + let binding = FusedSeriesBinding { + series_schema, + grid: statement_grid(&stmt).unwrap(), + metric_name: Some("some_metric".into()), + time_index: "ts".to_string(), + tags: Vec::new(), + }; + let output_schema: DFSchemaRef = Arc::new( + DFSchema::new_with_metadata( + vec![ + (None, Arc::new(Field::new("v0", DataType::Float64, true))), + ( + None, + Arc::new(Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + )), + ), + ], + Default::default(), + ) + .unwrap(), + ); + + LogicalPlan::Extension(Extension { + node: Arc::new( + FusedSeries::try_new( + Arc::from("some_metric"), + tree, + binding, + input, + output_schema, + ) + .unwrap(), + ), + }) + } + + /// The region only packs raw series frames for a per-series request, and the node below is + /// what reads them, so the node's own input asks for that distribution. + #[test] + fn a_scan_under_fused_series_asks_for_the_per_series_distribution() { + let provider = Arc::new(mock_table_provider(RegionId::new(1, 1))); + let plan = fused_series(scan_plan(provider.clone(), "t")); + + let rewritten = ScanHintRule + .rewrite(plan, &OptimizerContext::default()) + .unwrap() + .data; + + assert_eq!( + scan_requests(&rewritten)[0].distribution, + Some(TimeSeriesDistribution::PerSeries) + ); + // The hint belongs to this use-site only, never to the shared catalog provider. + assert_eq!(provider.scan_request().distribution, None); + } + + /// The ordinary path: a scan with no fused node above it is left exactly as it was. + #[test] + fn a_scan_not_under_fused_series_keeps_its_request() { + let provider = Arc::new(mock_table_provider(RegionId::new(1, 1))); + let plan = scan_plan(provider.clone(), "t"); + + let rewritten = ScanHintRule + .rewrite(plan, &OptimizerContext::default()) + .unwrap() + .data; + + assert_eq!(scan_requests(&rewritten)[0].distribution, None); + assert_eq!(provider.scan_request().distribution, None); + } + + /// Only the node's own input is a raw series reader; a scan further down feeds something else + /// and keeps the request it had. + #[test] + fn a_scan_below_fused_series_input_keeps_its_request() { + let provider = Arc::new(mock_table_provider(RegionId::new(1, 1))); + let scan = scan_plan(provider.clone(), "t"); + let input = LogicalPlanBuilder::from(scan) + .project(vec![col("ts"), col("v0")]) + .unwrap() + .build() + .unwrap(); + let plan = fused_series(input); + + let rewritten = ScanHintRule + .rewrite(plan, &OptimizerContext::default()) + .unwrap() + .data; + + assert_eq!(scan_requests(&rewritten)[0].distribution, None); + } }