From 8348d3d2d4600ade47dbbfb05fac36e4e5f2d53b Mon Sep 17 00:00:00 2001 From: discord9 Date: Mon, 2 Nov 2026 13:34:01 +0800 Subject: [PATCH] fix(promql): let a too-deep tree fall back instead of failing the query The depth limit was reported as `DataFusionError::Internal`, the error a compiler defect gets. The planner only keeps the legacy plan on `NotImplemented`, so enabling the native route turned an answer into an error: a statement compiling to more than 64 nodes -- 70 `time()` terms compile to 139 -- failed outright, where the legacy planner answers it. Switching a feature flag on must never make a query stop working. The two rejections now say what they mean. Exceeding the depth limit is "this statement is beyond what this route handles", which is `NotImplemented` and makes the planner fall back, from both the constructor and the deserializer. Every other rejection in `validate()` -- a mismatched binding, a map over a scalar, a window with the wrong grid -- stays `Internal`, because those mean the compiler built something inconsistent and reclassifying them would hide real defects as unsupported queries. A fallback test pins it: 70 `time()` terms must plan exactly as the legacy route does, with the node count in the message so a future change to the cap explains itself. `postcard` moves into the external-dependencies block. It is declared with `default-features = false`, and that block exists precisely for inherited dependencies whose features are narrowed. Signed-off-by: discord9 --- Cargo.toml | 2 +- src/promql/src/extension_plan/fused_series.rs | 43 +++++++------------ .../src/extension_plan/fused_series/tests.rs | 3 +- src/promql/src/native/plan.rs | 11 ++--- src/promql/src/native/plan/tests.rs | 28 +++++++++++- src/query/src/dist_plan/analyzer/test.rs | 4 +- src/query/src/dist_plan/commutativity.rs | 2 +- src/query/src/promql/planner/test.rs | 19 ++++++++ 8 files changed, 71 insertions(+), 41 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 45531bb2496..388c52067be 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -93,13 +93,13 @@ rust.unknown_lints = "deny" rust.unexpected_cfgs = { level = "warn", check-cfg = ['cfg(tokio_unstable)'] } [workspace.dependencies] -postcard = { version = "1.1.3", default-features = false, features = ["use-std"] } # DO_NOT_REMOVE_THIS: BEGIN_OF_EXTERNAL_DEPENDENCIES # We turn off default-features for some dependencies here so the workspaces which inherit them can # selectively turn them on if needed, since we can override default-features = true (from false) # for the inherited dependency but cannot do the reverse (override from true to false). # # See for more detaiils: https://github.com/rust-lang/cargo/issues/11329 +postcard = { version = "1.1.3", default-features = false, features = ["alloc", "use-std"] } ahash = { version = "0.8", features = ["compile-time-rng"] } aquamarine = "0.6" arrow = { version = "=59.2.0", features = ["prettyprint"] } diff --git a/src/promql/src/extension_plan/fused_series.rs b/src/promql/src/extension_plan/fused_series.rs index d4a7b62b896..a0f97418b3f 100644 --- a/src/promql/src/extension_plan/fused_series.rs +++ b/src/promql/src/extension_plan/fused_series.rs @@ -273,6 +273,7 @@ impl FusedSeries { .samples() .map_or(&[], |schema| &schema.fields), &self.binding.time_index, + matches!(self.tree.root().output, NativeType::Scalar(_)), ) }, }; @@ -547,7 +548,16 @@ impl FusedSeries { /// The label columns of the output, independent of the raw-input binding. fn output_tags(&self) -> Vec { - output_tags_for_tree(&self.tree) + output_tags_from_schema( + self.output_schema.as_arrow(), + self.tree + .root() + .output + .samples() + .map_or(&[], |schema| &schema.fields), + &self.binding.time_index, + matches!(self.tree.root().output, NativeType::Scalar(_)), + ) } } @@ -823,6 +833,7 @@ fn output_tags_from_schema( schema: &datafusion::arrow::datatypes::Schema, series_fields: &[FieldSpec], time: &str, + scalar_root: bool, ) -> Vec { let field_names = series_fields .iter() @@ -831,34 +842,10 @@ fn output_tags_from_schema( schema .fields() .iter() - .filter(|field| field.name() != time && !field_names.contains(field.name().as_str())) - .map(|field| field.name().clone()) - .collect() -} - -fn output_tags_for_tree(tree: &NativeTree) -> Vec { - let Some(schema) = tree.root().output.samples() else { - return Vec::new(); - }; - let field_names = schema - .fields - .iter() - .map(|field| field.name.as_ref()) - .collect::>(); - collect_source_schemas(tree) - .unwrap_or_default() - .into_iter() - .flat_map(|source| { - source - .fields - .iter() - .map(|field| field.name.clone()) - .collect::>() + .filter(|field| { + field.name() != time && !field_names.contains(field.name().as_str()) && !scalar_root }) - .filter(|name| !field_names.contains(name.as_ref())) - .map(|name| name.to_string()) - .collect::>() - .into_iter() + .map(|field| field.name().clone()) .collect() } diff --git a/src/promql/src/extension_plan/fused_series/tests.rs b/src/promql/src/extension_plan/fused_series/tests.rs index 5ee9b469fcb..04acfdb4ddf 100644 --- a/src/promql/src/extension_plan/fused_series/tests.rs +++ b/src/promql/src/extension_plan/fused_series/tests.rs @@ -1044,8 +1044,7 @@ fn plain_envelope( output_tags: Vec, ) -> Vec { let binding = binding_of(&tree, raw_schema(), binding_tags); - let schema_tags = super::output_tags_for_tree(&tree); - let schema = super::plain_rows_schema(&tree, &schema_tags, TIME).unwrap(); + let schema = super::plain_rows_schema(&tree, &output_tags, TIME).unwrap(); let node = FusedSeries::try_new( Arc::from(PROMQL), tree, diff --git a/src/promql/src/native/plan.rs b/src/promql/src/native/plan.rs index 01882fdaa53..c1c65e05168 100644 --- a/src/promql/src/native/plan.rs +++ b/src/promql/src/native/plan.rs @@ -538,9 +538,8 @@ pub struct NativeTree { impl NativeTree { /// Creates a tree over `root`, checking that every node's contract holds. /// - /// A rejected tree is a compiler defect rather than a query error: the compiler only builds a - /// node after deciding that this combination is supported, so the error is - /// [`DataFusionError::Internal`] and not a statement the client wrote wrong. + /// A structurally inconsistent tree is a compiler defect; a tree beyond the native depth + /// limit is unsupported by this route and can fall back to the existing planner. pub fn try_new(root: NativeNode) -> Result { ensure_native_node_depth(&root)?; root.validate()?; @@ -584,14 +583,16 @@ impl<'de> serde::Deserialize<'de> for NativeTree { root: NativeNode, } let wire = Wire::deserialize(deserializer)?; - Self::try_new(wire.root).map_err(serde::de::Error::custom) + ensure_native_node_depth(&wire.root).map_err(serde::de::Error::custom)?; + wire.root.validate().map_err(serde::de::Error::custom)?; + Ok(Self { root: wire.root }) } } fn ensure_native_node_depth(root: &NativeNode) -> Result<()> { fn visit(node: &NativeNode, depth: usize) -> Result<()> { if depth > MAX_NATIVE_NODE_DEPTH { - return Err(invalid(&format!( + return Err(DataFusionError::NotImplemented(format!( "native tree nesting exceeds {MAX_NATIVE_NODE_DEPTH}" ))); } diff --git a/src/promql/src/native/plan/tests.rs b/src/promql/src/native/plan/tests.rs index 1e2b3026645..1d47c4f7fe7 100644 --- a/src/promql/src/native/plan/tests.rs +++ b/src/promql/src/native/plan/tests.rs @@ -14,7 +14,7 @@ use std::sync::Arc; -use datafusion::common::DFSchema; +use datafusion::common::{DFSchema, DataFusionError}; use datatypes::series::SeriesOrdering; use datatypes::series::samples::{FieldSpec, SeriesSchema}; use datatypes::series::time::{EvalGrid, GridSpec}; @@ -1493,6 +1493,27 @@ fn a_vector_operation_cannot_consume_raw_samples() { assert!(error.to_string().contains("an evaluated vector"), "{error}"); } +/// A samples root without an input is refused by FusedSeries. +#[test] +fn a_samples_root_without_an_input_is_refused() { + let tree = NativeTree::try_new(input()).unwrap(); + let binding = FusedSeriesBinding { + grid: grid(0, 60_000, 4), + time_index: "time".to_string(), + tags: vec![], + metric_name: None, + }; + let error = FusedSeries::try_new( + Arc::from("some_metric"), + tree, + binding, + Vec::new(), + Arc::new(DFSchema::empty()), + ) + .expect_err("a samples root without an input is refused"); + assert!(error.to_string().contains("source indexes"), "{error}"); +} + /// FusedSeries must reject a tree constructed outside the checked NativeTree constructor. #[test] fn fused_series_rejects_a_tree_over_the_depth_limit() { @@ -1611,7 +1632,10 @@ fn native_tree_depth_limit_holds_across_lifecycle_on_named_worker_stack() { postcard::from_bytes::(&bytes).unwrap_err(), postcard::Error::SerdeDeCustom ); - assert!(NativeTree::try_new(over_limit.root.clone()).is_err()); + assert!(matches!( + NativeTree::try_new(over_limit.root.clone()), + Err(DataFusionError::NotImplemented(_)) + )); drop(over_limit); }) .unwrap() diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index a899c824679..361321f6702 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -98,11 +98,11 @@ fn ordinary_encodable_plan_keeps_in_process_encodability_gate() { let in_process_rewriter = PlanRewriter::new(&plan, true); assert!(in_process_rewriter.whole_plan_encodable); - assert!(!in_process_rewriter.skip_encodability_gate); + assert!(!in_process_rewriter.skip_in_process_encodability); let wire_rewriter = PlanRewriter::new(&plan, false); assert!(wire_rewriter.whole_plan_encodable); - assert!(!wire_rewriter.skip_encodability_gate); + assert!(!wire_rewriter.skip_in_process_encodability); } struct MergeScanRemoteDynFilterProducerIdCollector<'a> { diff --git a/src/query/src/dist_plan/commutativity.rs b/src/query/src/dist_plan/commutativity.rs index 60c1ecfb5e2..d4bad40f236 100644 --- a/src/query/src/dist_plan/commutativity.rs +++ b/src/query/src/dist_plan/commutativity.rs @@ -1012,7 +1012,7 @@ mod tests { .as_any() .downcast_ref::() .unwrap(); - let inputs = finalize.inputs(); + let inputs = datafusion_expr::UserDefinedLogicalNodeCore::inputs(finalize); let [producer] = inputs.as_slice() else { unreachable!() }; diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index a54e4d4bf44..755ff15c8f8 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -2596,6 +2596,25 @@ async fn native_fused_whole_vector_operation_falls_back() { } } +#[tokio::test] +async fn native_fused_depth_limit_falls_back_to_legacy_plan() { + let query = std::iter::repeat_n("time()", 70) + .collect::>() + .join(" + "); + let native = build_native_fused_plan(&query, true).await; + let legacy = build_native_fused_plan(&query, false).await; + let display = native.display_indent().to_string(); + assert!( + !display.contains("FusedSeries"), + "70 time() terms compile to 139 native nodes, exceeding the 64-node cap: {display}" + ); + assert_eq!( + display, + legacy.display_indent().to_string(), + "139 native nodes exceed the 64-node cap, so planning must preserve the legacy plan" + ); +} + #[tokio::test] async fn native_fused_unsupported_shape_falls_back() { for query in [