diff --git a/nodedb-cluster-tests/tests/common_suite/cases/shuffle_aggregate_cross_node.rs b/nodedb-cluster-tests/tests/common_suite/cases/shuffle_aggregate_cross_node.rs index d534d18ba..08beb3d1a 100644 --- a/nodedb-cluster-tests/tests/common_suite/cases/shuffle_aggregate_cross_node.rs +++ b/nodedb-cluster-tests/tests/common_suite/cases/shuffle_aggregate_cross_node.rs @@ -107,6 +107,8 @@ fn producer_plan(rows: &[&Row]) -> Vec { rows: msgpack_array(rows), filters: Vec::new(), projection: Vec::new(), + computed_columns: Vec::new(), + window_functions: Vec::new(), sort_keys: Vec::new(), limit: None, offset: 0, diff --git a/nodedb-cluster-tests/tests/common_suite/cases/shuffle_consume_cross_node.rs b/nodedb-cluster-tests/tests/common_suite/cases/shuffle_consume_cross_node.rs index 205772766..f343e19fe 100644 --- a/nodedb-cluster-tests/tests/common_suite/cases/shuffle_consume_cross_node.rs +++ b/nodedb-cluster-tests/tests/common_suite/cases/shuffle_consume_cross_node.rs @@ -104,6 +104,8 @@ fn provider_scan_plan(rows: &[&Row]) -> Vec { rows: msgpack_array(rows), filters: Vec::new(), projection: Vec::new(), + computed_columns: Vec::new(), + window_functions: Vec::new(), sort_keys: Vec::new(), limit: None, offset: 0, diff --git a/nodedb-cluster-tests/tests/common_suite/cases/shuffle_produce_cross_node.rs b/nodedb-cluster-tests/tests/common_suite/cases/shuffle_produce_cross_node.rs index 3f0602c28..49f8bf830 100644 --- a/nodedb-cluster-tests/tests/common_suite/cases/shuffle_produce_cross_node.rs +++ b/nodedb-cluster-tests/tests/common_suite/cases/shuffle_produce_cross_node.rs @@ -108,6 +108,8 @@ fn provider_scan_plan(rows: &[Vec]) -> Vec { rows: msgpack_array(rows), filters: Vec::new(), projection: Vec::new(), + computed_columns: Vec::new(), + window_functions: Vec::new(), sort_keys: Vec::new(), limit: None, offset: 0, diff --git a/nodedb-physical/src/physical_plan/collection.rs b/nodedb-physical/src/physical_plan/collection.rs index 50a6ef243..90001efca 100644 --- a/nodedb-physical/src/physical_plan/collection.rs +++ b/nodedb-physical/src/physical_plan/collection.rs @@ -113,6 +113,8 @@ impl PhysicalPlan { PhysicalPlan::Query(QueryOp::Exchange(op)) => op.child.collection(), // PostProcess: recurse into the materialized input plan. PhysicalPlan::Query(QueryOp::PostProcess { input, .. }) => input.collection(), + // SetOp merges N branches; no single collection names the node. + PhysicalPlan::Query(QueryOp::SetOp { .. }) => None, // ProviderScan is a catalog/constant source — no user collection. PhysicalPlan::Query(QueryOp::ProviderScan { .. }) => None, // KV ops carry their own collection (sorted-index-only ops → None). diff --git a/nodedb-physical/src/physical_plan/mod.rs b/nodedb-physical/src/physical_plan/mod.rs index 98967dfb2..0df5f1b35 100644 --- a/nodedb-physical/src/physical_plan/mod.rs +++ b/nodedb-physical/src/physical_plan/mod.rs @@ -22,6 +22,7 @@ pub mod plan; pub mod query; pub mod rls_write_check_accessor; pub mod routing; +pub mod set_op; pub mod sort_key; pub mod spatial; pub mod streaming; @@ -51,6 +52,7 @@ pub use meta::MetaOp; pub use plan::PhysicalPlan; pub use query::{AggregateSpec, GroupKeySpec, JoinProjection, QueryOp}; pub use routing::plan_contains_cluster_partitioned_leaf; +pub use set_op::SetOpKind; pub use sort_key::SortKeySpec; pub use spatial::{SpatialOp, SpatialPredicate}; pub use text::TextOp; diff --git a/nodedb-physical/src/physical_plan/query.rs b/nodedb-physical/src/physical_plan/query.rs index 96f087d3c..67e023498 100644 --- a/nodedb-physical/src/physical_plan/query.rs +++ b/nodedb-physical/src/physical_plan/query.rs @@ -78,6 +78,14 @@ pub enum QueryOp { /// Output column names to keep. Empty = emit all columns. #[serde(default)] projection: Vec, + /// Serialized `Vec` (MessagePack), same encoding as + /// `DocumentOp::Scan::computed_columns`. Empty = none. + #[serde(default)] + computed_columns: Vec, + /// Serialized `Vec` (MessagePack), same encoding as + /// `DocumentOp::Scan::window_functions`. Empty = none. + #[serde(default)] + window_functions: Vec, /// ORDER BY terms, each an expression. Empty = unordered. #[serde(default)] sort_keys: Vec, @@ -118,6 +126,14 @@ pub enum QueryOp { /// Output column names to keep. Empty = emit all columns. #[serde(default)] projection: Vec, + /// Serialized `Vec` (MessagePack), same encoding as + /// `DocumentOp::Scan::computed_columns`. Empty = none. + #[serde(default)] + computed_columns: Vec, + /// Serialized `Vec` (MessagePack), same encoding as + /// `DocumentOp::Scan::window_functions`. Empty = none. + #[serde(default)] + window_functions: Vec, /// ORDER BY terms, each an expression. Empty = unordered. #[serde(default)] sort_keys: Vec, @@ -133,18 +149,33 @@ pub enum QueryOp { distinct: bool, }, + /// Set operation over N materialized children. Coordinator-only: the + /// resolver materializes every child and merges the rows into one + /// `ProviderScan` before dispatch. A Data-Plane core never sees this node. + /// + /// Lowered from a derived-table body that is `UNION [ALL]`, + /// `INTERSECT [ALL]`, or `EXCEPT [ALL]`, so the body is one relation for + /// an outer [`QueryOp::PostProcess`] or input-sourced [`QueryOp::Aggregate`]. + SetOp { + /// Child relations in SQL order. Each sharded child is wrapped in + /// `Exchange{Gather}` by the converter so its gather runs once. + inputs: Vec, + /// Which set operation merges the inputs. + op: crate::physical_plan::SetOpKind, + }, + /// Aggregate: GROUP BY + aggregate functions. Aggregate { collection: QualifiedCollection, /// Optional sub-plan whose decoded rows are aggregated instead of - /// scanning `collection` per-shard. `Some` currently means EXACTLY a - /// catalog source (a `ProviderScan` lowered by the converter): the - /// aggregate runs over the coordinator-materialized catalog rows and is - /// therefore coordinator-local (never broadcast — see - /// `is_sharded_source`). `None` = legacy path: scan the named - /// `collection` on every shard. `collection` stays populated in both - /// cases so downstream RLS / permission / classification continue to - /// read it; the executor simply prefers `input` when present. + /// scanning `collection` per-shard. `Some` = an input-sourced + /// aggregate over a materialized relation: a catalog `ProviderScan`, + /// or any derived-table body the coordinator materializes into a + /// `ProviderScan` before dispatch. Coordinator-local, never broadcast + /// (see `is_sharded_source`). `None` = scan the named `collection` on + /// every shard. `collection` stays populated in both cases so + /// downstream RLS / permission / classification continue to read it; + /// the executor prefers `input` when present. #[serde(default)] input: Option>, group_by: Vec, diff --git a/nodedb-physical/src/physical_plan/routing.rs b/nodedb-physical/src/physical_plan/routing.rs index 0fda15961..ba224ae09 100644 --- a/nodedb-physical/src/physical_plan/routing.rs +++ b/nodedb-physical/src/physical_plan/routing.rs @@ -68,6 +68,11 @@ pub fn plan_contains_cluster_partitioned_leaf(plan: &PhysicalPlan) -> bool { plan_contains_cluster_partitioned_leaf(input) } + // Recurse through every SetOp branch for the same reason. + PhysicalPlan::Query(QueryOp::SetOp { inputs, .. }) => { + inputs.iter().any(plan_contains_cluster_partitioned_leaf) + } + // Recurse through lateral outer plans. PhysicalPlan::Query(QueryOp::LateralTopK { outer_plan, .. }) | PhysicalPlan::Query(QueryOp::LateralLoop { outer_plan, .. }) => { @@ -159,6 +164,9 @@ impl PhysicalPlan { right_bitmap, ) } + // Coordinator-local: the resolver materializes every branch and + // merges on the coordinator, so the node itself is never fanned out. + PhysicalPlan::Query(QueryOp::SetOp { .. }) => false, _ => self.is_sharded_source_leaf(), } } diff --git a/nodedb-physical/src/physical_plan/set_op.rs b/nodedb-physical/src/physical_plan/set_op.rs new file mode 100644 index 000000000..d4f3c1ae1 --- /dev/null +++ b/nodedb-physical/src/physical_plan/set_op.rs @@ -0,0 +1,33 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Set-operation kinds for [`crate::physical_plan::QueryOp::SetOp`]. + +/// Which SQL set operation a [`crate::physical_plan::QueryOp::SetOp`] node +/// applies over its materialized inputs. Coordinator-resolved, never reaches +/// a Data-Plane core. +#[derive( + Debug, + Clone, + Copy, + PartialEq, + Eq, + serde::Serialize, + serde::Deserialize, + zerompk::ToMessagePack, + zerompk::FromMessagePack, +)] +#[msgpack(c_enum)] +pub enum SetOpKind { + /// `UNION ALL`: concatenate every input in order. + UnionAll, + /// `UNION`: concatenate, then drop duplicate rows. + UnionDistinct, + /// `INTERSECT`: rows present in every input, deduplicated. + Intersect, + /// `INTERSECT ALL`: rows present in every input, bag semantics. + IntersectAll, + /// `EXCEPT`: rows of the first input absent from the rest, deduplicated. + Except, + /// `EXCEPT ALL`: rows of the first input absent from the rest, bag semantics. + ExceptAll, +} diff --git a/nodedb-physical/src/physical_plan/streaming.rs b/nodedb-physical/src/physical_plan/streaming.rs index 76d92634b..cd728fc93 100644 --- a/nodedb-physical/src/physical_plan/streaming.rs +++ b/nodedb-physical/src/physical_plan/streaming.rs @@ -42,8 +42,9 @@ impl PhysicalPlan { sort_keys, offset, distinct, + window_functions, .. - }) => sort_keys.is_empty() && *offset == 0 && !*distinct, + }) => sort_keys.is_empty() && *offset == 0 && !*distinct && window_functions.is_empty(), // Every other Document / Kv / Columnar / Timeseries op, plus all // other engines and query ops, are not unordered-streamable. diff --git a/nodedb-query/src/window/eval.rs b/nodedb-query/src/window/eval.rs index cf6085ad6..ad741cdf3 100644 --- a/nodedb-query/src/window/eval.rs +++ b/nodedb-query/src/window/eval.rs @@ -15,7 +15,8 @@ use super::spec::WindowFuncSpec; /// /// `rows` is the sorted result set. Each row is a `(doc_id, serde_json::Value)`. /// The same rows are mutated in place with window columns appended to each -/// document. +/// document. The row array keeps its input order; each spec's partitions are +/// ordered by that spec's own ORDER BY, independent of the row array order. /// /// Unknown window function names must be rejected by the planner before /// reaching this dispatcher; an unrecognised name here is an internal bug @@ -29,7 +30,7 @@ pub fn evaluate_window_functions( specs: &[WindowFuncSpec], ) -> Result<(), crate::expr::EvalError> { for spec in specs { - let partitions = build_partitions(rows, &spec.partition_by)?; + let partitions = build_partitions(rows, &spec.partition_by, &spec.order_by)?; for partition_indices in &partitions { match spec.func_name.as_str() { @@ -146,9 +147,11 @@ mod tests { frame: WindowFrame::default(), }; evaluate_window_functions(&mut rows, &[spec]).unwrap(); - assert_eq!(rows[0].1["running_total"], json!(100.0)); - assert_eq!(rows[1].1["running_total"], json!(220.0)); - assert_eq!(rows[2].1["running_total"], json!(310.0)); + // The frame runs in salary order within each dept, not in row + // arrival order: eng = Carol(90) → Alice(100) → Bob(120). + assert_eq!(rows[0].1["running_total"], json!(190.0)); + assert_eq!(rows[1].1["running_total"], json!(310.0)); + assert_eq!(rows[2].1["running_total"], json!(90.0)); assert_eq!(rows[3].1["running_total"], json!(80.0)); assert_eq!(rows[4].1["running_total"], json!(190.0)); } @@ -262,6 +265,33 @@ mod tests { assert_eq!(rows[4].1["nv"], json!(2)); } + #[test] + fn rank_orders_by_spec_order_by_not_row_arrival_order() { + // Rows arrive as Alice(100), Bob(120), Carol(90) within dept "eng" — + // not sorted by salary. RANK() OVER (ORDER BY salary DESC) must rank + // by salary, and the row array order must stay unchanged. + let mut rows = make_rows(); + let spec = WindowFuncSpec { + alias: "rnk".into(), + func_name: "rank".into(), + args: vec![], + partition_by: vec![SqlExpr::Column("dept".into())], + order_by: vec![(SqlExpr::Column("salary".into()), false)], + frame: WindowFrame::default(), + }; + evaluate_window_functions(&mut rows, &[spec]).unwrap(); + assert_eq!(rows[0].1["name"], json!("Alice")); + assert_eq!(rows[1].1["name"], json!("Bob")); + assert_eq!(rows[2].1["name"], json!("Carol")); + assert_eq!(rows[3].1["name"], json!("Dave")); + assert_eq!(rows[4].1["name"], json!("Eve")); + assert_eq!(rows[0].1["rnk"], json!(2)); // Alice, salary 100 + assert_eq!(rows[1].1["rnk"], json!(1)); // Bob, salary 120 + assert_eq!(rows[2].1["rnk"], json!(3)); // Carol, salary 90 + assert_eq!(rows[3].1["rnk"], json!(2)); // Dave, salary 80 + assert_eq!(rows[4].1["rnk"], json!(1)); // Eve, salary 110 + } + #[test] #[should_panic(expected = "should have been rejected at planning time")] fn unknown_function_panics_at_evaluator() { diff --git a/nodedb-query/src/window/helpers.rs b/nodedb-query/src/window/helpers.rs index 3bd7b36fb..487d30c4f 100644 --- a/nodedb-query/src/window/helpers.rs +++ b/nodedb-query/src/window/helpers.rs @@ -6,35 +6,112 @@ use std::collections::HashMap; use crate::expr::types::SqlExpr; -/// Group row indices by partition key, preserving first-seen partition order. +/// Group row indices by partition key, preserving first-seen partition order, +/// then sort each partition's indices by the spec's ORDER BY. /// -/// A division/modulo-by-zero in a PARTITION BY expression propagates as -/// `Err(EvalError::DivisionByZero)` rather than being folded to NULL. +/// The returned index lists are ordered by `order_by`; the `rows` array +/// itself keeps its input order — only the per-partition index lists move. +/// +/// A division/modulo-by-zero in a PARTITION BY or ORDER BY expression +/// propagates as `Err(EvalError::DivisionByZero)` rather than being folded to +/// NULL. pub(super) fn build_partitions( rows: &[(String, serde_json::Value)], partition_by: &[SqlExpr], + order_by: &[(SqlExpr, bool)], ) -> Result>, crate::expr::EvalError> { - if partition_by.is_empty() { - return Ok(vec![(0..rows.len()).collect()]); - } + let mut partitions = if partition_by.is_empty() { + vec![(0..rows.len()).collect::>()] + } else { + let mut groups: HashMap> = HashMap::new(); + let mut order = Vec::new(); + + for (i, (_id, doc)) in rows.iter().enumerate() { + let key: String = partition_by + .iter() + .map(|expr| eval_expr_on_json(expr, doc).map(|v| v.to_string())) + .collect::, _>>()? + .join("\x00"); + let entry = groups.entry(key.clone()).or_default(); + if entry.is_empty() { + order.push(key); + } + entry.push(i); + } + + order.iter().filter_map(|k| groups.remove(k)).collect() + }; - let mut groups: HashMap> = HashMap::new(); - let mut order = Vec::new(); + if !order_by.is_empty() { + let mut keys: Vec> = Vec::with_capacity(rows.len()); + for (_id, doc) in rows.iter() { + keys.push( + order_by + .iter() + .map(|(expr, _)| eval_expr_on_json(expr, doc)) + .collect::, _>>()?, + ); + } - for (i, (_id, doc)) in rows.iter().enumerate() { - let key: String = partition_by - .iter() - .map(|expr| eval_expr_on_json(expr, doc).map(|v| v.to_string())) - .collect::, _>>()? - .join("\x00"); - let entry = groups.entry(key.clone()).or_default(); - if entry.is_empty() { - order.push(key); + for partition in &mut partitions { + partition.sort_by(|&a, &b| compare_order_keys(&keys[a], &keys[b], order_by)); } - entry.push(i); } - Ok(order.iter().filter_map(|k| groups.remove(k)).collect()) + Ok(partitions) +} + +/// Decide NULL placement for one ORDER BY column, shared by every window +/// evaluator's `compare_order_keys`. +/// +/// NULL placement follows PostgreSQL's default: ASC places NULLs last, DESC +/// places NULLs first. A window spec carries no explicit NULLS FIRST/LAST +/// override, so this default is fixed by direction alone. Returns `None` +/// when neither value is NULL, leaving the non-null comparison to the +/// caller. +pub(super) fn null_order( + a_null: bool, + b_null: bool, + ascending: bool, +) -> Option { + use std::cmp::Ordering; + let nulls_first = !ascending; + match (a_null, b_null) { + (true, true) => Some(Ordering::Equal), + (true, false) => Some(if nulls_first { + Ordering::Less + } else { + Ordering::Greater + }), + (false, true) => Some(if nulls_first { + Ordering::Greater + } else { + Ordering::Less + }), + (false, false) => None, + } +} + +/// Compare two rows' pre-evaluated ORDER BY keys. +fn compare_order_keys( + a: &[serde_json::Value], + b: &[serde_json::Value], + order_by: &[(SqlExpr, bool)], +) -> std::cmp::Ordering { + use std::cmp::Ordering; + for (idx, (_, ascending)) in order_by.iter().enumerate() { + let (Some(va), Some(vb)) = (a.get(idx), b.get(idx)) else { + continue; + }; + let ord = null_order(va.is_null(), vb.is_null(), *ascending).unwrap_or_else(|| { + let c = crate::json_expr::compare_json(va, vb); + if *ascending { c } else { c.reverse() } + }); + if ord != Ordering::Equal { + return ord; + } + } + Ordering::Equal } pub(super) fn set_window_col(row: &mut serde_json::Value, alias: &str, val: serde_json::Value) { diff --git a/nodedb-query/src/window/mod.rs b/nodedb-query/src/window/mod.rs index 47cfe48b0..f1e4f5dda 100644 --- a/nodedb-query/src/window/mod.rs +++ b/nodedb-query/src/window/mod.rs @@ -16,6 +16,7 @@ pub mod running; pub mod spec; pub mod value_agg; pub mod value_eval; +pub mod value_partition; pub use eval::evaluate_window_functions; pub use spec::{FrameBound, WindowFrame, WindowFuncSpec}; diff --git a/nodedb-query/src/window/value_eval.rs b/nodedb-query/src/window/value_eval.rs index c22230d94..1ba2da874 100644 --- a/nodedb-query/src/window/value_eval.rs +++ b/nodedb-query/src/window/value_eval.rs @@ -12,6 +12,7 @@ use nodedb_types::Value; use super::spec::WindowFuncSpec; use super::value_agg::apply_v_aggregate; +use super::value_partition::build_value_partitions; use crate::expr::types::SqlExpr; use crate::value_ops::compare_values; @@ -35,7 +36,9 @@ pub enum WindowError { /// /// `column_index` maps column name → position in each row slice. /// For each spec, one `Value` is appended to every row. Returns the list of -/// new column names, one per spec in spec order. +/// new column names, one per spec in spec order. `rows` keeps its input +/// order; each spec's partitions are ordered by that spec's own ORDER BY, +/// independent of the row order. pub fn evaluate_window_functions_value( rows: &mut [Vec], column_index: &HashMap, @@ -91,48 +94,7 @@ pub fn evaluate_window_functions_value( Ok(new_cols) } -// ── Partition building ──────────────────────────────────────────────────────── - -fn build_value_partitions( - rows: &[Vec], - column_index: &HashMap, - spec: &WindowFuncSpec, -) -> Result>, WindowError> { - if spec.partition_by.is_empty() { - return Ok(vec![(0..rows.len()).collect()]); - } - - let mut groups: HashMap> = HashMap::new(); - let mut order: Vec = Vec::new(); - - for (i, row) in rows.iter().enumerate() { - let key = partition_key(row, column_index, &spec.partition_by)?; - let entry = groups.entry(key.clone()).or_default(); - if entry.is_empty() { - order.push(key); - } - entry.push(i); - } - - Ok(order.iter().filter_map(|k| groups.remove(k)).collect()) -} - -fn partition_key( - row: &[Value], - column_index: &HashMap, - partition_by: &[SqlExpr], -) -> Result { - Ok(partition_by - .iter() - .map(|expr| { - let v = eval_arg_for_row(expr, row, column_index)?; - Ok(format!("{v:?}")) - }) - .collect::, WindowError>>()? - .join("\x00")) -} - -// ── Value comparison helpers (pub(super) for value_agg) ─────────────────────── +// ── Value comparison helpers (pub(super) for value_agg, value_partition) ───── pub(super) fn cmp_values(a: &Value, b: &Value) -> std::cmp::Ordering { match (a, b) { @@ -619,6 +581,19 @@ mod tests { assert_eq!(out_int(&rows, 2), vec![1, 2, 1]); } + #[test] + fn rank_orders_by_spec_order_by_not_row_arrival_order() { + // Rows arrive as 10, 30, 20 — not sorted by value. RANK() OVER + // (ORDER BY v DESC) must rank 30(1), 20(2), 10(3); the row array + // order must stay 10, 30, 20. + let mut rows = rows_v(&[10, 30, 20]); + let cols = ci(&["v"]); + let s = spec("rank", vec![], vec![], vec![(col("v"), false)]); + evaluate_window_functions_value(&mut rows, &cols, &[s]).unwrap(); + assert_eq!(out_int(&rows, 0), vec![10, 30, 20]); + assert_eq!(out_int(&rows, 1), vec![3, 1, 2]); + } + #[test] fn unknown_function_errors() { let mut rows = rows_v(&[1]); diff --git a/nodedb-query/src/window/value_partition.rs b/nodedb-query/src/window/value_partition.rs new file mode 100644 index 000000000..3a7a5d0c3 --- /dev/null +++ b/nodedb-query/src/window/value_partition.rs @@ -0,0 +1,102 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Partition building and per-partition ORDER BY sort for the Value-native +//! window evaluator. + +use std::collections::HashMap; + +use nodedb_types::Value; + +use super::helpers::null_order; +use super::spec::WindowFuncSpec; +use super::value_eval::{WindowError, cmp_values, eval_arg_for_row}; +use crate::expr::types::SqlExpr; + +/// Group row indices by partition key, preserving first-seen partition +/// order, then sort each partition's indices by the spec's ORDER BY. +/// +/// The returned index lists are ordered by `spec.order_by`; `rows` itself +/// keeps its input order — only the per-partition index lists move. +pub(super) fn build_value_partitions( + rows: &[Vec], + column_index: &HashMap, + spec: &WindowFuncSpec, +) -> Result>, WindowError> { + let mut partitions = if spec.partition_by.is_empty() { + vec![(0..rows.len()).collect::>()] + } else { + let mut groups: HashMap> = HashMap::new(); + let mut order: Vec = Vec::new(); + + for (i, row) in rows.iter().enumerate() { + let key = partition_key(row, column_index, &spec.partition_by)?; + let entry = groups.entry(key.clone()).or_default(); + if entry.is_empty() { + order.push(key); + } + entry.push(i); + } + + order.iter().filter_map(|k| groups.remove(k)).collect() + }; + + if !spec.order_by.is_empty() { + let mut keys: Vec> = Vec::with_capacity(rows.len()); + for row in rows.iter() { + keys.push( + spec.order_by + .iter() + .map(|(expr, _)| eval_arg_for_row(expr, row, column_index)) + .collect::, _>>()?, + ); + } + + for partition in &mut partitions { + partition.sort_by(|&a, &b| compare_order_keys(&keys[a], &keys[b], &spec.order_by)); + } + } + + Ok(partitions) +} + +fn partition_key( + row: &[Value], + column_index: &HashMap, + partition_by: &[SqlExpr], +) -> Result { + Ok(partition_by + .iter() + .map(|expr| { + let v = eval_arg_for_row(expr, row, column_index)?; + Ok(format!("{v:?}")) + }) + .collect::, WindowError>>()? + .join("\x00")) +} + +/// Compare two rows' pre-evaluated ORDER BY keys. +fn compare_order_keys( + a: &[Value], + b: &[Value], + order_by: &[(SqlExpr, bool)], +) -> std::cmp::Ordering { + use std::cmp::Ordering; + for (idx, (_, ascending)) in order_by.iter().enumerate() { + let (Some(va), Some(vb)) = (a.get(idx), b.get(idx)) else { + continue; + }; + let ord = null_order( + matches!(va, Value::Null), + matches!(vb, Value::Null), + *ascending, + ) + .unwrap_or_else(|| { + let c = cmp_values(va, vb); + if *ascending { c } else { c.reverse() } + }); + if ord != Ordering::Equal { + return ord; + } + } + Ordering::Equal +} diff --git a/nodedb-sql/src/planner/catalog_fold.rs b/nodedb-sql/src/planner/catalog_fold.rs index d29db868b..5aca0895d 100644 --- a/nodedb-sql/src/planner/catalog_fold.rs +++ b/nodedb-sql/src/planner/catalog_fold.rs @@ -112,6 +112,7 @@ fn walk_plan( input, mut filters, mut projection, + mut window_functions, mut sort_keys, offset, distinct, @@ -121,11 +122,13 @@ fn walk_plan( fold_filter(f, catalog, database_id, tenant_id); } fold_projection(&mut projection, catalog, database_id, tenant_id); + fold_windows(&mut window_functions, catalog, database_id, tenant_id); fold_sort_keys(&mut sort_keys, catalog, database_id, tenant_id); SqlPlan::Subquery { input: Box::new(walk_plan(*input, catalog, database_id, tenant_id)), filters, projection, + window_functions, sort_keys, offset, distinct, diff --git a/nodedb-sql/src/planner/catalog_plan_validate.rs b/nodedb-sql/src/planner/catalog_plan_validate.rs index 8f947256e..fb0a24f18 100644 --- a/nodedb-sql/src/planner/catalog_plan_validate.rs +++ b/nodedb-sql/src/planner/catalog_plan_validate.rs @@ -106,12 +106,14 @@ pub(super) fn validate_catalog_exprs( input, filters, projection, + window_functions, sort_keys, .. } => { validate_catalog_exprs(input, catalog, database_id, tenant_id)?; validate_filters(filters, catalog, database_id, tenant_id)?; validate_projection(projection, catalog, database_id, tenant_id)?; + validate_windows(window_functions, catalog, database_id, tenant_id)?; validate_sort_keys(sort_keys, catalog, database_id, tenant_id)?; } SqlPlan::Join { diff --git a/nodedb-sql/src/planner/select/entry.rs b/nodedb-sql/src/planner/select/entry.rs index 48a41f8f6..70a2dc426 100644 --- a/nodedb-sql/src/planner/select/entry.rs +++ b/nodedb-sql/src/planner/select/entry.rs @@ -173,6 +173,7 @@ pub fn plan_query( SqlPlan::Subquery { filters, projection, + window_functions, sort_keys, offset, distinct, @@ -182,6 +183,7 @@ pub fn plan_query( input: Box::new(upgraded_leaf), filters, projection, + window_functions, sort_keys, offset, distinct, diff --git a/nodedb-sql/src/planner/select/post_process.rs b/nodedb-sql/src/planner/select/post_process.rs index 3d351cefa..fa3022587 100644 --- a/nodedb-sql/src/planner/select/post_process.rs +++ b/nodedb-sql/src/planner/select/post_process.rs @@ -37,6 +37,8 @@ pub(in crate::planner::select) fn post_process( input: Box::new(input), filters: Vec::new(), projection, + // The body keeps its own window specs; the tail evaluates none. + window_functions: Vec::new(), sort_keys, offset, distinct: false, diff --git a/nodedb-sql/src/types/plan/variants.rs b/nodedb-sql/src/types/plan/variants.rs index 7439de518..87ed75533 100644 --- a/nodedb-sql/src/types/plan/variants.rs +++ b/nodedb-sql/src/types/plan/variants.rs @@ -537,6 +537,8 @@ pub enum SqlPlan { filters: Vec, /// Outer projection (target list). Empty = inherit the body's columns. projection: Vec, + /// Window functions evaluated over the post-processed rows. Empty = none. + window_functions: Vec, /// Outer `ORDER BY` keys applied over the materialized rows. sort_keys: Vec, /// Outer `OFFSET` (0 = none). diff --git a/nodedb-sql/src/visitor/plan_visitor/args.rs b/nodedb-sql/src/visitor/plan_visitor/args.rs index f3678a436..2c89436e7 100644 --- a/nodedb-sql/src/visitor/plan_visitor/args.rs +++ b/nodedb-sql/src/visitor/plan_visitor/args.rs @@ -37,6 +37,7 @@ pub struct SubqueryVisitArgs<'a> { pub input: &'a SqlPlan, pub filters: &'a [Filter], pub projection: &'a [Projection], + pub window_functions: &'a [WindowSpec], pub sort_keys: &'a [SortKey], pub offset: usize, pub distinct: bool, diff --git a/nodedb-sql/src/visitor/plan_visitor/dispatch_rest.rs b/nodedb-sql/src/visitor/plan_visitor/dispatch_rest.rs index 0f5f61452..1a24ee07b 100644 --- a/nodedb-sql/src/visitor/plan_visitor/dispatch_rest.rs +++ b/nodedb-sql/src/visitor/plan_visitor/dispatch_rest.rs @@ -26,6 +26,7 @@ pub(super) fn dispatch_rest( input, filters, projection, + window_functions, sort_keys, offset, distinct, @@ -34,6 +35,7 @@ pub(super) fn dispatch_rest( input, filters, projection, + window_functions, sort_keys, offset: *offset, distinct: *distinct, diff --git a/nodedb/src/control/clone/resolver/rewrite.rs b/nodedb/src/control/clone/resolver/rewrite.rs index 237ee8f23..f1189ebbd 100644 --- a/nodedb/src/control/clone/resolver/rewrite.rs +++ b/nodedb/src/control/clone/resolver/rewrite.rs @@ -8,7 +8,7 @@ use nodedb_types::TenantId; use crate::control::state::SharedState; use nodedb_physical::physical_plan::{ - ColumnarOp, DocumentOp, ExchangeOp, KvOp, PhysicalPlan, QueryOp, TimeseriesOp, + ColumnarOp, DocumentOp, ExchangeOp, KvOp, PhysicalPlan, QueryOp, SetOpKind, TimeseriesOp, }; use nodedb_types::SystemTimeScope; @@ -117,6 +117,8 @@ pub fn rewrite_plan_for_source(params: RewriteForSourceParams<'_>) -> crate::Res input, filters, projection, + computed_columns, + window_functions, sort_keys, limit, offset, @@ -139,6 +141,8 @@ pub fn rewrite_plan_for_source(params: RewriteForSourceParams<'_>) -> crate::Res input: child, filters: filters.clone(), projection: projection.clone(), + computed_columns: computed_columns.clone(), + window_functions: window_functions.clone(), sort_keys: sort_keys.clone(), limit: *limit, offset: *offset, @@ -149,6 +153,54 @@ pub fn rewrite_plan_for_source(params: RewriteForSourceParams<'_>) -> crate::Res }) } + // SetOp: rewrite every branch. Branches that do not read the cloned + // collection yield no source task and are dropped, so the source-side + // node carries only the rows the target side is missing. That is + // sound for `UNION ALL` (the merge appends). Any other kind dedups or + // subtracts by exact row match against the target rows, which is + // unsound across an unmaterialized clone; refuse it the same way the + // task-level `post_set_op` refusal does. + PhysicalPlan::Query(QueryOp::SetOp { inputs, op }) => { + let mut rewritten_inputs = Vec::with_capacity(inputs.len()); + for input in inputs { + let rewritten = rewrite_plan_for_source(RewriteForSourceParams { + plan: input, + target_db_id, + source_db_id, + tenant_id, + target_coll, + source_coll, + effective_source_ms, + kv_surrogate_ceiling, + state, + })?; + if let SourceRewrite::Task(child) = rewritten { + rewritten_inputs.push(*child); + } + } + if rewritten_inputs.is_empty() { + return Ok(SourceRewrite::NoSourceTask); + } + match op { + SetOpKind::UnionAll => { + Ok(SourceRewrite::task(PhysicalPlan::Query(QueryOp::SetOp { + inputs: rewritten_inputs, + op: SetOpKind::UnionAll, + }))) + } + SetOpKind::UnionDistinct + | SetOpKind::Intersect + | SetOpKind::IntersectAll + | SetOpKind::Except + | SetOpKind::ExceptAll => Err(crate::Error::PlanError { + detail: format!( + "a set operation over '{target_coll}' cannot be read through an \ + unmaterialized clone; run ALTER DATABASE MATERIALIZE first" + ), + }), + } + } + PhysicalPlan::Document(DocumentOp::Scan { collection, limit, @@ -626,6 +678,8 @@ mod tests { input: Box::new(gather(plan)), filters: Vec::new(), projection: Vec::new(), + computed_columns: Vec::new(), + window_functions: Vec::new(), sort_keys: Vec::new(), limit: None, offset: 0, diff --git a/nodedb/src/control/exec_receiver/support.rs b/nodedb/src/control/exec_receiver/support.rs index b6c802352..c0a8b72cf 100644 --- a/nodedb/src/control/exec_receiver/support.rs +++ b/nodedb/src/control/exec_receiver/support.rs @@ -55,6 +55,9 @@ pub(super) fn plan_contains_exchange(plan: &PhysicalPlan) -> bool { // resolution, still carries an `Exchange{Gather}` — recurse so an // unresolved PostProcess is correctly flagged as Exchange-bearing. QueryOp::PostProcess { input, .. } => plan_contains_exchange(input), + // SetOp inputs are unresolved bodies; any of them can carry a + // Gather. + QueryOp::SetOp { inputs, .. } => inputs.iter().any(plan_contains_exchange), // Aggregate may carry a sub-plan input (catalog `ProviderScan`), // which could in principle nest an Exchange — recurse when present. QueryOp::Aggregate { input, .. } => { diff --git a/nodedb/src/control/gateway/version_set.rs b/nodedb/src/control/gateway/version_set.rs index cd5639ba5..fffb0cdb9 100644 --- a/nodedb/src/control/gateway/version_set.rs +++ b/nodedb/src/control/gateway/version_set.rs @@ -479,6 +479,13 @@ pub fn touched_collections(plan: &PhysicalPlan) -> Vec { out.extend(touched_collections(input)); } + // SetOp: every branch is a body that reads its own collections. + SetOp { inputs, .. } => { + for input in inputs { + out.extend(touched_collections(input)); + } + } + // ProviderScan is a catalog/constant source — no user collection. ProviderScan { .. } => {} diff --git a/nodedb/src/control/planner/redaction_refusal/plan.rs b/nodedb/src/control/planner/redaction_refusal/plan.rs index dd71f71eb..aab8e64ab 100644 --- a/nodedb/src/control/planner/redaction_refusal/plan.rs +++ b/nodedb/src/control/planner/redaction_refusal/plan.rs @@ -198,6 +198,8 @@ fn walk_query(op: &QueryOp, ctx: &RefusalCtx<'_>) -> crate::Result<()> { QueryOp::PostProcess { input, .. } => walk(input, ctx), + QueryOp::SetOp { inputs, .. } => inputs.iter().try_for_each(|input| walk(input, ctx)), + QueryOp::Aggregate { collection, input, @@ -461,6 +463,8 @@ mod tests { input: Box::new(aggregate_plan("users", vec![agg_spec("min", "ssn")])), filters: Vec::new(), projection: Vec::new(), + computed_columns: Vec::new(), + window_functions: Vec::new(), sort_keys: Vec::new(), limit: None, offset: 0, diff --git a/nodedb/src/control/planner/rls_injection/permission_tree/plan.rs b/nodedb/src/control/planner/rls_injection/permission_tree/plan.rs index 742d8716f..0ada83e52 100644 --- a/nodedb/src/control/planner/rls_injection/permission_tree/plan.rs +++ b/nodedb/src/control/planner/rls_injection/permission_tree/plan.rs @@ -414,6 +414,8 @@ mod tests { input: Box::new(columnar_scan("events")), filters: Vec::new(), projection: Vec::new(), + computed_columns: Vec::new(), + window_functions: Vec::new(), sort_keys: Vec::new(), limit: None, offset: 0, diff --git a/nodedb/src/control/planner/rls_injection/permission_tree/query.rs b/nodedb/src/control/planner/rls_injection/permission_tree/query.rs index 3667a3723..854a2e0e6 100644 --- a/nodedb/src/control/planner/rls_injection/permission_tree/query.rs +++ b/nodedb/src/control/planner/rls_injection/permission_tree/query.rs @@ -21,6 +21,10 @@ pub(super) fn apply_query(ctx: &PermCtx<'_>, op: &mut QueryOp) -> crate::Result< // tree restricts. QueryOp::PostProcess { input, .. } => walk(ctx, input), + // Recurse: every set-operation branch is its own body whose rows the + // policy restricts. + QueryOp::SetOp { inputs, .. } => inputs.iter_mut().try_for_each(|input| walk(ctx, input)), + // Filter and recurse: the aggregate handler evaluates `filters` // against both row sources — the per-shard collection scan and the // rows decoded from an embedded sub-plan — so the subtree filter goes diff --git a/nodedb/src/control/planner/rls_injection/plan.rs b/nodedb/src/control/planner/rls_injection/plan.rs index 57cf47739..671884b88 100644 --- a/nodedb/src/control/planner/rls_injection/plan.rs +++ b/nodedb/src/control/planner/rls_injection/plan.rs @@ -412,6 +412,8 @@ mod tests { input: Box::new(rag_fusion("docs")), filters: Vec::new(), projection: Vec::new(), + computed_columns: Vec::new(), + window_functions: Vec::new(), sort_keys: Vec::new(), limit: None, offset: 0, diff --git a/nodedb/src/control/planner/rls_injection/query.rs b/nodedb/src/control/planner/rls_injection/query.rs index 2c0bd05dd..5ba24f397 100644 --- a/nodedb/src/control/planner/rls_injection/query.rs +++ b/nodedb/src/control/planner/rls_injection/query.rs @@ -21,6 +21,10 @@ pub(super) fn inject_query(ctx: &RlsCtx<'_>, op: &mut QueryOp) -> crate::Result< // policy restricts. QueryOp::PostProcess { input, .. } => walk(ctx, input), + // Recurse: every set-operation branch is its own body whose rows the + // policy restricts. + QueryOp::SetOp { inputs, .. } => inputs.iter_mut().try_for_each(|input| walk(ctx, input)), + // Inject or recurse: a catalog aggregate (`input: Some`) sources rows // from the embedded sub-plan, so the policy belongs in that input // rather than in the aggregate's own (empty) filters. A legacy diff --git a/nodedb/src/control/planner/sql_plan_convert/aggregate/input_sourced.rs b/nodedb/src/control/planner/sql_plan_convert/aggregate/input_sourced.rs new file mode 100644 index 000000000..bdafd78c3 --- /dev/null +++ b/nodedb/src/control/planner/sql_plan_convert/aggregate/input_sourced.rs @@ -0,0 +1,84 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Aggregate over an input-sourced body: a derived table, constant result, or +//! union that carries no routing collection. The body lowers to one physical +//! plan, a sharded body is gathered, and one coordinator-local `Aggregate` +//! task runs over the materialized rows. + +use nodedb_sql::types::{AggregateExpr, Filter, SqlExpr, SqlPlan}; + +use crate::types::TenantId; +use nodedb_physical::physical_plan::*; +use nodedb_physical::physical_task::PhysicalTask; + +use super::super::body::convert_body_to_single_plan; +use super::super::convert::ConvertContext; +use super::super::filter::serialize_filters; +use super::spec::{InputSourcedTaskParams, build_input_sourced_aggregate_task}; + +pub(super) struct InputSourcedAggregateParams<'a> { + pub input: &'a SqlPlan, + pub group_by: &'a [SqlExpr], + pub aggregates: &'a [AggregateExpr], + pub having: &'a [Filter], + pub limit: usize, + pub grouping_sets: Option<&'a [Vec]>, + /// Post-aggregate sort keys, already lowered to bridge specs. + pub sort_keys: Vec, + pub tenant_id: TenantId, + pub ctx: &'a ConvertContext, +} + +/// Lower an aggregate whose input is a materialized relation. +/// +/// The body lowers to ONE relation through `convert_body_to_single_plan`: a +/// set-operation body becomes a coordinator-resolved `SetOp`, and a sharded +/// body is wrapped in `Exchange{Gather}` so the coordinator resolves it to a +/// `ProviderScan` before the aggregate runs. The emitted task is +/// coordinator-local: an empty collection keeps it on the coordinator vshard +/// and `is_sharded_source` reports the `Some(input)` aggregate as +/// non-sharded, so it runs once and is never broadcast. +pub(super) fn convert_input_sourced_aggregate( + p: InputSourcedAggregateParams<'_>, +) -> crate::Result> { + let InputSourcedAggregateParams { + input, + group_by, + aggregates, + having, + limit, + grouping_sets, + sort_keys, + tenant_id, + ctx, + } = p; + + // The input-sourced aggregate executor does not expand ROLLUP / CUBE / + // GROUPING SETS. A typed error beats a silent base-grouping-only answer. + if grouping_sets.is_some_and(|sets| !sets.is_empty()) { + return Err(crate::Error::PlanError { + detail: "ROLLUP / CUBE / GROUPING SETS over a derived-table body is not supported" + .to_string(), + }); + } + + // The body is ONE relation, already gathered when sharded, so the + // aggregate observes the full union exactly once. + let child = convert_body_to_single_plan(input, tenant_id, ctx)?; + + let having_bytes = serialize_filters(having)?; + + Ok(vec![build_input_sourced_aggregate_task( + InputSourcedTaskParams { + tenant_id, + ctx, + raw_collection: String::new(), + child, + group_by, + aggregates, + having_bytes, + limit, + sort_keys, + }, + )]) +} diff --git a/nodedb/src/control/planner/sql_plan_convert/aggregate/mod.rs b/nodedb/src/control/planner/sql_plan_convert/aggregate/mod.rs index ac20801bd..bdc156c37 100644 --- a/nodedb/src/control/planner/sql_plan_convert/aggregate/mod.rs +++ b/nodedb/src/control/planner/sql_plan_convert/aggregate/mod.rs @@ -4,11 +4,13 @@ //! //! Split by concern so each file stays under the project's hard size limit: //! `plan` (the `convert_aggregate` entry point and its join / catalog / -//! timeseries lowering), `spec` (aggregate-spec + collection/alias helpers and +//! timeseries lowering), `input_sourced` (aggregate over a materialized +//! derived-table body), `spec` (aggregate-spec + collection/alias helpers and //! join-side embedding), and `projection` (projection / computed-column / //! window-function serialization). mod cost; +mod input_sourced; mod plan; mod projection; mod spec; diff --git a/nodedb/src/control/planner/sql_plan_convert/aggregate/plan.rs b/nodedb/src/control/planner/sql_plan_convert/aggregate/plan.rs index 2f9c4c3bc..ae9ba6e9b 100644 --- a/nodedb/src/control/planner/sql_plan_convert/aggregate/plan.rs +++ b/nodedb/src/control/planner/sql_plan_convert/aggregate/plan.rs @@ -1,7 +1,14 @@ // SPDX-License-Identifier: BUSL-1.1 -//! The `convert_aggregate` entry point: join-sourced, catalog (input-sourced), -//! timeseries, and standard single-collection aggregate lowering. +//! The `convert_aggregate` entry point. Four lowering shapes: +//! +//! - join: `HashJoin` with post-join grouping +//! - input-sourced body (derived table, constant result, union): the body is +//! materialized and the aggregate runs over its rows on the coordinator +//! (`input_sourced.rs`) +//! - catalog: `ProviderScan` input, coordinator-local +//! - single collection: per-shard `Aggregate` (timeseries routes through +//! `TimeseriesOp::Scan`) use nodedb_sql::types::{EngineType, Filter, SortKey, SqlExpr, SqlPlan}; @@ -14,8 +21,9 @@ use super::super::convert::{ConvertContext, db_qualified}; use super::super::expr::convert_sort_keys; use super::super::filter::serialize_filters; use super::spec::{ - agg_expr_to_pair, agg_expr_to_spec, extract_collection_name, extract_scan_alias, - group_by_to_specs, group_by_to_strings, inline_join_side, join_side_collection, + InputSourcedTaskParams, agg_expr_to_pair, agg_expr_to_spec, build_input_sourced_aggregate_task, + extract_collection_name, extract_scan_alias, group_by_to_specs, group_by_to_strings, + inline_join_side, join_side_collection, }; use nodedb_sql::types::AggregateExpr; @@ -123,8 +131,23 @@ pub(in crate::control::planner::sql_plan_convert) fn convert_aggregate( }]); } - // Standard aggregate on a single collection. - let raw_collection = extract_collection_name(input); + // A body with no routing collection (derived table, constant result, + // union) is materialized and aggregated on the coordinator. + let Some(raw_collection) = extract_collection_name(input) else { + return super::input_sourced::convert_input_sourced_aggregate( + super::input_sourced::InputSourcedAggregateParams { + input, + group_by, + aggregates, + having, + limit, + grouping_sets, + sort_keys: bridge_sort_keys, + tenant_id, + ctx, + }, + ); + }; let (filters_ref, engine) = match input { SqlPlan::Scan { filters, engine, .. @@ -157,45 +180,33 @@ pub(in crate::control::planner::sql_plan_convert) fn convert_aggregate( ), }); } - let group_specs = group_by_to_specs(group_by); - let agg_specs: Vec = aggregates.iter().map(agg_expr_to_spec).collect(); let provider_scan = PhysicalPlan::Query(QueryOp::ProviderScan { provider: Some(raw_collection.clone()), rows: Vec::new(), // WHERE predicates on the catalog are applied by the ProviderScan // before the rows reach the aggregate. - filters: filter_bytes.clone(), + filters: filter_bytes, projection: Vec::new(), + computed_columns: Vec::new(), + window_functions: Vec::new(), sort_keys: Vec::new(), limit: None, offset: 0, distinct: false, }); - return Ok(vec![PhysicalTask { - tenant_id, - // Coordinator-local: empty collection keeps the task on the - // coordinator vshard (catalog rows are not per-shard). - vshard_id: VShardId::from_collection_in_database(ctx.database_id, ""), - database_id: ctx.database_id, - plan: PhysicalPlan::Query(QueryOp::Aggregate { - collection: nodedb_types::QualifiedCollection::from_stored(raw_collection), - input: Some(Box::new(provider_scan)), - group_by: group_specs, - aggregates: agg_specs, - // Filters live on the ProviderScan input; the aggregate node - // applies none of its own over the already-filtered rows. - filters: Vec::new(), - having: having_bytes, + return Ok(vec![build_input_sourced_aggregate_task( + InputSourcedTaskParams { + tenant_id, + ctx, + raw_collection, + child: provider_scan, + group_by, + aggregates, + having_bytes, limit, - sub_group_by: Vec::new(), - sub_aggregates: Vec::new(), - // Guarded above: catalog aggregates never carry grouping sets. - grouping_sets: Vec::new(), sort_keys: bridge_sort_keys, - }), - post_set_op: PostSetOp::None, - txn_id: None, - }]); + }, + )]); } let collection = db_qualified(ctx.database_id, &raw_collection); @@ -217,7 +228,7 @@ pub(in crate::control::planner::sql_plan_convert) fn convert_aggregate( collection: qualified_collection, // Derived in the Data Plane against the declared TIME_KEY. time_range: UNBOUNDED_TIME_RANGE, - sort_keys: bridge_sort_keys.clone(), + sort_keys: bridge_sort_keys, projection: Vec::new(), limit, filters: filter_bytes, @@ -399,8 +410,8 @@ mod tests { vec!["name".to_string(), "rn".to_string()] ); - let computed_bytes = - extract_computed_columns(&projection, &window_functions).expect("serialize computed"); + let computed_bytes = extract_computed_columns(&projection, &window_functions, false) + .expect("serialize computed"); let computed: Vec = zerompk::from_msgpack(&computed_bytes).expect("deserialize computed"); diff --git a/nodedb/src/control/planner/sql_plan_convert/aggregate/projection.rs b/nodedb/src/control/planner/sql_plan_convert/aggregate/projection.rs index 2591a3d9f..bdb5152a5 100644 --- a/nodedb/src/control/planner/sql_plan_convert/aggregate/projection.rs +++ b/nodedb/src/control/planner/sql_plan_convert/aggregate/projection.rs @@ -125,10 +125,28 @@ fn encode_computed_columns( }) } +/// Pick the bridge-expression converter for the row shape the expression runs +/// against. A join / lateral body emits merged documents keyed by `t.col`, so +/// its column references keep their table qualifier. +fn bridge_expr_converter(qualified: bool) -> fn(&SqlExpr) -> crate::bridge::expr_eval::SqlExpr { + if qualified { + sql_expr_to_bridge_expr_qualified + } else { + sql_expr_to_bridge_expr + } +} + +/// Serialize the computed (non-window) projection entries. +/// +/// `qualified` selects the column-key convention of the rows the columns are +/// evaluated over: `true` for a join / lateral merged document, `false` for a +/// single-collection row. pub(in crate::control::planner::sql_plan_convert) fn extract_computed_columns( proj: &[Projection], window_functions: &[WindowSpec], + qualified: bool, ) -> crate::Result> { + let convert = bridge_expr_converter(qualified); let computed: Vec = proj .iter() .filter_map(|p| match p { @@ -137,7 +155,7 @@ pub(in crate::control::planner::sql_plan_convert) fn extract_computed_columns( { Some(crate::bridge::expr_eval::ComputedColumn { alias: alias.clone(), - expr: sql_expr_to_bridge_expr(expr), + expr: convert(expr), }) } _ => None, @@ -151,23 +169,27 @@ pub(in crate::control::planner::sql_plan_convert) fn extract_computed_columns( }) } +/// Serialize window specs. `qualified` follows the same convention as +/// [`extract_computed_columns`]. pub(in crate::control::planner::sql_plan_convert) fn serialize_window_functions( - specs: &[nodedb_sql::types::WindowSpec], + specs: &[WindowSpec], + qualified: bool, ) -> crate::Result> { if specs.is_empty() { return Ok(Vec::new()); } + let convert = bridge_expr_converter(qualified); let bridge_specs: Vec = specs .iter() .map(|s| crate::bridge::window_func::WindowFuncSpec { alias: s.alias.clone(), func_name: s.function.clone(), - args: s.args.iter().map(sql_expr_to_bridge_expr).collect(), - partition_by: s.partition_by.iter().map(sql_expr_to_bridge_expr).collect(), + args: s.args.iter().map(convert).collect(), + partition_by: s.partition_by.iter().map(convert).collect(), order_by: s .order_by .iter() - .map(|k| (sql_expr_to_bridge_expr(&k.expr), k.ascending)) + .map(|k| (convert(&k.expr), k.ascending)) .collect(), frame: s.frame.clone(), }) diff --git a/nodedb/src/control/planner/sql_plan_convert/aggregate/spec.rs b/nodedb/src/control/planner/sql_plan_convert/aggregate/spec.rs index e51e3a898..65f4497f4 100644 --- a/nodedb/src/control/planner/sql_plan_convert/aggregate/spec.rs +++ b/nodedb/src/control/planner/sql_plan_convert/aggregate/spec.rs @@ -6,8 +6,9 @@ use nodedb_sql::types::{AggregateExpr, SqlExpr, SqlPlan}; use crate::bridge::envelope::PhysicalPlan; -use crate::types::TenantId; +use crate::types::{TenantId, VShardId}; use nodedb_physical::physical_plan::*; +use nodedb_physical::physical_task::{PhysicalTask, PostSetOp}; use super::super::convert::{ConvertContext, convert_one, db_qualified}; use super::super::expr::sql_expr_to_bridge_expr; @@ -86,26 +87,31 @@ pub(in crate::control::planner::sql_plan_convert) fn join_side_collection( plan: &SqlPlan, database_id: crate::types::DatabaseId, ) -> String { - let raw = extract_collection_name(plan); - if scan_is_catalog(&raw) { - String::new() - } else { - db_qualified(database_id, &raw) + // An input-sourced join side carries no routing collection; its rows come + // from `left_input` / `right_input`. + match extract_collection_name(plan) { + Some(raw) if !scan_is_catalog(&raw) => db_qualified(database_id, &raw), + Some(_) | None => String::new(), } } +/// Routing collection for a scan-shaped input. `None` for an input-sourced +/// (materialized) body. +/// +/// A `Join` yields its left side's collection as the routing hint for a join +/// over scans. An `Aggregate` input is a materialized relation, not a scan. pub(in crate::control::planner::sql_plan_convert) fn extract_collection_name( plan: &SqlPlan, -) -> String { +) -> Option { match plan { - SqlPlan::Scan { collection, .. } => collection.clone(), - SqlPlan::PointGet { collection, .. } => collection.clone(), + SqlPlan::Scan { collection, .. } => Some(collection.clone()), + SqlPlan::PointGet { collection, .. } => Some(collection.clone()), SqlPlan::Join { left, .. } => extract_collection_name(left), - SqlPlan::Aggregate { input, .. } => extract_collection_name(input), - _ => String::new(), + _ => None, } } +/// Scan alias for a scan-shaped input. `None` for an input-sourced body. pub(in crate::control::planner::sql_plan_convert) fn extract_scan_alias( plan: &SqlPlan, ) -> Option { @@ -113,7 +119,6 @@ pub(in crate::control::planner::sql_plan_convert) fn extract_scan_alias( SqlPlan::Scan { alias, .. } => alias.clone(), SqlPlan::PointGet { alias, .. } => alias.clone(), SqlPlan::Join { left, .. } => extract_scan_alias(left), - SqlPlan::Aggregate { input, .. } => extract_scan_alias(input), _ => None, } } @@ -197,6 +202,70 @@ pub(super) fn group_by_to_strings(exprs: &[SqlExpr]) -> Vec { .collect() } +/// Build the coordinator-local `Aggregate{input: Some(child)}` task shared by +/// the catalog and derived-table-body lowerings: an aggregate that runs once +/// over an already-materialized `child` plan instead of scanning a per-shard +/// collection. `raw_collection` is the RAW catalog source name for a catalog +/// body, or `String::new()` for a body with no routing collection (derived +/// table, constant result, union) — either way the task's vshard is the +/// coordinator's empty-collection vshard, so it is never broadcast. +/// Inputs to [`build_input_sourced_aggregate_task`]. +pub(in crate::control::planner::sql_plan_convert) struct InputSourcedTaskParams<'a> { + pub tenant_id: TenantId, + pub ctx: &'a ConvertContext, + pub raw_collection: String, + pub child: PhysicalPlan, + pub group_by: &'a [SqlExpr], + pub aggregates: &'a [AggregateExpr], + pub having_bytes: Vec, + pub limit: usize, + pub sort_keys: Vec, +} + +pub(in crate::control::planner::sql_plan_convert) fn build_input_sourced_aggregate_task( + p: InputSourcedTaskParams<'_>, +) -> PhysicalTask { + let InputSourcedTaskParams { + tenant_id, + ctx, + raw_collection, + child, + group_by, + aggregates, + having_bytes, + limit, + sort_keys, + } = p; + let group_specs = group_by_to_specs(group_by); + let agg_specs: Vec = aggregates.iter().map(agg_expr_to_spec).collect(); + PhysicalTask { + tenant_id, + // Coordinator-local: empty collection keeps the task on the + // coordinator vshard (the child's rows are not per-shard). + vshard_id: VShardId::from_collection_in_database(ctx.database_id, ""), + database_id: ctx.database_id, + plan: PhysicalPlan::Query(QueryOp::Aggregate { + collection: nodedb_types::QualifiedCollection::from_stored(raw_collection), + input: Some(Box::new(child)), + group_by: group_specs, + aggregates: agg_specs, + // The child carries its own WHERE / ProviderScan filters; the + // aggregate node applies none of its own. + filters: Vec::new(), + having: having_bytes, + limit, + sub_group_by: Vec::new(), + sub_aggregates: Vec::new(), + // Guarded by the caller: input-sourced aggregates never carry + // grouping sets. + grouping_sets: Vec::new(), + sort_keys, + }), + post_set_op: PostSetOp::None, + txn_id: None, + } +} + /// Lower GROUP BY expressions to Data-Plane group-key specs. /// /// A bare `Column` key extracts from, and is emitted under, its own column diff --git a/nodedb/src/control/planner/sql_plan_convert/body.rs b/nodedb/src/control/planner/sql_plan_convert/body.rs new file mode 100644 index 000000000..38720030c --- /dev/null +++ b/nodedb/src/control/planner/sql_plan_convert/body.rs @@ -0,0 +1,241 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Derived-table body lowering: one `SqlPlan` body to ONE physical relation, +//! for a post-processor or an input-sourced aggregate. + +use nodedb_sql::types::SqlPlan; + +use crate::bridge::envelope::PhysicalPlan; +use crate::types::TenantId; +use nodedb_physical::physical_plan::{ExchangeMode, ExchangeOp, QueryOp, SetOpKind}; + +use super::convert::{ConvertContext, convert_one}; + +/// Lower a derived-table body to ONE physical relation for a post-processor +/// or an input-sourced aggregate. +/// +/// A set-operation body lowers to a coordinator-resolved `QueryOp::SetOp` +/// whose branches recurse through this function, so nested set operations +/// (`(a UNION b) INTERSECT c`) stay one relation. Every other body lowers +/// through `convert_one` and must yield exactly one task. +/// +/// A sharded relation (the body itself, or any set-operation branch) is +/// wrapped in `Exchange{Gather}` so its gather runs exactly once over the +/// full union before the enclosing tail or merge observes it. `SetOp` is +/// coordinator-local, so it is never wrapped itself. +pub(super) fn convert_body_to_single_plan( + input: &SqlPlan, + tenant_id: TenantId, + ctx: &ConvertContext, +) -> crate::Result { + match input { + SqlPlan::Union { inputs, distinct } => { + let op = if *distinct { + SetOpKind::UnionDistinct + } else { + SetOpKind::UnionAll + }; + let inputs = inputs + .iter() + .map(|branch| convert_body_to_single_plan(branch, tenant_id, ctx)) + .collect::>>()?; + Ok(PhysicalPlan::Query(QueryOp::SetOp { inputs, op })) + } + SqlPlan::Intersect { left, right, all } => { + let op = if *all { + SetOpKind::IntersectAll + } else { + SetOpKind::Intersect + }; + convert_binary_set_op(left, right, op, tenant_id, ctx) + } + SqlPlan::Except { left, right, all } => { + let op = if *all { + SetOpKind::ExceptAll + } else { + SetOpKind::Except + }; + convert_binary_set_op(left, right, op, tenant_id, ctx) + } + // Every other body lowers through the ordinary converter. Listed in + // full so a new `SqlPlan` variant forces a decision here. + SqlPlan::ConstantResult { .. } + | SqlPlan::Scan { .. } + | SqlPlan::PointGet { .. } + | SqlPlan::DocumentIndexLookup { .. } + | SqlPlan::RangeScan { .. } + | SqlPlan::Insert { .. } + | SqlPlan::KvInsert { .. } + | SqlPlan::Upsert { .. } + | SqlPlan::InsertSelect { .. } + | SqlPlan::Update { .. } + | SqlPlan::UpdateFrom { .. } + | SqlPlan::Delete { .. } + | SqlPlan::Truncate { .. } + | SqlPlan::Join { .. } + | SqlPlan::Aggregate { .. } + | SqlPlan::TimeseriesScan { .. } + | SqlPlan::TimeseriesIngest { .. } + | SqlPlan::VectorSearch { .. } + | SqlPlan::MultiVectorSearch { .. } + | SqlPlan::SparseSearch { .. } + | SqlPlan::TextSearch { .. } + | SqlPlan::HybridSearch { .. } + | SqlPlan::HybridSearchTriple { .. } + | SqlPlan::SpatialScan { .. } + | SqlPlan::RecursiveScan { .. } + | SqlPlan::RecursiveValue { .. } + | SqlPlan::Cte { .. } + | SqlPlan::Subquery { .. } + | SqlPlan::CreateArray { .. } + | SqlPlan::DropArray { .. } + | SqlPlan::AlterArray { .. } + | SqlPlan::InsertArray { .. } + | SqlPlan::DeleteArray { .. } + | SqlPlan::ArraySlice { .. } + | SqlPlan::ArrayProject { .. } + | SqlPlan::ArrayAgg { .. } + | SqlPlan::ArrayElementwise { .. } + | SqlPlan::ArrayFlush { .. } + | SqlPlan::ArrayCompact { .. } + | SqlPlan::Merge { .. } + | SqlPlan::LateralTopK { .. } + | SqlPlan::LateralLoop { .. } + | SqlPlan::VectorPrimaryInsert { .. } + | SqlPlan::CreateIndex { .. } + | SqlPlan::DropIndex { .. } => { + let mut tasks = convert_one(input, tenant_id, ctx)?; + let plan = match (tasks.len(), tasks.pop()) { + (1, Some(task)) => task.plan, + (n, _) => { + return Err(crate::Error::PlanError { + detail: format!( + "derived-table body lowers to {n} physical tasks; the body must \ + produce a single relation" + ), + }); + } + }; + Ok(gather_if_sharded(plan)) + } + } +} + +/// Lower the two sides of `INTERSECT` / `EXCEPT` to a two-branch `SetOp`. +fn convert_binary_set_op( + left: &SqlPlan, + right: &SqlPlan, + op: SetOpKind, + tenant_id: TenantId, + ctx: &ConvertContext, +) -> crate::Result { + let inputs = vec![ + convert_body_to_single_plan(left, tenant_id, ctx)?, + convert_body_to_single_plan(right, tenant_id, ctx)?, + ]; + Ok(PhysicalPlan::Query(QueryOp::SetOp { inputs, op })) +} + +/// Wrap a sharded relation in `Exchange{Gather}` so its gather runs exactly +/// once over the full union. The enclosing node is coordinator-local, so the +/// top-level `convert()` wrap loop does not gather it. +fn gather_if_sharded(plan: PhysicalPlan) -> PhysicalPlan { + if !plan.is_sharded_source() { + return plan; + } + let as_aggregate = matches!( + &plan, + PhysicalPlan::Query(QueryOp::Aggregate { .. }) + | PhysicalPlan::Query(QueryOp::PartialAggregate { .. }) + ); + PhysicalPlan::Query(QueryOp::Exchange(ExchangeOp { + child: Box::new(plan), + mode: ExchangeMode::Gather { as_aggregate }, + })) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::control::planner::sql_plan_convert::PlanningPurpose; + use nodedb_sql::types::SqlValue; + + fn ctx() -> ConvertContext { + ConvertContext { + purpose: PlanningPurpose::Execute, + retention_registry: None, + array_catalog: None, + credentials: None, + wal: None, + surrogate_assigner: None, + cluster_enabled: false, + bitemporal_retention_registry: None, + max_vector_dim: 0, + force_shuffle_join: false, + shuffle_num_parts: 0, + force_shuffle_agg: false, + shuffle_agg_num_parts: 0, + broadcast_threshold_bytes: 8 * 1024 * 1024, + shuffle_agg_threshold: 10_000, + database_id: crate::types::DatabaseId::DEFAULT, + tenant_id: crate::types::TenantId::new(0), + } + } + + fn constant(x: i64) -> SqlPlan { + SqlPlan::ConstantResult { + columns: vec!["x".into()], + values: vec![SqlValue::Int(x)], + volatile: false, + } + } + + #[test] + fn union_all_body_lowers_to_one_set_op() { + let body = SqlPlan::Union { + inputs: vec![constant(1), constant(2)], + distinct: false, + }; + let plan = convert_body_to_single_plan(&body, TenantId::new(1), &ctx()) + .expect("union body lowers"); + match plan { + PhysicalPlan::Query(QueryOp::SetOp { inputs, op }) => { + assert_eq!(op, SetOpKind::UnionAll); + assert_eq!(inputs.len(), 2); + for input in &inputs { + assert!(matches!( + input, + PhysicalPlan::Query(QueryOp::ProviderScan { provider: None, .. }) + )); + } + } + other => panic!("expected SetOp, got {other:?}"), + } + } + + #[test] + fn nested_set_ops_stay_one_relation() { + let body = SqlPlan::Intersect { + left: Box::new(SqlPlan::Union { + inputs: vec![constant(1), constant(2)], + distinct: true, + }), + right: Box::new(constant(2)), + all: false, + }; + let plan = convert_body_to_single_plan(&body, TenantId::new(1), &ctx()) + .expect("nested set-op body lowers"); + let PhysicalPlan::Query(QueryOp::SetOp { inputs, op }) = plan else { + panic!("expected SetOp"); + }; + assert_eq!(op, SetOpKind::Intersect); + assert!(matches!( + &inputs[0], + PhysicalPlan::Query(QueryOp::SetOp { + op: SetOpKind::UnionDistinct, + .. + }) + )); + assert!(!PhysicalPlan::Query(QueryOp::SetOp { inputs, op }).is_sharded_source()); + } +} diff --git a/nodedb/src/control/planner/sql_plan_convert/convert.rs b/nodedb/src/control/planner/sql_plan_convert/convert.rs index 1b38821ea..03a901dfc 100644 --- a/nodedb/src/control/planner/sql_plan_convert/convert.rs +++ b/nodedb/src/control/planner/sql_plan_convert/convert.rs @@ -251,6 +251,8 @@ pub fn convert( rows: Vec::new(), filters: Vec::new(), projection: Vec::new(), + computed_columns: Vec::new(), + window_functions: Vec::new(), sort_keys: Vec::new(), limit: None, offset: 0, diff --git a/nodedb/src/control/planner/sql_plan_convert/expr.rs b/nodedb/src/control/planner/sql_plan_convert/expr.rs deleted file mode 100644 index 80119f4ff..000000000 --- a/nodedb/src/control/planner/sql_plan_convert/expr.rs +++ /dev/null @@ -1,716 +0,0 @@ -// SPDX-License-Identifier: BUSL-1.1 - -//! Expression conversion and CTE inlining. - -use nodedb_physical::physical_plan::SortKeySpec; -use nodedb_sql::types::{SortKey, SqlExpr, SqlPlan}; - -use super::value::sql_value_to_nodedb_value; - -/// Convert a `nodedb_sql::types::SqlExpr` (parser AST) to a -/// `nodedb_query::expr::SqlExpr` (bridge evaluation type). -/// -/// Column references use the **bare** name (no table qualifier) for -/// single-collection evaluation contexts (WHERE, CHECK, GENERATED). -/// For join contexts where the merged document uses qualified keys -/// (`"t1.col"`), use [`sql_expr_to_bridge_expr_qualified`] instead. -pub(super) fn sql_expr_to_bridge_expr(expr: &SqlExpr) -> crate::bridge::expr_eval::SqlExpr { - convert_expr_inner(expr, false) -} - -/// Like [`sql_expr_to_bridge_expr`] but qualifies column references -/// with their table name (`t.col` → `"t.col"`) for join merged docs. -pub(super) fn sql_expr_to_bridge_expr_qualified( - expr: &SqlExpr, -) -> crate::bridge::expr_eval::SqlExpr { - convert_expr_inner(expr, true) -} - -fn convert_expr_inner(expr: &SqlExpr, qualify: bool) -> crate::bridge::expr_eval::SqlExpr { - use crate::bridge::expr_eval::SqlExpr as BExpr; - match expr { - SqlExpr::Column { table, name } => { - // `EXCLUDED.col` references the row proposed for insertion in - // `INSERT ... ON CONFLICT DO UPDATE`. Emit the dedicated - // variant so the upsert handler can resolve against the - // incoming row via `eval_with_excluded`. The table qualifier - // comes in already-normalized (lowercased) from the parser. - if table - .as_deref() - .is_some_and(|t| t.eq_ignore_ascii_case("excluded")) - { - return BExpr::ExcludedColumn(name.clone()); - } - if qualify { - BExpr::Column(nodedb_sql::planner::qualified_name(table.as_deref(), name)) - } else { - BExpr::Column(name.clone()) - } - } - SqlExpr::Literal(v) => BExpr::Literal(sql_value_to_nodedb_value(v)), - SqlExpr::BinaryOp { left, op, right } => BExpr::BinaryOp { - left: Box::new(convert_expr_inner(left, qualify)), - op: match op { - nodedb_sql::types::BinaryOp::Add => crate::bridge::expr_eval::BinaryOp::Add, - nodedb_sql::types::BinaryOp::Sub => crate::bridge::expr_eval::BinaryOp::Sub, - nodedb_sql::types::BinaryOp::Mul => crate::bridge::expr_eval::BinaryOp::Mul, - nodedb_sql::types::BinaryOp::Div => crate::bridge::expr_eval::BinaryOp::Div, - nodedb_sql::types::BinaryOp::Mod => crate::bridge::expr_eval::BinaryOp::Mod, - nodedb_sql::types::BinaryOp::Eq => crate::bridge::expr_eval::BinaryOp::Eq, - nodedb_sql::types::BinaryOp::Ne => crate::bridge::expr_eval::BinaryOp::NotEq, - nodedb_sql::types::BinaryOp::Gt => crate::bridge::expr_eval::BinaryOp::Gt, - nodedb_sql::types::BinaryOp::Ge => crate::bridge::expr_eval::BinaryOp::GtEq, - nodedb_sql::types::BinaryOp::Lt => crate::bridge::expr_eval::BinaryOp::Lt, - nodedb_sql::types::BinaryOp::Le => crate::bridge::expr_eval::BinaryOp::LtEq, - nodedb_sql::types::BinaryOp::And => crate::bridge::expr_eval::BinaryOp::And, - nodedb_sql::types::BinaryOp::Or => crate::bridge::expr_eval::BinaryOp::Or, - nodedb_sql::types::BinaryOp::Concat => crate::bridge::expr_eval::BinaryOp::Concat, - }, - right: Box::new(convert_expr_inner(right, qualify)), - }, - SqlExpr::Function { name, args, .. } => BExpr::Function { - name: name.clone(), - args: args - .iter() - .map(|a| convert_expr_inner(a, qualify)) - .collect(), - }, - SqlExpr::Case { - operand, - when_then, - else_expr, - } => BExpr::Case { - operand: operand - .as_ref() - .map(|e| Box::new(convert_expr_inner(e, qualify))), - when_thens: when_then - .iter() - .map(|(w, t)| { - ( - convert_expr_inner(w, qualify), - convert_expr_inner(t, qualify), - ) - }) - .collect(), - else_expr: else_expr - .as_ref() - .map(|e| Box::new(convert_expr_inner(e, qualify))), - }, - SqlExpr::Cast { expr, to_type } => { - let cast_type = match to_type.to_uppercase().as_str() { - "INT" | "INTEGER" | "BIGINT" | "SMALLINT" => { - crate::bridge::expr_eval::CastType::Int - } - "FLOAT" | "DOUBLE" | "REAL" | "NUMERIC" | "DECIMAL" => { - crate::bridge::expr_eval::CastType::Float - } - "BOOL" | "BOOLEAN" => crate::bridge::expr_eval::CastType::Bool, - _ => crate::bridge::expr_eval::CastType::String, - }; - BExpr::Cast { - expr: Box::new(convert_expr_inner(expr, qualify)), - to_type: cast_type, - } - } - SqlExpr::Wildcard => BExpr::Column("*".into()), - - // NOT e / -e → evaluator's Negate (handles both bool and numeric). - SqlExpr::UnaryOp { expr, .. } => BExpr::Negate(Box::new(convert_expr_inner(expr, qualify))), - - // `e IS NULL` / `e IS NOT NULL` — direct passthrough. - SqlExpr::IsNull { expr, negated } => BExpr::IsNull { - expr: Box::new(convert_expr_inner(expr, qualify)), - negated: *negated, - }, - - // `e BETWEEN low AND high` desugars to `e >= low AND e <= high` - // (or `e < low OR e > high` when negated). The evaluator has no - // native Between variant, so the planner must lower it here. - SqlExpr::Between { - expr, - low, - high, - negated, - } => { - let e = convert_expr_inner(expr, qualify); - let l = convert_expr_inner(low, qualify); - let h = convert_expr_inner(high, qualify); - if *negated { - let lt = BExpr::BinaryOp { - left: Box::new(e.clone()), - op: crate::bridge::expr_eval::BinaryOp::Lt, - right: Box::new(l), - }; - let gt = BExpr::BinaryOp { - left: Box::new(e), - op: crate::bridge::expr_eval::BinaryOp::Gt, - right: Box::new(h), - }; - BExpr::BinaryOp { - left: Box::new(lt), - op: crate::bridge::expr_eval::BinaryOp::Or, - right: Box::new(gt), - } - } else { - let ge = BExpr::BinaryOp { - left: Box::new(e.clone()), - op: crate::bridge::expr_eval::BinaryOp::GtEq, - right: Box::new(l), - }; - let le = BExpr::BinaryOp { - left: Box::new(e), - op: crate::bridge::expr_eval::BinaryOp::LtEq, - right: Box::new(h), - }; - BExpr::BinaryOp { - left: Box::new(ge), - op: crate::bridge::expr_eval::BinaryOp::And, - right: Box::new(le), - } - } - } - - // `e IN (a, b, c)` desugars to `e = a OR e = b OR e = c` — each - // element may itself be a non-literal expression, so we must - // recursively convert and OR the comparisons together. `NOT IN` - // is `e <> a AND e <> b AND e <> c`. - SqlExpr::InList { - expr, - list, - negated, - } => { - let target = convert_expr_inner(expr, qualify); - if list.is_empty() { - // Empty list: `e IN ()` = false, `e NOT IN ()` = true. - return BExpr::Literal(nodedb_types::Value::Bool(*negated)); - } - let (eq_op, combine_op) = if *negated { - ( - crate::bridge::expr_eval::BinaryOp::NotEq, - crate::bridge::expr_eval::BinaryOp::And, - ) - } else { - ( - crate::bridge::expr_eval::BinaryOp::Eq, - crate::bridge::expr_eval::BinaryOp::Or, - ) - }; - // Empty list is handled above, so `list` is guaranteed non-empty - // here: we reduce `(target eq list[0]) op (target eq list[1]) op ...` - // without touching `.unwrap()` or `.expect()`. - list.iter() - .map(|item| BExpr::BinaryOp { - left: Box::new(target.clone()), - op: eq_op, - right: Box::new(convert_expr_inner(item, qualify)), - }) - .reduce(|acc, next| BExpr::BinaryOp { - left: Box::new(acc), - op: combine_op, - right: Box::new(next), - }) - // Unreachable: `list.is_empty()` returns early above. - .unwrap_or(BExpr::Literal(nodedb_types::Value::Bool(*negated))) - } - - // `e LIKE pattern` — no direct evaluator variant; route through a - // function call so the shared function dispatcher handles it. - SqlExpr::Like { - expr, - pattern, - negated, - case_insensitive, - } => { - let fn_name = if *case_insensitive { "ilike" } else { "like" }; - let call = BExpr::Function { - name: fn_name.into(), - args: vec![ - convert_expr_inner(expr, qualify), - convert_expr_inner(pattern, qualify), - ], - }; - if *negated { - BExpr::Negate(Box::new(call)) - } else { - call - } - } - - // `ARRAY['a', 'b', ...]` — lower each element and, when all resolve to - // `BExpr::Literal`, fold into a single `Value::Array` literal so that - // functions like `pg_json_has_any_key` / `pg_json_has_all_keys` receive - // a proper `Value::Array` argument rather than `Value::Null`. - SqlExpr::ArrayLiteral(elems) => { - let mut values = Vec::with_capacity(elems.len()); - let mut all_literal = true; - for elem in elems { - match convert_expr_inner(elem, qualify) { - BExpr::Literal(v) => values.push(v), - other => { - all_literal = false; - // Non-literal element: fall back to Null for that slot. - let _ = other; - values.push(nodedb_types::Value::Null); - } - } - } - if all_literal { - BExpr::Literal(nodedb_types::Value::Array(values)) - } else { - BExpr::Literal(nodedb_types::Value::Null) - } - } - - _ => BExpr::Literal(nodedb_types::Value::Null), - } -} - -/// Lower planner sort keys to their physical form. -/// -/// Every key is carried, expression and all. Dropping a key the Data Plane -/// cannot name as a stored column would silently answer -/// `ORDER BY 100 / weight` with rows in storage order. -pub(super) fn convert_sort_keys(keys: &[SortKey]) -> Vec { - keys.iter() - .map(|k| SortKeySpec { - expr: sql_expr_to_bridge_expr(&k.expr), - ascending: k.ascending, - nulls_first: k.nulls_first, - }) - .collect() -} - -/// Replace scans on `cte_name` with the CTE's actual subquery plan. -/// -/// Outer constraints on the CTE reference are merged onto the CTE body as far -/// as the body can carry them: a `Scan` body takes all of them; a -/// `VectorSearch` body takes filters, projection, and an unordered LIMIT (as -/// `top_k`). Constraints a body has no slot for — an outer `ORDER BY`, OFFSET, -/// or DISTINCT over a non-`Scan` body — are not applied. -pub(super) fn inline_cte(plan: &SqlPlan, cte_name: &str, cte_plan: &SqlPlan) -> SqlPlan { - match plan { - // Direct scan on CTE name → replace with CTE plan. - SqlPlan::Scan { - collection, - filters, - projection, - sort_keys, - limit, - offset, - distinct, - .. - } if collection == cte_name => { - // If the outer query adds filters/sort/limit, wrap the CTE plan. - // For simple SELECT * FROM cte, just return the CTE plan directly. - if filters.is_empty() - && sort_keys.is_empty() - && limit.is_none() - && !distinct - && projection.is_empty() - { - cte_plan.clone() - } else { - // Merge outer constraints onto the CTE plan if it's also a Scan. - if let SqlPlan::Scan { - collection: inner_col, - alias: inner_alias, - engine: inner_eng, - filters: inner_f, - projection: inner_p, - sort_keys: inner_s, - limit: inner_l, - offset: inner_o, - distinct: inner_d, - window_functions: inner_w, - temporal: inner_t, - } = cte_plan - { - let mut merged_filters = inner_f.clone(); - merged_filters.extend(filters.iter().cloned()); - SqlPlan::Scan { - collection: inner_col.clone(), - alias: inner_alias.clone(), - engine: *inner_eng, - filters: merged_filters, - // Outer projection overrides inner; empty means "inherit from CTE". - projection: if projection.is_empty() { - inner_p.clone() - } else { - projection.clone() - }, - sort_keys: if sort_keys.is_empty() { - inner_s.clone() - } else { - sort_keys.clone() - }, - limit: limit.or(*inner_l), - // offset 0 = unspecified → inherit CTE's offset. - offset: if *offset > 0 { *offset } else { *inner_o }, - distinct: *distinct || *inner_d, - window_functions: inner_w.clone(), - temporal: *inner_t, - } - } else if let SqlPlan::VectorSearch { .. } = cte_plan { - // A k-NN body carries its own post-filter list and top-k. An - // outer `WHERE` merges into the engine post-filter so the cut - // counts MATCHING rows, and — when nothing reorders the - // result — an unordered `LIMIT` folds into `top_k` and the - // projection rides along. An outer `ORDER BY` / `OFFSET` / - // `DISTINCT` reorders the k rows, which the search leaf has no - // slot for; those (and a `LIMIT` that must apply after the - // reorder) run in a `Subquery` post-processor over the k rows. - let needs_reorder = !sort_keys.is_empty() || *offset > 0 || *distinct; - let mut leaf = cte_plan.clone(); - if let SqlPlan::VectorSearch { - filters: body_filters, - projection: body_projection, - top_k, - .. - } = &mut leaf - { - body_filters.extend(filters.iter().cloned()); - if !needs_reorder { - if !projection.is_empty() { - body_projection.clone_from(projection); - } - if let Some(outer_limit) = limit { - *top_k = (*top_k).min(*outer_limit); - } - } - } - if needs_reorder { - // Filters already run in the engine; the tail applies the - // reorder-dependent constraints over the k rows. It sorts - // before projecting, so ORDER BY may reference any column. - SqlPlan::Subquery { - input: Box::new(leaf), - filters: Vec::new(), - projection: projection.clone(), - sort_keys: sort_keys.clone(), - offset: *offset, - distinct: *distinct, - limit: *limit, - } - } else { - leaf - } - } else if filters.is_empty() - && sort_keys.is_empty() - && *offset == 0 - && !*distinct - && limit.is_none() - { - // Any other non-`Scan` body (Aggregate, Join, TextSearch, - // HybridSearch, SparseSearch, SpatialScan, MultiVectorSearch, - // ...) with only an outer projection: the response boundary - // projects by output schema, so no post-processor is needed. - cte_plan.clone() - } else { - // The body has no slot for these outer constraints. Apply - // them over its materialized rows in a `Subquery` - // post-processor — previously they were silently dropped. - SqlPlan::Subquery { - input: Box::new(cte_plan.clone()), - filters: filters.clone(), - projection: projection.clone(), - sort_keys: sort_keys.clone(), - offset: *offset, - distinct: *distinct, - limit: *limit, - } - } - } - } - - // Aggregate referencing CTE → inline into the input. - SqlPlan::Aggregate { - input, - group_by, - group_by_aliases, - output_order, - aggregates, - having, - limit, - grouping_sets, - sort_keys, - } => SqlPlan::Aggregate { - input: Box::new(inline_cte(input, cte_name, cte_plan)), - group_by: group_by.clone(), - group_by_aliases: group_by_aliases.clone(), - output_order: output_order.clone(), - aggregates: aggregates.clone(), - having: having.clone(), - limit: *limit, - grouping_sets: grouping_sets.clone(), - sort_keys: sort_keys.clone(), - }, - - // JOIN referencing CTE on either side. - SqlPlan::Join { - left, - right, - on, - join_type, - condition, - limit, - projection, - filters, - } => SqlPlan::Join { - left: Box::new(inline_cte(left, cte_name, cte_plan)), - right: Box::new(inline_cte(right, cte_name, cte_plan)), - on: on.clone(), - join_type: *join_type, - condition: condition.clone(), - limit: *limit, - projection: projection.clone(), - filters: filters.clone(), - }, - - // Union referencing CTE → inline into all inputs. - SqlPlan::Union { inputs, distinct } => SqlPlan::Union { - inputs: inputs - .iter() - .map(|i| inline_cte(i, cte_name, cte_plan)) - .collect(), - distinct: *distinct, - }, - - // Intersect referencing CTE → inline into both sides. - SqlPlan::Intersect { left, right, all } => SqlPlan::Intersect { - left: Box::new(inline_cte(left, cte_name, cte_plan)), - right: Box::new(inline_cte(right, cte_name, cte_plan)), - all: *all, - }, - - // Except referencing CTE → inline into both sides. - SqlPlan::Except { left, right, all } => SqlPlan::Except { - left: Box::new(inline_cte(left, cte_name, cte_plan)), - right: Box::new(inline_cte(right, cte_name, cte_plan)), - all: *all, - }, - - // INSERT ... SELECT referencing CTE → inline into the source subquery. - SqlPlan::InsertSelect { - target, - source, - limit, - column_map, - } => SqlPlan::InsertSelect { - target: target.clone(), - source: Box::new(inline_cte(source, cte_name, cte_plan)), - limit: *limit, - column_map: column_map.clone(), - }, - - // A post-processor produced by an earlier CTE definition: recurse into - // its body so a later definition's references inside it still inline. - SqlPlan::Subquery { - input, - filters, - projection, - sort_keys, - offset, - distinct, - limit, - } => SqlPlan::Subquery { - input: Box::new(inline_cte(input, cte_name, cte_plan)), - filters: filters.clone(), - projection: projection.clone(), - sort_keys: sort_keys.clone(), - offset: *offset, - distinct: *distinct, - limit: *limit, - }, - - // No CTE reference — return as-is. - _ => plan.clone(), - } -} - -#[cfg(test)] -mod tests { - use super::*; - use nodedb_sql::types::{CompareOp, EngineType, Filter, FilterExpr, SortKey, SqlValue}; - - fn vector_search_body() -> SqlPlan { - SqlPlan::VectorSearch { - collection: "docs".to_string(), - field: "embedding".to_string(), - query_vector: vec![0.1, 0.2], - top_k: 3, - ef_search: 64, - metric: nodedb_sql::types::DistanceMetric::L2, - filters: Vec::new(), - array_prefilter: None, - ann_options: nodedb_sql::types::VectorAnnOptions::default(), - skip_payload_fetch: false, - payload_filters: Vec::new(), - projection: Vec::new(), - } - } - - fn scan_on_cte(filters: Vec, limit: Option) -> SqlPlan { - SqlPlan::Scan { - collection: "knn".to_string(), - alias: None, - engine: EngineType::DocumentSchemaless, - filters, - projection: Vec::new(), - sort_keys: Vec::new(), - limit, - offset: 0, - distinct: false, - window_functions: Vec::new(), - temporal: nodedb_sql::TemporalScope::default(), - } - } - - fn tag_filter() -> Filter { - Filter { - expr: FilterExpr::Comparison { - field: "tag".to_string(), - op: CompareOp::Eq, - value: SqlValue::String("keep".to_string()), - }, - } - } - - fn expect_vector_search(plan: SqlPlan) -> (Vec, usize) { - match plan { - SqlPlan::VectorSearch { filters, top_k, .. } => (filters, top_k), - other => panic!("expected VectorSearch, got {other:?}"), - } - } - - #[test] - fn outer_filter_merges_onto_vector_search_cte_body() { - let (filters, top_k) = expect_vector_search(inline_cte( - &scan_on_cte(vec![tag_filter()], None), - "knn", - &vector_search_body(), - )); - assert_eq!( - filters.len(), - 1, - "the outer WHERE must survive inlining, else the k-NN result comes back unfiltered" - ); - assert_eq!(top_k, 3, "a filter alone must not change the requested k"); - } - - #[test] - fn outer_limit_narrows_the_vector_search_top_k() { - let (_, top_k) = expect_vector_search(inline_cte( - &scan_on_cte(Vec::new(), Some(1)), - "knn", - &vector_search_body(), - )); - assert_eq!(top_k, 1, "an outer LIMIT below k must narrow the k-NN cut"); - } - - #[test] - fn outer_limit_above_k_leaves_top_k_untouched() { - let (_, top_k) = expect_vector_search(inline_cte( - &scan_on_cte(Vec::new(), Some(99)), - "knn", - &vector_search_body(), - )); - assert_eq!(top_k, 3, "an outer LIMIT above k cannot widen the k-NN cut"); - } - - #[test] - fn unconstrained_reference_returns_the_vector_search_body_verbatim() { - let (filters, top_k) = expect_vector_search(inline_cte( - &scan_on_cte(Vec::new(), None), - "knn", - &vector_search_body(), - )); - assert!(filters.is_empty()); - assert_eq!(top_k, 3); - } - - /// A CTE-referencing scan carrying an outer ORDER BY / OFFSET / DISTINCT. - fn scan_on_cte_reorder(sort_keys: Vec, offset: usize, distinct: bool) -> SqlPlan { - SqlPlan::Scan { - collection: "knn".to_string(), - alias: None, - engine: EngineType::DocumentSchemaless, - filters: Vec::new(), - projection: Vec::new(), - sort_keys, - limit: None, - offset, - distinct, - window_functions: Vec::new(), - temporal: nodedb_sql::TemporalScope::default(), - } - } - - fn id_sort_key() -> SortKey { - SortKey { - expr: nodedb_sql::types::SqlExpr::Column { - table: Some("s".to_string()), - name: "id".to_string(), - }, - ascending: true, - nulls_first: false, - } - } - - #[test] - fn outer_order_by_wraps_vector_search_in_subquery() { - // An outer ORDER BY cannot fold into the k-NN leaf; it must become a - // post-processor over the search, and the leaf keeps its own top_k. - match inline_cte( - &scan_on_cte_reorder(vec![id_sort_key()], 0, false), - "knn", - &vector_search_body(), - ) { - SqlPlan::Subquery { - input, sort_keys, .. - } => { - assert_eq!( - sort_keys.len(), - 1, - "the outer ORDER BY must ride the wrapper" - ); - assert!( - matches!(*input, SqlPlan::VectorSearch { top_k: 3, .. }), - "the search leaf keeps its own top_k under the wrapper" - ); - } - other => panic!("expected Subquery, got {other:?}"), - } - } - - #[test] - fn outer_distinct_and_offset_wrap_vector_search_in_subquery() { - match inline_cte( - &scan_on_cte_reorder(Vec::new(), 2, true), - "knn", - &vector_search_body(), - ) { - SqlPlan::Subquery { - offset, distinct, .. - } => { - assert_eq!(offset, 2, "the outer OFFSET must ride the wrapper"); - assert!(distinct, "the outer DISTINCT must ride the wrapper"); - } - other => panic!("expected Subquery, got {other:?}"), - } - } - - #[test] - fn plain_limit_does_not_wrap_vector_search() { - // A LIMIT with no reorder still folds into top_k (fast path), NOT a - // Subquery wrapper. - let plan = inline_cte( - &scan_on_cte(Vec::new(), Some(1)), - "knn", - &vector_search_body(), - ); - assert!( - matches!(plan, SqlPlan::VectorSearch { top_k: 1, .. }), - "an unordered LIMIT must fold into top_k, not wrap: {plan:?}" - ); - } -} diff --git a/nodedb/src/control/planner/sql_plan_convert/expr/bridge_expr.rs b/nodedb/src/control/planner/sql_plan_convert/expr/bridge_expr.rs new file mode 100644 index 000000000..4924eba8b --- /dev/null +++ b/nodedb/src/control/planner/sql_plan_convert/expr/bridge_expr.rs @@ -0,0 +1,265 @@ +// SPDX-License-Identifier: BUSL-1.1 + +use nodedb_sql::types::SqlExpr; + +use super::super::value::sql_value_to_nodedb_value; + +/// Convert a `nodedb_sql::types::SqlExpr` (parser AST) to a +/// `nodedb_query::expr::SqlExpr` (bridge evaluation type). +/// +/// Column references use the **bare** name (no table qualifier) for +/// single-collection evaluation contexts (WHERE, CHECK, GENERATED). +/// For join contexts where the merged document uses qualified keys +/// (`"t1.col"`), use [`sql_expr_to_bridge_expr_qualified`] instead. +pub(in crate::control::planner::sql_plan_convert) fn sql_expr_to_bridge_expr( + expr: &SqlExpr, +) -> crate::bridge::expr_eval::SqlExpr { + convert_expr_inner(expr, false) +} + +/// Like [`sql_expr_to_bridge_expr`] but qualifies column references +/// with their table name (`t.col` → `"t.col"`) for join merged docs. +pub(in crate::control::planner::sql_plan_convert) fn sql_expr_to_bridge_expr_qualified( + expr: &SqlExpr, +) -> crate::bridge::expr_eval::SqlExpr { + convert_expr_inner(expr, true) +} + +fn convert_expr_inner(expr: &SqlExpr, qualify: bool) -> crate::bridge::expr_eval::SqlExpr { + use crate::bridge::expr_eval::SqlExpr as BExpr; + match expr { + SqlExpr::Column { table, name } => { + // `EXCLUDED.col` references the row proposed for insertion in + // `INSERT ... ON CONFLICT DO UPDATE`. Emit the dedicated + // variant so the upsert handler can resolve against the + // incoming row via `eval_with_excluded`. The table qualifier + // comes in already-normalized (lowercased) from the parser. + if table + .as_deref() + .is_some_and(|t| t.eq_ignore_ascii_case("excluded")) + { + return BExpr::ExcludedColumn(name.clone()); + } + if qualify { + BExpr::Column(nodedb_sql::planner::qualified_name(table.as_deref(), name)) + } else { + BExpr::Column(name.clone()) + } + } + SqlExpr::Literal(v) => BExpr::Literal(sql_value_to_nodedb_value(v)), + SqlExpr::BinaryOp { left, op, right } => BExpr::BinaryOp { + left: Box::new(convert_expr_inner(left, qualify)), + op: match op { + nodedb_sql::types::BinaryOp::Add => crate::bridge::expr_eval::BinaryOp::Add, + nodedb_sql::types::BinaryOp::Sub => crate::bridge::expr_eval::BinaryOp::Sub, + nodedb_sql::types::BinaryOp::Mul => crate::bridge::expr_eval::BinaryOp::Mul, + nodedb_sql::types::BinaryOp::Div => crate::bridge::expr_eval::BinaryOp::Div, + nodedb_sql::types::BinaryOp::Mod => crate::bridge::expr_eval::BinaryOp::Mod, + nodedb_sql::types::BinaryOp::Eq => crate::bridge::expr_eval::BinaryOp::Eq, + nodedb_sql::types::BinaryOp::Ne => crate::bridge::expr_eval::BinaryOp::NotEq, + nodedb_sql::types::BinaryOp::Gt => crate::bridge::expr_eval::BinaryOp::Gt, + nodedb_sql::types::BinaryOp::Ge => crate::bridge::expr_eval::BinaryOp::GtEq, + nodedb_sql::types::BinaryOp::Lt => crate::bridge::expr_eval::BinaryOp::Lt, + nodedb_sql::types::BinaryOp::Le => crate::bridge::expr_eval::BinaryOp::LtEq, + nodedb_sql::types::BinaryOp::And => crate::bridge::expr_eval::BinaryOp::And, + nodedb_sql::types::BinaryOp::Or => crate::bridge::expr_eval::BinaryOp::Or, + nodedb_sql::types::BinaryOp::Concat => crate::bridge::expr_eval::BinaryOp::Concat, + }, + right: Box::new(convert_expr_inner(right, qualify)), + }, + SqlExpr::Function { name, args, .. } => BExpr::Function { + name: name.clone(), + args: args + .iter() + .map(|a| convert_expr_inner(a, qualify)) + .collect(), + }, + SqlExpr::Case { + operand, + when_then, + else_expr, + } => BExpr::Case { + operand: operand + .as_ref() + .map(|e| Box::new(convert_expr_inner(e, qualify))), + when_thens: when_then + .iter() + .map(|(w, t)| { + ( + convert_expr_inner(w, qualify), + convert_expr_inner(t, qualify), + ) + }) + .collect(), + else_expr: else_expr + .as_ref() + .map(|e| Box::new(convert_expr_inner(e, qualify))), + }, + SqlExpr::Cast { expr, to_type } => { + let cast_type = match to_type.to_uppercase().as_str() { + "INT" | "INTEGER" | "BIGINT" | "SMALLINT" => { + crate::bridge::expr_eval::CastType::Int + } + "FLOAT" | "DOUBLE" | "REAL" | "NUMERIC" | "DECIMAL" => { + crate::bridge::expr_eval::CastType::Float + } + "BOOL" | "BOOLEAN" => crate::bridge::expr_eval::CastType::Bool, + _ => crate::bridge::expr_eval::CastType::String, + }; + BExpr::Cast { + expr: Box::new(convert_expr_inner(expr, qualify)), + to_type: cast_type, + } + } + SqlExpr::Wildcard => BExpr::Column("*".into()), + + // NOT e / -e → evaluator's Negate (handles both bool and numeric). + SqlExpr::UnaryOp { expr, .. } => BExpr::Negate(Box::new(convert_expr_inner(expr, qualify))), + + // `e IS NULL` / `e IS NOT NULL` — direct passthrough. + SqlExpr::IsNull { expr, negated } => BExpr::IsNull { + expr: Box::new(convert_expr_inner(expr, qualify)), + negated: *negated, + }, + + // `e BETWEEN low AND high` desugars to `e >= low AND e <= high` + // (or `e < low OR e > high` when negated). The evaluator has no + // native Between variant, so the planner must lower it here. + SqlExpr::Between { + expr, + low, + high, + negated, + } => { + let e = convert_expr_inner(expr, qualify); + let l = convert_expr_inner(low, qualify); + let h = convert_expr_inner(high, qualify); + if *negated { + let lt = BExpr::BinaryOp { + left: Box::new(e.clone()), + op: crate::bridge::expr_eval::BinaryOp::Lt, + right: Box::new(l), + }; + let gt = BExpr::BinaryOp { + left: Box::new(e), + op: crate::bridge::expr_eval::BinaryOp::Gt, + right: Box::new(h), + }; + BExpr::BinaryOp { + left: Box::new(lt), + op: crate::bridge::expr_eval::BinaryOp::Or, + right: Box::new(gt), + } + } else { + let ge = BExpr::BinaryOp { + left: Box::new(e.clone()), + op: crate::bridge::expr_eval::BinaryOp::GtEq, + right: Box::new(l), + }; + let le = BExpr::BinaryOp { + left: Box::new(e), + op: crate::bridge::expr_eval::BinaryOp::LtEq, + right: Box::new(h), + }; + BExpr::BinaryOp { + left: Box::new(ge), + op: crate::bridge::expr_eval::BinaryOp::And, + right: Box::new(le), + } + } + } + + // `e IN (a, b, c)` desugars to `e = a OR e = b OR e = c` — each + // element may itself be a non-literal expression, so we must + // recursively convert and OR the comparisons together. `NOT IN` + // is `e <> a AND e <> b AND e <> c`. + SqlExpr::InList { + expr, + list, + negated, + } => { + let target = convert_expr_inner(expr, qualify); + if list.is_empty() { + // Empty list: `e IN ()` = false, `e NOT IN ()` = true. + return BExpr::Literal(nodedb_types::Value::Bool(*negated)); + } + let (eq_op, combine_op) = if *negated { + ( + crate::bridge::expr_eval::BinaryOp::NotEq, + crate::bridge::expr_eval::BinaryOp::And, + ) + } else { + ( + crate::bridge::expr_eval::BinaryOp::Eq, + crate::bridge::expr_eval::BinaryOp::Or, + ) + }; + // Empty list is handled above, so `list` is guaranteed non-empty + // here: we reduce `(target eq list[0]) op (target eq list[1]) op ...` + // without touching `.unwrap()` or `.expect()`. + list.iter() + .map(|item| BExpr::BinaryOp { + left: Box::new(target.clone()), + op: eq_op, + right: Box::new(convert_expr_inner(item, qualify)), + }) + .reduce(|acc, next| BExpr::BinaryOp { + left: Box::new(acc), + op: combine_op, + right: Box::new(next), + }) + // Unreachable: `list.is_empty()` returns early above. + .unwrap_or(BExpr::Literal(nodedb_types::Value::Bool(*negated))) + } + + // `e LIKE pattern` — no direct evaluator variant; route through a + // function call so the shared function dispatcher handles it. + SqlExpr::Like { + expr, + pattern, + negated, + case_insensitive, + } => { + let fn_name = if *case_insensitive { "ilike" } else { "like" }; + let call = BExpr::Function { + name: fn_name.into(), + args: vec![ + convert_expr_inner(expr, qualify), + convert_expr_inner(pattern, qualify), + ], + }; + if *negated { + BExpr::Negate(Box::new(call)) + } else { + call + } + } + + // `ARRAY['a', 'b', ...]` — lower each element and, when all resolve to + // `BExpr::Literal`, fold into a single `Value::Array` literal so that + // functions like `pg_json_has_any_key` / `pg_json_has_all_keys` receive + // a proper `Value::Array` argument rather than `Value::Null`. + SqlExpr::ArrayLiteral(elems) => { + let mut values = Vec::with_capacity(elems.len()); + let mut all_literal = true; + for elem in elems { + match convert_expr_inner(elem, qualify) { + BExpr::Literal(v) => values.push(v), + other => { + all_literal = false; + // Non-literal element: fall back to Null for that slot. + let _ = other; + values.push(nodedb_types::Value::Null); + } + } + } + if all_literal { + BExpr::Literal(nodedb_types::Value::Array(values)) + } else { + BExpr::Literal(nodedb_types::Value::Null) + } + } + + _ => BExpr::Literal(nodedb_types::Value::Null), + } +} diff --git a/nodedb/src/control/planner/sql_plan_convert/expr/inline_cte.rs b/nodedb/src/control/planner/sql_plan_convert/expr/inline_cte.rs new file mode 100644 index 000000000..bcd334df5 --- /dev/null +++ b/nodedb/src/control/planner/sql_plan_convert/expr/inline_cte.rs @@ -0,0 +1,758 @@ +// SPDX-License-Identifier: BUSL-1.1 + +use nodedb_sql::types::{Filter, Projection, SortKey, SqlPlan, WindowSpec}; + +/// Replace scans on `cte_name` with the CTE's actual subquery plan. +/// +/// Outer constraints on the CTE reference merge onto the body only where the +/// body can carry them without changing the rows it produces: +/// +/// - A plain filtered `Scan` body (no computed projection, window functions, +/// LIMIT, OFFSET, DISTINCT, or ORDER BY) takes every outer constraint. +/// - A `VectorSearch` body takes filters, a column-only projection, and an +/// unordered LIMIT (as `top_k`). +/// - A column-only outer projection over any other body resolves by output +/// schema, so the body is returned as is. +/// +/// Every other combination lowers into a `Subquery` post-processor over the +/// body's materialized rows. Computed projection entries and window functions +/// always take that path: the body does not evaluate them, and a name lookup +/// at the response boundary yields NULL. +pub(in crate::control::planner::sql_plan_convert) fn inline_cte( + plan: &SqlPlan, + cte_name: &str, + cte_plan: &SqlPlan, +) -> SqlPlan { + match plan { + // Direct scan on CTE name → replace with CTE plan. + SqlPlan::Scan { + collection, + filters, + projection, + sort_keys, + limit, + offset, + distinct, + window_functions, + .. + } if collection == cte_name => inline_cte_scan_ref( + ScanRef { + filters, + projection, + sort_keys, + limit: *limit, + offset: *offset, + distinct: *distinct, + window_functions, + has_computed: has_computed_projection(projection), + }, + cte_plan, + ), + + // Aggregate referencing CTE → inline into the input. + SqlPlan::Aggregate { + input, + group_by, + group_by_aliases, + output_order, + aggregates, + having, + limit, + grouping_sets, + sort_keys, + } => SqlPlan::Aggregate { + input: Box::new(inline_cte(input, cte_name, cte_plan)), + group_by: group_by.clone(), + group_by_aliases: group_by_aliases.clone(), + output_order: output_order.clone(), + aggregates: aggregates.clone(), + having: having.clone(), + limit: *limit, + grouping_sets: grouping_sets.clone(), + sort_keys: sort_keys.clone(), + }, + + // JOIN referencing CTE on either side. + SqlPlan::Join { + left, + right, + on, + join_type, + condition, + limit, + projection, + filters, + } => SqlPlan::Join { + left: Box::new(inline_cte(left, cte_name, cte_plan)), + right: Box::new(inline_cte(right, cte_name, cte_plan)), + on: on.clone(), + join_type: *join_type, + condition: condition.clone(), + limit: *limit, + projection: projection.clone(), + filters: filters.clone(), + }, + + // Union referencing CTE → inline into all inputs. + SqlPlan::Union { inputs, distinct } => SqlPlan::Union { + inputs: inputs + .iter() + .map(|i| inline_cte(i, cte_name, cte_plan)) + .collect(), + distinct: *distinct, + }, + + // Intersect referencing CTE → inline into both sides. + SqlPlan::Intersect { left, right, all } => SqlPlan::Intersect { + left: Box::new(inline_cte(left, cte_name, cte_plan)), + right: Box::new(inline_cte(right, cte_name, cte_plan)), + all: *all, + }, + + // Except referencing CTE → inline into both sides. + SqlPlan::Except { left, right, all } => SqlPlan::Except { + left: Box::new(inline_cte(left, cte_name, cte_plan)), + right: Box::new(inline_cte(right, cte_name, cte_plan)), + all: *all, + }, + + // INSERT ... SELECT referencing CTE → inline into the source subquery. + SqlPlan::InsertSelect { + target, + source, + limit, + column_map, + } => SqlPlan::InsertSelect { + target: target.clone(), + source: Box::new(inline_cte(source, cte_name, cte_plan)), + limit: *limit, + column_map: column_map.clone(), + }, + + // A post-processor produced by an earlier CTE definition: recurse into + // its body so a later definition's references inside it still inline. + SqlPlan::Subquery { + input, + filters, + projection, + window_functions, + sort_keys, + offset, + distinct, + limit, + } => SqlPlan::Subquery { + input: Box::new(inline_cte(input, cte_name, cte_plan)), + filters: filters.clone(), + projection: projection.clone(), + window_functions: window_functions.clone(), + sort_keys: sort_keys.clone(), + offset: *offset, + distinct: *distinct, + limit: *limit, + }, + + // No CTE reference — return as-is. + _ => plan.clone(), + } +} + +/// `true` if any projection entry is a computed expression (`price * qty AS +/// total`) rather than a bare column or star. +fn has_computed_projection(projection: &[Projection]) -> bool { + projection + .iter() + .any(|p| matches!(p, Projection::Computed { .. })) +} + +/// The outer constraints carried on a `Scan` that references the CTE by +/// name. `has_computed` is precomputed once so the sub-cases below don't +/// each re-walk `projection`. +#[derive(Clone, Copy)] +struct ScanRef<'a> { + filters: &'a Vec, + projection: &'a Vec, + sort_keys: &'a Vec, + limit: Option, + offset: usize, + distinct: bool, + window_functions: &'a Vec, + has_computed: bool, +} + +impl ScanRef<'_> { + /// No filter, sort, limit, offset, distinct, window, or computed entry — + /// the outer reference adds nothing the body's output schema doesn't + /// already answer. + fn is_unconstrained(&self) -> bool { + self.filters.is_empty() + && self.sort_keys.is_empty() + && self.limit.is_none() + && !self.distinct + && self.offset == 0 + && self.window_functions.is_empty() + && !self.has_computed + } +} + +/// Wrap `input` in a `Subquery` post-processor carrying the outer +/// constraints that `input` has no slot for. +fn wrap_in_subquery(input: SqlPlan, filters: Vec, outer: ScanRef<'_>) -> SqlPlan { + SqlPlan::Subquery { + input: Box::new(input), + filters, + projection: outer.projection.clone(), + window_functions: outer.window_functions.clone(), + sort_keys: outer.sort_keys.clone(), + offset: outer.offset, + distinct: outer.distinct, + limit: outer.limit, + } +} + +/// Resolve a CTE reference for a `Scan { collection: cte_name, .. }` node, +/// merging the outer constraints onto `cte_plan` as far as its body kind +/// can carry them. +fn inline_cte_scan_ref(outer: ScanRef<'_>, cte_plan: &SqlPlan) -> SqlPlan { + // A column-only projection resolves by the body's output schema. + if outer.is_unconstrained() { + return cte_plan.clone(); + } + + if let Some(merged) = merge_into_scan_body(outer, cte_plan) { + return merged; + } + + if let Some(merged) = merge_into_vector_search_body(outer, cte_plan) { + return merged; + } + + // Any other body (Aggregate, Join, TextSearch, HybridSearch, + // SparseSearch, SpatialScan, MultiVectorSearch, a constrained Scan, + // ...) has no slot for the outer constraints reaching this point — the + // unconstrained case already returned at the top of the function. Apply + // them over the body's materialized rows in a `Subquery` + // post-processor, which evaluates computed columns and window + // functions itself. + wrap_in_subquery(cte_plan.clone(), outer.filters.clone(), outer) +} + +/// A plain filtered `Scan` body takes every outer constraint. A body that +/// limits, offsets, dedups, orders, computes, or windows changes which rows +/// the outer constraints see if they merge into it: an outer WHERE inside an +/// inner LIMIT changes the cut, and an inner `qty AS x` alias replaced by the +/// outer projection makes `x` NULL. +fn merge_into_scan_body(outer: ScanRef<'_>, cte_plan: &SqlPlan) -> Option { + let SqlPlan::Scan { + collection: inner_col, + alias: inner_alias, + engine: inner_eng, + filters: inner_f, + projection: inner_p, + sort_keys: inner_s, + limit: inner_l, + offset: inner_o, + distinct: inner_d, + window_functions: inner_w, + temporal: inner_t, + } = cte_plan + else { + return None; + }; + if has_computed_projection(inner_p) + || !inner_w.is_empty() + || inner_l.is_some() + || *inner_o != 0 + || *inner_d + || !inner_s.is_empty() + { + return None; + } + + let mut merged_filters = inner_f.clone(); + merged_filters.extend(outer.filters.iter().cloned()); + Some(SqlPlan::Scan { + collection: inner_col.clone(), + alias: inner_alias.clone(), + engine: *inner_eng, + filters: merged_filters, + // A named outer projection overrides the inner one. An empty or + // star-only outer projection inherits the CTE's own column list, so + // `SELECT * FROM (SELECT a FROM t)` emits `a` alone. + projection: if outer + .projection + .iter() + .all(|p| matches!(p, Projection::Star | Projection::QualifiedStar(_))) + { + inner_p.clone() + } else { + outer.projection.clone() + }, + sort_keys: outer.sort_keys.clone(), + limit: outer.limit, + offset: outer.offset, + distinct: outer.distinct, + window_functions: outer.window_functions.clone(), + temporal: *inner_t, + }) +} + +/// A k-NN body carries its own post-filter list and top-k. An outer `WHERE` +/// merges into the engine post-filter so the cut counts MATCHING rows. When +/// nothing reorders the result and the outer projection is column-only, an +/// unordered `LIMIT` folds into `top_k` and the projection rides along. An +/// outer `ORDER BY` / `OFFSET` / `DISTINCT` reorders the k rows, and a +/// computed projection or window function evaluates over them; the search +/// leaf has no slot for any of those, so they (and a `LIMIT` that must apply +/// after the reorder) run in a `Subquery` post-processor over the k rows. +fn merge_into_vector_search_body(outer: ScanRef<'_>, cte_plan: &SqlPlan) -> Option { + if !matches!(cte_plan, SqlPlan::VectorSearch { .. }) { + return None; + } + + let needs_reorder = !outer.sort_keys.is_empty() + || outer.offset > 0 + || outer.distinct + || outer.has_computed + || !outer.window_functions.is_empty(); + let mut leaf = cte_plan.clone(); + if let SqlPlan::VectorSearch { + filters: body_filters, + projection: body_projection, + top_k, + .. + } = &mut leaf + { + body_filters.extend(outer.filters.iter().cloned()); + if !needs_reorder { + if !outer.projection.is_empty() { + body_projection.clone_from(outer.projection); + } + if let Some(outer_limit) = outer.limit { + *top_k = (*top_k).min(outer_limit); + } + } + } + if !needs_reorder { + return Some(leaf); + } + // Filters already run in the engine; the tail applies the + // reorder-dependent constraints over the k rows. It sorts before + // projecting, so ORDER BY may reference any column. + Some(wrap_in_subquery(leaf, Vec::new(), outer)) +} + +#[cfg(test)] +mod tests { + use super::*; + use nodedb_sql::types::{CompareOp, EngineType, Filter, FilterExpr, SortKey, SqlValue}; + + fn vector_search_body() -> SqlPlan { + SqlPlan::VectorSearch { + collection: "docs".to_string(), + field: "embedding".to_string(), + query_vector: vec![0.1, 0.2], + top_k: 3, + ef_search: 64, + metric: nodedb_sql::types::DistanceMetric::L2, + filters: Vec::new(), + array_prefilter: None, + ann_options: nodedb_sql::types::VectorAnnOptions::default(), + skip_payload_fetch: false, + payload_filters: Vec::new(), + projection: Vec::new(), + } + } + + fn scan_on_cte(filters: Vec, limit: Option) -> SqlPlan { + SqlPlan::Scan { + collection: "knn".to_string(), + alias: None, + engine: EngineType::DocumentSchemaless, + filters, + projection: Vec::new(), + sort_keys: Vec::new(), + limit, + offset: 0, + distinct: false, + window_functions: Vec::new(), + temporal: nodedb_sql::TemporalScope::default(), + } + } + + fn tag_filter() -> Filter { + Filter { + expr: FilterExpr::Comparison { + field: "tag".to_string(), + op: CompareOp::Eq, + value: SqlValue::String("keep".to_string()), + }, + } + } + + fn expect_vector_search(plan: SqlPlan) -> (Vec, usize) { + match plan { + SqlPlan::VectorSearch { filters, top_k, .. } => (filters, top_k), + other => panic!("expected VectorSearch, got {other:?}"), + } + } + + #[test] + fn outer_filter_merges_onto_vector_search_cte_body() { + let (filters, top_k) = expect_vector_search(inline_cte( + &scan_on_cte(vec![tag_filter()], None), + "knn", + &vector_search_body(), + )); + assert_eq!( + filters.len(), + 1, + "the outer WHERE must survive inlining, else the k-NN result comes back unfiltered" + ); + assert_eq!(top_k, 3, "a filter alone must not change the requested k"); + } + + #[test] + fn outer_limit_narrows_the_vector_search_top_k() { + let (_, top_k) = expect_vector_search(inline_cte( + &scan_on_cte(Vec::new(), Some(1)), + "knn", + &vector_search_body(), + )); + assert_eq!(top_k, 1, "an outer LIMIT below k must narrow the k-NN cut"); + } + + #[test] + fn outer_limit_above_k_leaves_top_k_untouched() { + let (_, top_k) = expect_vector_search(inline_cte( + &scan_on_cte(Vec::new(), Some(99)), + "knn", + &vector_search_body(), + )); + assert_eq!(top_k, 3, "an outer LIMIT above k cannot widen the k-NN cut"); + } + + #[test] + fn unconstrained_reference_returns_the_vector_search_body_verbatim() { + let (filters, top_k) = expect_vector_search(inline_cte( + &scan_on_cte(Vec::new(), None), + "knn", + &vector_search_body(), + )); + assert!(filters.is_empty()); + assert_eq!(top_k, 3); + } + + /// A CTE-referencing scan carrying an outer ORDER BY / OFFSET / DISTINCT. + fn scan_on_cte_reorder(sort_keys: Vec, offset: usize, distinct: bool) -> SqlPlan { + SqlPlan::Scan { + collection: "knn".to_string(), + alias: None, + engine: EngineType::DocumentSchemaless, + filters: Vec::new(), + projection: Vec::new(), + sort_keys, + limit: None, + offset, + distinct, + window_functions: Vec::new(), + temporal: nodedb_sql::TemporalScope::default(), + } + } + + fn id_sort_key() -> SortKey { + SortKey { + expr: nodedb_sql::types::SqlExpr::Column { + table: Some("s".to_string()), + name: "id".to_string(), + }, + ascending: true, + nulls_first: false, + } + } + + #[test] + fn outer_order_by_wraps_vector_search_in_subquery() { + // An outer ORDER BY cannot fold into the k-NN leaf; it must become a + // post-processor over the search, and the leaf keeps its own top_k. + match inline_cte( + &scan_on_cte_reorder(vec![id_sort_key()], 0, false), + "knn", + &vector_search_body(), + ) { + SqlPlan::Subquery { + input, sort_keys, .. + } => { + assert_eq!( + sort_keys.len(), + 1, + "the outer ORDER BY must ride the wrapper" + ); + assert!( + matches!(*input, SqlPlan::VectorSearch { top_k: 3, .. }), + "the search leaf keeps its own top_k under the wrapper" + ); + } + other => panic!("expected Subquery, got {other:?}"), + } + } + + #[test] + fn outer_distinct_and_offset_wrap_vector_search_in_subquery() { + match inline_cte( + &scan_on_cte_reorder(Vec::new(), 2, true), + "knn", + &vector_search_body(), + ) { + SqlPlan::Subquery { + offset, distinct, .. + } => { + assert_eq!(offset, 2, "the outer OFFSET must ride the wrapper"); + assert!(distinct, "the outer DISTINCT must ride the wrapper"); + } + other => panic!("expected Subquery, got {other:?}"), + } + } + + #[test] + fn plain_limit_does_not_wrap_vector_search() { + // A LIMIT with no reorder still folds into top_k (fast path), NOT a + // Subquery wrapper. + let plan = inline_cte( + &scan_on_cte(Vec::new(), Some(1)), + "knn", + &vector_search_body(), + ); + assert!( + matches!(plan, SqlPlan::VectorSearch { top_k: 1, .. }), + "an unordered LIMIT must fold into top_k, not wrap: {plan:?}" + ); + } + + fn doubled_x() -> Projection { + Projection::Computed { + expr: nodedb_sql::types::SqlExpr::BinaryOp { + left: Box::new(nodedb_sql::types::SqlExpr::Column { + table: None, + name: "x".to_string(), + }), + op: nodedb_sql::types::BinaryOp::Mul, + right: Box::new(nodedb_sql::types::SqlExpr::Literal(SqlValue::Int(2))), + }, + alias: "y".to_string(), + } + } + + fn row_number_spec() -> nodedb_sql::types::WindowSpec { + nodedb_sql::types::WindowSpec { + function: "row_number".to_string(), + args: Vec::new(), + partition_by: Vec::new(), + order_by: Vec::new(), + alias: "rn".to_string(), + frame: Default::default(), + } + } + + /// A CTE-referencing scan carrying an outer projection and window list. + fn scan_on_cte_projected( + projection: Vec, + window_functions: Vec, + ) -> SqlPlan { + SqlPlan::Scan { + collection: "knn".to_string(), + alias: None, + engine: EngineType::DocumentSchemaless, + filters: Vec::new(), + projection, + sort_keys: Vec::new(), + limit: None, + offset: 0, + distinct: false, + window_functions, + temporal: nodedb_sql::TemporalScope::default(), + } + } + + /// A `SELECT 1 AS x` body: a constant result, not a scan. + fn constant_body() -> SqlPlan { + SqlPlan::ConstantResult { + columns: vec!["x".to_string()], + values: vec![SqlValue::Int(1)], + volatile: false, + } + } + + #[test] + fn computed_projection_over_vector_search_wraps_in_subquery() { + match inline_cte( + &scan_on_cte_projected(vec![doubled_x()], Vec::new()), + "knn", + &vector_search_body(), + ) { + SqlPlan::Subquery { + input, projection, .. + } => { + assert_eq!(projection.len(), 1, "the computed entry rides the wrapper"); + assert!( + matches!(&*input, SqlPlan::VectorSearch { projection, .. } if projection.is_empty()), + "a computed projection never folds into the search leaf" + ); + } + other => panic!("expected Subquery, got {other:?}"), + } + } + + #[test] + fn column_only_projection_over_non_scan_body_returns_body() { + let plan = inline_cte( + &scan_on_cte_projected(vec![Projection::Column("x".to_string())], Vec::new()), + "knn", + &constant_body(), + ); + assert!( + matches!(plan, SqlPlan::ConstantResult { .. }), + "a column-only projection resolves by output schema: {plan:?}" + ); + } + + #[test] + fn computed_projection_over_constant_body_wraps_in_subquery() { + match inline_cte( + &scan_on_cte_projected(vec![doubled_x()], Vec::new()), + "knn", + &constant_body(), + ) { + SqlPlan::Subquery { + input, + projection, + window_functions, + .. + } => { + assert!(matches!(*input, SqlPlan::ConstantResult { .. })); + assert_eq!(projection.len(), 1); + assert!(window_functions.is_empty()); + } + other => panic!("expected Subquery, got {other:?}"), + } + } + + #[test] + fn window_function_over_constant_body_rides_the_subquery() { + match inline_cte( + &scan_on_cte_projected(Vec::new(), vec![row_number_spec()]), + "knn", + &constant_body(), + ) { + SqlPlan::Subquery { + window_functions, .. + } => assert_eq!(window_functions[0].alias, "rn"), + other => panic!("expected Subquery, got {other:?}"), + } + } + + fn limited_scan_body() -> SqlPlan { + SqlPlan::Scan { + collection: "orders".to_string(), + alias: None, + engine: EngineType::DocumentSchemaless, + filters: Vec::new(), + projection: Vec::new(), + sort_keys: Vec::new(), + limit: Some(5), + offset: 0, + distinct: false, + window_functions: Vec::new(), + temporal: nodedb_sql::TemporalScope::default(), + } + } + + #[test] + fn outer_filter_over_limited_scan_body_wraps_instead_of_merging() { + // Merging the WHERE under the inner LIMIT changes which rows the + // limit sees; the filter must run over the limited rows instead. + match inline_cte( + &scan_on_cte(vec![tag_filter()], None), + "knn", + &limited_scan_body(), + ) { + SqlPlan::Subquery { input, filters, .. } => { + assert_eq!(filters.len(), 1); + assert!(matches!(*input, SqlPlan::Scan { limit: Some(5), .. })); + } + other => panic!("expected Subquery, got {other:?}"), + } + } + + #[test] + fn outer_filter_over_plain_scan_body_merges() { + let body = SqlPlan::Scan { + collection: "orders".to_string(), + alias: None, + engine: EngineType::DocumentSchemaless, + filters: Vec::new(), + projection: vec![Projection::Column("qty".to_string())], + sort_keys: Vec::new(), + limit: None, + offset: 0, + distinct: false, + window_functions: Vec::new(), + temporal: nodedb_sql::TemporalScope::default(), + }; + match inline_cte(&scan_on_cte(vec![tag_filter()], Some(2)), "knn", &body) { + SqlPlan::Scan { + collection, + filters, + projection, + limit, + .. + } => { + assert_eq!(collection, "orders"); + assert_eq!(filters.len(), 1); + assert_eq!(limit, Some(2)); + assert_eq!(projection.len(), 1, "the inner column list is inherited"); + } + other => panic!("expected merged Scan, got {other:?}"), + } + } + + #[test] + fn aliased_inner_scan_body_keeps_alias_under_outer_computed() { + // `SELECT x * 2 AS y FROM (SELECT qty AS x FROM orders) s`: the inner + // alias must be evaluated by the body before the outer expression reads it. + let body = SqlPlan::Scan { + collection: "orders".to_string(), + alias: None, + engine: EngineType::DocumentSchemaless, + filters: Vec::new(), + projection: vec![Projection::Computed { + expr: nodedb_sql::types::SqlExpr::Column { + table: None, + name: "qty".to_string(), + }, + alias: "x".to_string(), + }], + sort_keys: Vec::new(), + limit: None, + offset: 0, + distinct: false, + window_functions: Vec::new(), + temporal: nodedb_sql::TemporalScope::default(), + }; + match inline_cte( + &scan_on_cte_projected(vec![doubled_x()], Vec::new()), + "knn", + &body, + ) { + SqlPlan::Subquery { + input, projection, .. + } => { + assert!(matches!( + &*input, + SqlPlan::Scan { projection, .. } if projection.len() == 1 + )); + assert_eq!(projection.len(), 1); + } + other => panic!("expected Subquery, got {other:?}"), + } + } +} diff --git a/nodedb/src/control/planner/sql_plan_convert/expr/mod.rs b/nodedb/src/control/planner/sql_plan_convert/expr/mod.rs new file mode 100644 index 000000000..a2e40b80f --- /dev/null +++ b/nodedb/src/control/planner/sql_plan_convert/expr/mod.rs @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Expression conversion and CTE inlining. + +mod bridge_expr; +mod inline_cte; +mod sort_keys; + +pub(super) use bridge_expr::{sql_expr_to_bridge_expr, sql_expr_to_bridge_expr_qualified}; +pub(super) use inline_cte::inline_cte; +pub(super) use sort_keys::convert_sort_keys; diff --git a/nodedb/src/control/planner/sql_plan_convert/expr/sort_keys.rs b/nodedb/src/control/planner/sql_plan_convert/expr/sort_keys.rs new file mode 100644 index 000000000..2f80dbf6a --- /dev/null +++ b/nodedb/src/control/planner/sql_plan_convert/expr/sort_keys.rs @@ -0,0 +1,23 @@ +// SPDX-License-Identifier: BUSL-1.1 + +use nodedb_physical::physical_plan::SortKeySpec; +use nodedb_sql::types::SortKey; + +use super::bridge_expr::sql_expr_to_bridge_expr; + +/// Lower planner sort keys to their physical form. +/// +/// Every key is carried, expression and all. Dropping a key the Data Plane +/// cannot name as a stored column would silently answer +/// `ORDER BY 100 / weight` with rows in storage order. +pub(in crate::control::planner::sql_plan_convert) fn convert_sort_keys( + keys: &[SortKey], +) -> Vec { + keys.iter() + .map(|k| SortKeySpec { + expr: sql_expr_to_bridge_expr(&k.expr), + ascending: k.ascending, + nulls_first: k.nulls_first, + }) + .collect() +} diff --git a/nodedb/src/control/planner/sql_plan_convert/mod.rs b/nodedb/src/control/planner/sql_plan_convert/mod.rs index ba341a353..c08ece90b 100644 --- a/nodedb/src/control/planner/sql_plan_convert/mod.rs +++ b/nodedb/src/control/planner/sql_plan_convert/mod.rs @@ -4,6 +4,7 @@ pub mod aggregate; pub mod array_alter_convert; pub mod array_convert; pub mod array_fn_convert; +pub mod body; pub mod cache_verdict; pub mod convert; pub mod dml; diff --git a/nodedb/src/control/planner/sql_plan_convert/output_schema/build.rs b/nodedb/src/control/planner/sql_plan_convert/output_schema/build.rs index 06733ce39..1d9a5dd14 100644 --- a/nodedb/src/control/planner/sql_plan_convert/output_schema/build.rs +++ b/nodedb/src/control/planner/sql_plan_convert/output_schema/build.rs @@ -173,21 +173,29 @@ pub fn build_output_schema( let types = super::join_types::join_column_types(left, right, catalog, database_id); schema_from_projection(projection, &types, &[]) } - SqlPlan::ConstantResult { columns, .. } => { + SqlPlan::ConstantResult { + columns, values, .. + } => { // The row payload keys each cell by the unique per-column key // (`cell_keys`), not the raw display name: two constant columns may // share a name (`SELECT nextval('s'), nextval('s')`), and a single - // JSON object would collapse them. `display_name` keeps the + // object would collapse them. `display_name` keeps the // client-facing name; `lookup_key` is the cell key. + // + // The type mirrors the cell `convert_constant_result` encodes: + // `Int`/`Float`/`Bool` keep their typed cell, every other variant + // (`String`/`Null`/`Decimal`/`Bytes`/`Array`/`Timestamp`/ + // `Timestamptz`) is encoded as text. let lookup_keys = crate::control::server::response_shape::project::cell_keys(columns); OutputSchema { columns: columns .iter() .zip(lookup_keys) - .map(|(c, lookup_key)| OutputColumn { + .enumerate() + .map(|(index, (c, lookup_key))| OutputColumn { display_name: c.clone(), lookup_key, - ty: DdlColType::Text, + ty: constant_cell_type(values.get(index)), }) .collect(), is_star: false, @@ -379,6 +387,27 @@ pub fn build_output_schema( } } +/// Wire type of one constant cell. A column with no value (a plan built +/// without values) is `Text`. +fn constant_cell_type(value: Option<&nodedb_sql::types_expr::SqlValue>) -> DdlColType { + use nodedb_sql::types_expr::SqlValue; + match value { + Some(SqlValue::Int(_)) => DdlColType::Int8, + Some(SqlValue::Float(_)) => DdlColType::Float8, + Some(SqlValue::Bool(_)) => DdlColType::Bool, + Some( + SqlValue::String(_) + | SqlValue::Null + | SqlValue::Decimal(_) + | SqlValue::Bytes(_) + | SqlValue::Array(_) + | SqlValue::Timestamp(_) + | SqlValue::Timestamptz(_), + ) + | None => DdlColType::Text, + } +} + #[cfg(test)] mod tests { use super::*; @@ -402,18 +431,24 @@ mod tests { } #[test] - fn constant_result_columns_map_to_text_output_columns() { + fn constant_result_columns_are_typed_from_their_values() { + use nodedb_sql::types_expr::SqlValue; let plans = vec![SqlPlan::ConstantResult { - columns: vec!["a".to_string(), "b".to_string()], - values: vec![], + columns: vec!["a".to_string(), "b".to_string(), "c".to_string()], + values: vec![SqlValue::Int(1), SqlValue::String("x".into())], volatile: false, }]; let schema = build_output_schema(&plans, &NoCatalog, nodedb_types::DatabaseId::DEFAULT, None); - assert_eq!(schema.columns.len(), 2); + assert_eq!(schema.columns.len(), 3); assert_eq!(schema.columns[0].display_name, "a"); assert_eq!(schema.columns[0].lookup_key, "a"); + assert_eq!(schema.columns[0].ty, DdlColType::Int8); assert_eq!(schema.columns[1].display_name, "b"); + assert_eq!(schema.columns[1].ty, DdlColType::Text); + // A column without a value keeps its slot and types as text. + assert_eq!(schema.columns[2].display_name, "c"); + assert_eq!(schema.columns[2].ty, DdlColType::Text); assert!(!schema.is_star); } diff --git a/nodedb/src/control/planner/sql_plan_convert/scan/core.rs b/nodedb/src/control/planner/sql_plan_convert/scan/core.rs index 5e56c2697..4367a9719 100644 --- a/nodedb/src/control/planner/sql_plan_convert/scan/core.rs +++ b/nodedb/src/control/planner/sql_plan_convert/scan/core.rs @@ -43,6 +43,9 @@ pub(in crate::control::planner::sql_plan_convert) fn convert_scan( // `rows` is left empty here; the coordinator fills it post-cache via // `materialize_providers`. Using an empty-coordinator vshard (empty // collection string) keeps the task coordinator-local. + let computed_bytes = extract_computed_columns(projection, window_functions, false)?; + let window_bytes = serialize_window_functions(window_functions, false)?; + if crate::control::server::pgwire::catalog::schema::catalog_collection_info(collection) .is_some() { @@ -58,6 +61,8 @@ pub(in crate::control::planner::sql_plan_convert) fn convert_scan( rows: Vec::new(), filters: filter_bytes, projection: proj_names, + computed_columns: computed_bytes, + window_functions: window_bytes, sort_keys: sort, limit: *limit, offset: *offset, @@ -75,8 +80,6 @@ pub(in crate::control::planner::sql_plan_convert) fn convert_scan( let proj_names = extract_projection_names(projection, window_functions); let sort = convert_sort_keys(sort_keys); let vshard = VShardId::from_collection_in_database(database_id, collection); - let computed_bytes = extract_computed_columns(projection, window_functions)?; - let window_bytes = serialize_window_functions(window_functions)?; let physical = match engine { EngineType::Timeseries => { diff --git a/nodedb/src/control/planner/sql_plan_convert/scan/join.rs b/nodedb/src/control/planner/sql_plan_convert/scan/join.rs index f201e513f..2523ad394 100644 --- a/nodedb/src/control/planner/sql_plan_convert/scan/join.rs +++ b/nodedb/src/control/planner/sql_plan_convert/scan/join.rs @@ -108,6 +108,8 @@ pub(in crate::control::planner::sql_plan_convert) fn convert_join( // keys its persisted column stats by the bare collection name, so the // shuffle cost model must look them up by the same raw name (not the // db-qualified token used for storage routing). + // An input-sourced join side carries no routing collection; its rows come + // from `left_input` / `right_input`. let mut left_raw = super::super::aggregate::extract_collection_name(left); let mut right_raw = super::super::aggregate::extract_collection_name(right); let mut left_alias = extract_scan_alias(left); @@ -210,7 +212,11 @@ pub(in crate::control::planner::sql_plan_convert) fn convert_join( ); let shuffle_eligible = structurally_shufflable && (p.ctx.force_shuffle_join - || super::join_cost::cost_model_picks_shuffle(p.ctx, &left_raw, &right_raw)); + || super::join_cost::cost_model_picks_shuffle( + p.ctx, + left_raw.as_deref().unwrap_or(""), + right_raw.as_deref().unwrap_or(""), + )); // Shuffle hash keys mirror the resolver's per-side split: the LEFT column of // each `on` pair partitions the probe side, the RIGHT column the build side diff --git a/nodedb/src/control/planner/sql_plan_convert/scan/timeseries.rs b/nodedb/src/control/planner/sql_plan_convert/scan/timeseries.rs index 8333cd819..0d8521248 100644 --- a/nodedb/src/control/planner/sql_plan_convert/scan/timeseries.rs +++ b/nodedb/src/control/planner/sql_plan_convert/scan/timeseries.rs @@ -64,7 +64,7 @@ pub(in crate::control::planner::sql_plan_convert) fn convert_timeseries_scan( } let proj_names = extract_projection_names(projection, &[]); - let computed_bytes = extract_computed_columns(projection, &[])?; + let computed_bytes = extract_computed_columns(projection, &[], false)?; let vshard = VShardId::from_collection_in_database(ctx.database_id, collection); Ok(vec![PhysicalTask { tenant_id, diff --git a/nodedb/src/control/planner/sql_plan_convert/set_ops.rs b/nodedb/src/control/planner/sql_plan_convert/set_ops.rs index 9e43b00ef..d8beeea04 100644 --- a/nodedb/src/control/planner/sql_plan_convert/set_ops.rs +++ b/nodedb/src/control/planner/sql_plan_convert/set_ops.rs @@ -2,16 +2,18 @@ //! Set operations and miscellaneous plan conversions (UNION, INTERSECT, EXCEPT, CTE, etc.). -use nodedb_sql::types::{Projection, SortKey, SqlExpr, SqlPlan, SqlValue}; +use nodedb_sql::types::{Projection, SortKey, SqlExpr, SqlPlan, SqlValue, WindowSpec}; use crate::bridge::envelope::PhysicalPlan; use crate::types::{TenantId, VShardId}; use nodedb_physical::physical_plan::*; +use super::body::convert_body_to_single_plan; use super::convert::{ConvertContext, convert_one}; use super::expr::inline_cte; -use super::value::sql_value_to_string; +use super::value::{sql_value_to_nodedb_value, sql_value_to_string}; use nodedb_physical::physical_task::{PhysicalTask, PostSetOp}; +use nodedb_types::Value; pub(super) fn convert_constant_result( columns: &[String], @@ -19,25 +21,38 @@ pub(super) fn convert_constant_result( tenant_id: TenantId, ctx: &ConvertContext, ) -> crate::Result> { - // A constant row is one JSON object, which cannot hold two cells under one + // A constant row is one object, which cannot hold two cells under one // key. `SELECT nextval('s'), nextval('s')` legally repeats an output name; // keying both cells by the name would collapse them to the last value. Use // the same unique per-column keys every response encoder derives, so each // column keeps its own cell. let cell_keys = crate::control::server::response_shape::project::cell_keys(columns); - let mut obj = serde_json::Map::new(); - for ((_col, val), key) in columns.iter().zip(values.iter()).zip(cell_keys.iter()) { - let json_val = match val { - SqlValue::Null => serde_json::Value::Null, - other => serde_json::Value::String(sql_value_to_string(other)), + let mut obj = std::collections::HashMap::with_capacity(columns.len()); + for ((_col, val), key) in columns.iter().zip(values.iter()).zip(cell_keys) { + let cell = match val { + SqlValue::Int(_) + | SqlValue::Float(_) + | SqlValue::Bool(_) + | SqlValue::Null + | SqlValue::String(_) => sql_value_to_nodedb_value(val), + // The shaper has no typed renderer that reproduces PostgreSQL's + // text form for these — `\x..` for bytes, `{1,2}` for arrays, the + // ISO string for timestamps — from a typed value, so they keep + // that text form under a `Text` column instead. + SqlValue::Decimal(_) + | SqlValue::Bytes(_) + | SqlValue::Array(_) + | SqlValue::Timestamp(_) + | SqlValue::Timestamptz(_) => Value::String(sql_value_to_string(val)), }; - obj.insert(key.clone(), json_val); + obj.insert(key, cell); } - let arr = serde_json::Value::Array(vec![serde_json::Value::Object(obj)]); - let payload = nodedb_types::json_to_msgpack(&arr).map_err(|e| crate::Error::Serialization { - format: "msgpack".into(), - detail: format!("constant result: {e}"), - })?; + let arr = Value::Array(vec![Value::Object(obj)]); + let payload = + nodedb_types::value_to_msgpack(&arr).map_err(|e| crate::Error::Serialization { + format: "msgpack".into(), + detail: format!("constant result: {e}"), + })?; Ok(vec![PhysicalTask { tenant_id, vshard_id: VShardId::from_collection_in_database(ctx.database_id, ""), @@ -47,6 +62,8 @@ pub(super) fn convert_constant_result( rows: payload, filters: Vec::new(), projection: Vec::new(), + computed_columns: Vec::new(), + window_functions: Vec::new(), sort_keys: Vec::new(), limit: None, offset: 0, @@ -241,9 +258,11 @@ pub(super) fn convert_cte( /// whose leaf could not absorb the outer constraints — into a coordinator- /// resolved `QueryOp::PostProcess`. /// -/// The body is converted to a single physical plan and, when it is a sharded -/// source, wrapped in `Exchange{Gather}` so the sort/distinct/offset/limit tail -/// runs exactly once over the full union at resolve time. +/// The body lowers to ONE physical relation through +/// `convert_body_to_single_plan`: a set-operation body becomes a +/// coordinator-resolved `SetOp`, and a sharded body is wrapped in +/// `Exchange{Gather}` so the sort/distinct/offset/limit tail runs exactly +/// once over the full union at resolve time. pub(super) fn convert_subquery( args: nodedb_sql::SubqueryVisitArgs<'_>, tenant_id: TenantId, @@ -253,59 +272,25 @@ pub(super) fn convert_subquery( input, filters, projection, + window_functions, sort_keys, offset, distinct, limit, } = args; - // Materialize the body as a single physical plan. A subquery/derived-table - // body is one relation; a body that lowers to multiple tasks (e.g. a set - // operation) has no single row stream to post-process here. - let mut body_tasks = convert_one(input, tenant_id, ctx)?; - if body_tasks.len() != 1 { - return Err(crate::Error::PlanError { - detail: format!( - "ORDER BY / OFFSET / DISTINCT over a subquery whose body lowers to {} physical \ - tasks is not supported; the body must produce a single relation", - body_tasks.len() - ), - }); - } - let mut child = body_tasks.pop().expect("checked len == 1").plan; + // The body is ONE relation, already gathered when sharded. + let child = convert_body_to_single_plan(input, tenant_id, ctx)?; // A join / lateral body emits ONE merged document per output row whose // columns keep their table prefix (`a.attnum`), which is why the response - // shaper looks those rows up by the qualified name. The tail's sort keys - // must address the same shape — an unqualified key resolves to NULL on - // every merged row, and a sort where every key is NULL is a no-op that - // silently answers an ordered query in the body's own order. - let merged_doc_body = matches!( - child, - PhysicalPlan::Query( - QueryOp::HashJoin { .. } - | QueryOp::NestedLoopJoin { .. } - | QueryOp::SortMergeJoin { .. } - | QueryOp::LateralTopK { .. } - | QueryOp::LateralLoop { .. } - ) - ); - - // A sharded body must be gathered before the relational tail runs, so the - // sort/distinct/offset/limit observe the FULL union exactly once. - // PostProcess is itself coordinator-local (`is_sharded_source() == false`), - // so the top-level `convert()` wrap loop will not gather the child for us. - if child.is_sharded_source() { - let as_aggregate = matches!( - &child, - PhysicalPlan::Query(QueryOp::Aggregate { .. }) - | PhysicalPlan::Query(QueryOp::PartialAggregate { .. }) - ); - child = PhysicalPlan::Query(QueryOp::Exchange(ExchangeOp { - child: Box::new(child), - mode: ExchangeMode::Gather { as_aggregate }, - })); - } + // shaper looks those rows up by the qualified name. The tail's sort keys, + // computed columns, and window specs must address the same shape — an + // unqualified key resolves to NULL on every merged row, and a sort where + // every key is NULL is a no-op that silently answers an ordered query in + // the body's own order. The body may sit under the `Exchange{Gather}` + // wrapper, so the detection looks through it. + let merged_doc_body = is_merged_doc_body(&child); Ok(vec![PhysicalTask { tenant_id, @@ -316,7 +301,16 @@ pub(super) fn convert_subquery( plan: PhysicalPlan::Query(QueryOp::PostProcess { input: Box::new(child), filters: super::filter::serialize_filters(filters)?, - projection: lower_subquery_projection(projection)?, + projection: lower_subquery_projection(projection, window_functions)?, + computed_columns: super::aggregate::extract_computed_columns( + projection, + window_functions, + merged_doc_body, + )?, + window_functions: super::aggregate::serialize_window_functions( + window_functions, + merged_doc_body, + )?, sort_keys: lower_subquery_sort_keys(sort_keys, merged_doc_body), limit, offset, @@ -327,19 +321,65 @@ pub(super) fn convert_subquery( }]) } +/// Whether a body plan is a join / lateral whose rows keep their table +/// prefix on every column, looking through the converter's +/// `Exchange{Gather}` wrapper. +fn is_merged_doc_body(plan: &PhysicalPlan) -> bool { + match plan { + PhysicalPlan::Query(QueryOp::Exchange(ExchangeOp { child, .. })) => { + is_merged_doc_body(child) + } + PhysicalPlan::Query( + QueryOp::HashJoin { .. } + | QueryOp::NestedLoopJoin { .. } + | QueryOp::SortMergeJoin { .. } + | QueryOp::LateralTopK { .. } + | QueryOp::LateralLoop { .. }, + ) => true, + PhysicalPlan::Query( + QueryOp::ProviderScan { .. } + | QueryOp::PostProcess { .. } + | QueryOp::SetOp { .. } + | QueryOp::Aggregate { .. } + | QueryOp::PartialAggregate { .. } + | QueryOp::PartialAggregateState { .. } + | QueryOp::ShuffleJoinConsume { .. } + | QueryOp::ShuffleAggregateConsume { .. } + | QueryOp::FacetCounts { .. } + | QueryOp::RecursiveScan { .. } + | QueryOp::RecursiveValue { .. }, + ) + | PhysicalPlan::Document(_) + | PhysicalPlan::Vector(_) + | PhysicalPlan::Graph(_) + | PhysicalPlan::Text(_) + | PhysicalPlan::Columnar(_) + | PhysicalPlan::Timeseries(_) + | PhysicalPlan::Spatial(_) + | PhysicalPlan::Kv(_) + | PhysicalPlan::Crdt(_) + | PhysicalPlan::Meta(_) + | PhysicalPlan::Array(_) + | PhysicalPlan::ClusterArray(_) + | PhysicalPlan::ClusterEvent(_) => false, + } +} + /// Lower outer projection items to the row keys the relational tail matches. /// /// A bare column keeps its unqualified name (the flattened row's column key); a /// star selects every column, so no column pruning is applied (empty = all). /// -/// A computed item is projected under its alias: the body evaluates the -/// expression and emits the value under that name before the tail runs, which -/// is the same key the response shaper reads it back by. Erroring here instead -/// would reject `SELECT a, f(b) … ORDER BY c` outright, since the wrapper -/// carries the original SELECT list whenever the body had to be widened to keep -/// the sort column. -fn lower_subquery_projection(projection: &[Projection]) -> crate::Result> { - let mut names = Vec::with_capacity(projection.len()); +/// A computed item is projected under its alias: the tail evaluates the +/// expression over the materialized rows and emits the value under that name, +/// which is the same key the response shaper reads it back by. Every window +/// alias is kept too, so a window output the SELECT list does not repeat as a +/// computed entry survives the column pruning. +fn lower_subquery_projection( + projection: &[Projection], + window_functions: &[WindowSpec], +) -> crate::Result> { + let mut names = Vec::with_capacity(projection.len() + window_functions.len()); for p in projection { match p { Projection::Column(qname) => { @@ -349,6 +389,11 @@ fn lower_subquery_projection(projection: &[Projection]) -> crate::Result names.push(alias.clone()), } } + for spec in window_functions { + if !names.contains(&spec.alias) { + names.push(spec.alias.clone()); + } + } Ok(names) } diff --git a/nodedb/src/control/security/identity/plan_permission.rs b/nodedb/src/control/security/identity/plan_permission.rs index ec68a472c..7461d7132 100644 --- a/nodedb/src/control/security/identity/plan_permission.rs +++ b/nodedb/src/control/security/identity/plan_permission.rs @@ -92,6 +92,14 @@ pub fn required_permission(plan: &crate::bridge::envelope::PhysicalPlan) -> Perm // PostProcess only reshapes child rows; permission is the child's — recurse, don't assume Read. PhysicalPlan::Query(QueryOp::PostProcess { input, .. }) => required_permission(input), + // SetOp merges N child relations; it requires the strictest permission + // any branch requires. An empty input list is unreachable from the + // converter and maps to Read, the weakest tier. + PhysicalPlan::Query(QueryOp::SetOp { inputs, .. }) => inputs + .iter() + .map(required_permission) + .fold(Permission::Read, strictest), + PhysicalPlan::Text( TextOp::Search { .. } | TextOp::BM25ScoreScan { .. } @@ -363,3 +371,30 @@ pub fn required_permission(plan: &crate::bridge::envelope::PhysicalPlan) -> Perm } } } + +/// The stricter of two permissions under the tier order used to fold a +/// multi-input node. Exhaustive so a new `Permission` variant forces a +/// placement here. +fn strictest(a: Permission, b: Permission) -> Permission { + if strictness_rank(b) > strictness_rank(a) { + b + } else { + a + } +} + +/// Tier order from weakest to strictest. Read-class tiers come first, then +/// write, then schema, then cluster-wide control. +fn strictness_rank(permission: Permission) -> u8 { + match permission { + Permission::Read => 0, + Permission::Monitor => 1, + Permission::Execute => 2, + Permission::Write => 3, + Permission::Create => 4, + Permission::Drop => 5, + Permission::Alter => 6, + Permission::Backup => 7, + Permission::Admin => 8, + } +} diff --git a/nodedb/src/control/server/exchange/resolve/exchange/aggregate_input_arm.rs b/nodedb/src/control/server/exchange/resolve/exchange/aggregate_input_arm.rs new file mode 100644 index 000000000..63667c288 --- /dev/null +++ b/nodedb/src/control/server/exchange/resolve/exchange/aggregate_input_arm.rs @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Input-sourced `Aggregate` resolution: materialize the aggregate's child on +//! the coordinator, then hand the Data Plane an aggregate over a +//! `ProviderScan` of those rows. + +use nodedb_physical::physical_plan::{ + AggregateSpec, GroupKeySpec, PhysicalPlan, QueryOp, SortKeySpec, +}; +use nodedb_types::QualifiedCollection; + +use crate::control::server::exchange::resolve::capture::DistributedReadCapture; +use crate::control::state::SharedState; + +use super::dispatch::ResolveCtx; +use super::entry::Resolved; +use super::post_process_arm::{ChildRows, materialize_child_rows, provider_scan_of_rows}; + +/// Fields of a `QueryOp::Aggregate { input: Some(_) }` plan node, carried +/// through resolution as one value. +pub(super) struct AggregateFields { + pub collection: QualifiedCollection, + pub input: Box, + pub group_by: Vec, + pub aggregates: Vec, + pub filters: Vec, + pub having: Vec, + pub limit: usize, + pub sub_group_by: Vec, + pub sub_aggregates: Vec, + pub grouping_sets: Vec>, + pub sort_keys: Vec, +} + +/// Resolve an input-sourced `QueryOp::Aggregate`. +/// +/// A child that is already a materialized `ProviderScan{provider: None}` +/// (a catalog source filled by pass 1) passes through unchanged. Any other +/// child — an `Exchange{Gather}` over a sharded body, a `PostProcess`, a +/// constant result — is materialized on the coordinator and replaced by a +/// `ProviderScan` over its rows, so the aggregate runs exactly once over the +/// full relation and no Exchange reaches a Data-Plane core. +pub(super) async fn resolve_aggregate_input( + state: &SharedState, + ctx: ResolveCtx, + captures: &mut Vec, + fields: AggregateFields, +) -> crate::Result { + let AggregateFields { + collection, + input, + group_by, + aggregates, + filters, + having, + limit, + sub_group_by, + sub_aggregates, + grouping_sets, + sort_keys, + } = fields; + + let rebuild = |input: Box| { + Resolved::Plan(Box::new(PhysicalPlan::Query(QueryOp::Aggregate { + collection, + input: Some(input), + group_by, + aggregates, + filters, + having, + limit, + sub_group_by, + sub_aggregates, + grouping_sets, + sort_keys, + }))) + }; + + // Fast path: the child is already materialized rows. + if matches!( + *input, + PhysicalPlan::Query(QueryOp::ProviderScan { provider: None, .. }) + ) { + return Ok(rebuild(input)); + } + + let rows = match materialize_child_rows(state, ctx, captures, *input).await? { + ChildRows::Rows(rows) => rows, + ChildRows::Passthrough(resolved) => return Ok(resolved), + }; + Ok(rebuild(Box::new(provider_scan_of_rows(rows)))) +} diff --git a/nodedb/src/control/server/exchange/resolve/exchange/dispatch.rs b/nodedb/src/control/server/exchange/resolve/exchange/dispatch.rs index c312c65c4..29212c4a4 100644 --- a/nodedb/src/control/server/exchange/resolve/exchange/dispatch.rs +++ b/nodedb/src/control/server/exchange/resolve/exchange/dispatch.rs @@ -9,10 +9,13 @@ use crate::control::server::exchange::resolve::capture::DistributedReadCapture; use crate::control::state::SharedState; use crate::types::{DatabaseId, TenantId, TraceId, TxnId}; +use super::aggregate_input_arm::AggregateFields; use super::entry::Resolved; use super::hash_join_arm::HashJoinFields; use super::post_process_arm::PostProcessFields; -use super::{gather_arm, hash_join_arm, post_process_arm, shuffle_arm}; +use super::{ + aggregate_input_arm, gather_arm, hash_join_arm, post_process_arm, set_op_arm, shuffle_arm, +}; /// Request-scoped identifiers threaded through every arm resolver, bundled /// to keep each resolver's argument list within the clippy default arity. @@ -32,6 +35,10 @@ pub(super) struct ResolveCtx { /// - Root-level `Shuffle` wrapping a `HashJoin` → orchestrate a cross-node /// grace hash join, return `Resolved::Gathered`. `Shuffle` as a join input is /// a typed error. +/// - `Aggregate{input: Some}` → materialize the child on the coordinator and +/// embed it as `ProviderScan{None, rows}`, return `Resolved::Plan`. +/// - `SetOp{inputs, op}` → materialize every branch, merge with `op`, and +/// embed as `ProviderScan{None, rows}`, return `Resolved::Plan`. /// - Anything else → `Resolved::Plan` unchanged. /// /// `captures` accumulates one [`DistributedReadCapture`] per base collection an @@ -157,6 +164,8 @@ pub(super) async fn resolve_exchange( input, filters, projection, + computed_columns, + window_functions, sort_keys, limit, offset, @@ -170,6 +179,8 @@ pub(super) async fn resolve_exchange( input, filters, projection, + computed_columns, + window_functions, sort_keys, limit, offset, @@ -179,6 +190,51 @@ pub(super) async fn resolve_exchange( .await } + // Input-sourced Aggregate: materialize the child on the coordinator + // (unless it is already a `ProviderScan` of rows) and aggregate over + // those rows once. The aggregate is coordinator-local, so the root + // Gather arm never sees it; the child is gathered here instead. + PhysicalPlan::Query(QueryOp::Aggregate { + collection, + input: Some(input), + group_by, + aggregates, + filters, + having, + limit, + sub_group_by, + sub_aggregates, + grouping_sets, + sort_keys, + }) => { + aggregate_input_arm::resolve_aggregate_input( + state, + ctx, + captures, + AggregateFields { + collection, + input, + group_by, + aggregates, + filters, + having, + limit, + sub_group_by, + sub_aggregates, + grouping_sets, + sort_keys, + }, + ) + .await + } + + // SetOp: materialize every branch on the coordinator, merge with the + // set operation, and lower to a `ProviderScan` of the merged rows. + // The node is coordinator-local, so the root Gather arm never sees it. + PhysicalPlan::Query(QueryOp::SetOp { inputs, op }) => { + set_op_arm::resolve_set_op(state, ctx, captures, inputs, op).await + } + // All other plan variants: pass through unchanged. other => Ok(Resolved::Plan(Box::new(other))), } diff --git a/nodedb/src/control/server/exchange/resolve/exchange/mod.rs b/nodedb/src/control/server/exchange/resolve/exchange/mod.rs index 2bf036327..7356e5fb5 100644 --- a/nodedb/src/control/server/exchange/resolve/exchange/mod.rs +++ b/nodedb/src/control/server/exchange/resolve/exchange/mod.rs @@ -12,13 +12,19 @@ //! cross-node grace hash join (`super::shuffle`) and return the merged rows //! as `Resolved::Gathered`. `Shuffle` as a join INPUT is a typed error (it //! only ever wraps a complete join). +//! - `Aggregate{input: Some}` whose child is not yet materialized rows → +//! materialize the child on the coordinator and embed it as +//! `ProviderScan{provider: None, rows}`; return `Resolved::Plan`. //! - No Exchange / no empty ProviderScan → `Resolved::Plan` unchanged. +mod aggregate_input_arm; mod dispatch; mod entry; mod gather_arm; mod hash_join_arm; mod post_process_arm; +mod set_op_arm; mod shuffle_arm; pub use entry::{Resolved, resolve_and_materialize, resolve_exchange_in_plan}; +pub(crate) use post_process_arm::provider_scan_of_rows; diff --git a/nodedb/src/control/server/exchange/resolve/exchange/post_process_arm.rs b/nodedb/src/control/server/exchange/resolve/exchange/post_process_arm.rs index 9e4920724..8230cffec 100644 --- a/nodedb/src/control/server/exchange/resolve/exchange/post_process_arm.rs +++ b/nodedb/src/control/server/exchange/resolve/exchange/post_process_arm.rs @@ -5,12 +5,14 @@ use nodedb_physical::physical_plan::{ ExchangeMode, ExchangeOp, PhysicalPlan, QueryOp, SortKeySpec, TextOp, VectorOp, + plan_contains_cluster_partitioned_leaf, }; use crate::control::server::exchange::full_scan::{ScanSide, full_scan_plan_for_collection}; use crate::control::server::exchange::gather::{ GatherOutcome, finalize_aggregate, gather_all_vshards, }; +use crate::control::server::exchange::owning_core::gather_single_owning_core; use crate::control::server::exchange::resolve::capture::DistributedReadCapture; use crate::control::server::response_translate::hit_key::parse_surrogate_hex; use crate::control::server::response_translate::vector::resolve_surrogate_pk; @@ -19,6 +21,7 @@ use crate::data::executor::response_codec::{ flatten_hybrid_hits_to_relational_rows, flatten_to_relational_rows, flatten_vector_hits_to_relational_rows, }; +use crate::types::VShardId; use super::dispatch::{ResolveCtx, resolve_exchange}; use super::entry::Resolved; @@ -29,6 +32,8 @@ pub(super) struct PostProcessFields { pub input: Box, pub filters: Vec, pub projection: Vec, + pub computed_columns: Vec, + pub window_functions: Vec, pub sort_keys: Vec, pub limit: Option, pub offset: usize, @@ -86,37 +91,59 @@ fn hit_collection_name(plan: &PhysicalPlan) -> Option { } } -/// Resolve a `QueryOp::PostProcess` node: materialize the child's rows on the -/// coordinator, then lower to a `ProviderScan` that applies filter → offset → -/// sort → distinct → project → limit on a single core (its existing tail). -/// This keeps "run exactly once over the full union" correct: the child is -/// gathered here, so the relational tail never runs per-shard. -pub(super) async fn resolve_post_process( +/// A `ProviderScan` carrying final `rows` and an empty relational tail: +/// no filter, projection, computed column, window, sort, limit, offset, or +/// distinct. The shape every coordinator-materialized child is embedded as. +pub(crate) fn provider_scan_of_rows(rows: Vec) -> PhysicalPlan { + PhysicalPlan::Query(QueryOp::ProviderScan { + provider: None, + rows, + filters: Vec::new(), + projection: Vec::new(), + computed_columns: Vec::new(), + window_functions: Vec::new(), + sort_keys: Vec::new(), + limit: None, + offset: 0, + distinct: false, + }) +} + +/// Rows of a materialized child, or a resolution the caller returns as-is. +pub(super) enum ChildRows { + /// The child's rows, flattened to the bare relational row shape a + /// `ProviderScan{provider: None}` consumes. + Rows(Vec), + /// The child resolved to a root `Gathered` / `Stream` result. The caller + /// returns it unchanged. + Passthrough(Resolved), +} + +/// Materialize a coordinator-side child plan into relational rows. +/// +/// Unwraps the converter's `Exchange{Gather}` wrapper, resolves any Exchange +/// nested inside the body (a `HashJoin` build-side `Broadcast`), gathers the +/// body across all vShards, finalizes a partial-aggregate payload, and +/// flattens hit-shaped payloads (vector / hybrid) to columned rows with the +/// surrogate resolved to the user PK. An in-transaction read records the +/// child's base collection in `captures` at its observed read-version. +pub(super) async fn materialize_child_rows( state: &SharedState, ctx: ResolveCtx, captures: &mut Vec, - fields: PostProcessFields, -) -> crate::Result { + input: PhysicalPlan, +) -> crate::Result { let ResolveCtx { database_id, tenant_id, trace_id, txn_id, } = ctx; - let PostProcessFields { - input, - filters, - projection, - sort_keys, - limit, - offset, - distinct, - } = fields; // The converter wraps a sharded body in `Exchange{Gather}`; unwrap // it so the child is the real body plan (a plain body has no // wrapper and routes to its owning vShard directly). - let (child, as_aggregate) = match *input { + let (child, as_aggregate) = match input { PhysicalPlan::Query(QueryOp::Exchange(ExchangeOp { child, mode: ExchangeMode::Gather { as_aggregate }, @@ -140,11 +167,11 @@ pub(super) async fn resolve_post_process( { Resolved::Plan(p) => *p, // The unwrapped body is not itself a root Gather / stream; - // surface these defensively without dropping post-processing. + // surface these without dropping the caller's tail. Resolved::Gathered(resp, wms, caps) => { - return Ok(Resolved::Gathered(resp, wms, caps)); + return Ok(ChildRows::Passthrough(Resolved::Gathered(resp, wms, caps))); } - Resolved::Stream(s) => return Ok(Resolved::Stream(s)), + Resolved::Stream(s) => return Ok(ChildRows::Passthrough(Resolved::Stream(s))), }; // Classify the body's row shape so the gathered payload is @@ -174,8 +201,27 @@ pub(super) async fn resolve_post_process( None }; - let outcome: GatherOutcome = - gather_all_vshards(state, tenant_id, database_id, child, trace_id, txn_id).await?; + // A coordinator-local body (a `ProviderScan` carrying embedded rows, a + // nested `PostProcess`) reads no per-shard collection. It runs exactly + // once on the coordinator vshard: fanning it to every core returns its + // rows once per core. + let coordinator_local = child.collection().is_none() + && !child.is_sharded_source() + && !plan_contains_cluster_partitioned_leaf(&child); + let outcome: GatherOutcome = if coordinator_local { + gather_single_owning_core( + state, + tenant_id, + database_id, + child, + VShardId::from_collection_in_database(database_id, ""), + trace_id, + txn_id, + ) + .await? + } else { + gather_all_vshards(state, tenant_id, database_id, child, trace_id, txn_id).await? + }; if let Some(coll) = probe_collection && let Some(scan_plan) = full_scan_plan_for_collection( @@ -217,12 +263,44 @@ pub(super) async fn resolve_post_process( }), HitShape::None => flatten_to_relational_rows(&merged), }; + Ok(ChildRows::Rows(rows)) +} + +/// Resolve a `QueryOp::PostProcess` node: materialize the child's rows on the +/// coordinator, then lower to a `ProviderScan` that applies filter → offset → +/// sort → distinct → project → limit on a single core (its existing tail). +/// This keeps "run exactly once over the full union" correct: the child is +/// gathered here, so the relational tail never runs per-shard. +pub(super) async fn resolve_post_process( + state: &SharedState, + ctx: ResolveCtx, + captures: &mut Vec, + fields: PostProcessFields, +) -> crate::Result { + let PostProcessFields { + input, + filters, + projection, + computed_columns, + window_functions, + sort_keys, + limit, + offset, + distinct, + } = fields; + + let rows = match materialize_child_rows(state, ctx, captures, *input).await? { + ChildRows::Rows(rows) => rows, + ChildRows::Passthrough(resolved) => return Ok(resolved), + }; Ok(Resolved::Plan(Box::new(PhysicalPlan::Query( QueryOp::ProviderScan { provider: None, rows, filters, projection, + computed_columns, + window_functions, sort_keys, limit, offset, diff --git a/nodedb/src/control/server/exchange/resolve/exchange/set_op_arm.rs b/nodedb/src/control/server/exchange/resolve/exchange/set_op_arm.rs new file mode 100644 index 000000000..814038084 --- /dev/null +++ b/nodedb/src/control/server/exchange/resolve/exchange/set_op_arm.rs @@ -0,0 +1,120 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! `SetOp` exchange resolution: materialize every branch on the coordinator, +//! merge the rows with the set operation, and lower to a `ProviderScan`. + +use nodedb_physical::physical_plan::{PhysicalPlan, SetOpKind}; + +use crate::control::server::exchange::resolve::capture::DistributedReadCapture; +use crate::control::server::payload_merge::merge_msgpack_arrays; +use crate::control::server::set_op_merge::{ + SetMergeMode, dedup_union_payloads, merge_set_op_payloads, +}; +use crate::control::state::SharedState; + +use super::dispatch::ResolveCtx; +use super::entry::Resolved; +use super::post_process_arm::{ChildRows, materialize_child_rows, provider_scan_of_rows}; + +/// Resolve a `QueryOp::SetOp` node. +/// +/// Every branch is an independent body, so all of them materialize +/// concurrently, each into its own read-capture list. The captures are then +/// appended to `captures` in branch order, so the in-transaction read-set +/// sees every branch's base collection exactly once. A branch that resolves +/// to a root `Gathered` / `Stream` result is returned as-is, the way the +/// post-processor returns it. +/// +/// The merged rows are embedded as a `ProviderScan{provider: None}` with an +/// empty relational tail; the enclosing post-processor or input-sourced +/// aggregate supplies its own tail over these rows. +pub(super) async fn resolve_set_op( + state: &SharedState, + ctx: ResolveCtx, + captures: &mut Vec, + inputs: Vec, + op: SetOpKind, +) -> crate::Result { + let branches = inputs.into_iter().map(|input| async move { + let mut branch_captures = Vec::new(); + let rows = materialize_child_rows(state, ctx, &mut branch_captures, input).await?; + Ok::<_, crate::Error>((rows, branch_captures)) + }); + let materialized = futures::future::try_join_all(branches).await?; + + let mut payloads = Vec::with_capacity(materialized.len()); + for (rows, branch_captures) in materialized { + captures.extend(branch_captures); + match rows { + ChildRows::Rows(rows) => payloads.push(rows), + ChildRows::Passthrough(resolved) => return Ok(resolved), + } + } + + let merged = merge_set_op_rows(&payloads, op); + Ok(Resolved::Plan(Box::new(provider_scan_of_rows(merged)))) +} + +/// Merge materialized branch payloads with `op`. Each payload is one msgpack +/// array of flat row maps. The mapping from kind to merge mirrors the pgwire +/// per-task set-op path: `INTERSECT [ALL]` and `EXCEPT [ALL]` share one +/// value-keyed merge each, `UNION` dedups on raw bytes, and `UNION ALL` +/// concatenates. +fn merge_set_op_rows(payloads: &[Vec], op: SetOpKind) -> Vec { + match op { + SetOpKind::UnionAll => merge_msgpack_arrays(payloads), + SetOpKind::UnionDistinct => dedup_union_payloads(payloads), + SetOpKind::Intersect | SetOpKind::IntersectAll => { + merge_set_op_payloads(payloads, SetMergeMode::Intersect) + } + SetOpKind::Except | SetOpKind::ExceptAll => { + merge_set_op_payloads(payloads, SetMergeMode::Except) + } + } +} + +#[cfg(test)] +mod tests { + use super::merge_set_op_rows; + use nodedb_physical::physical_plan::SetOpKind; + + fn encode_array(rows: &[serde_json::Value]) -> Vec { + nodedb_types::json_to_msgpack(&serde_json::Value::Array(rows.to_vec())).unwrap() + } + + fn decode(payload: &[u8]) -> String { + crate::data::executor::response_codec::decode_payload_to_json(payload) + } + + #[test] + fn union_all_keeps_every_row_in_branch_order() { + let left = encode_array(&[serde_json::json!({"x": 1}), serde_json::json!({"x": 2})]); + let right = encode_array(&[serde_json::json!({"x": 2})]); + let merged = merge_set_op_rows(&[left, right], SetOpKind::UnionAll); + assert_eq!(decode(&merged), r#"[{"x":1},{"x":2},{"x":2}]"#); + } + + #[test] + fn union_distinct_drops_duplicates() { + let left = encode_array(&[serde_json::json!({"x": 1}), serde_json::json!({"x": 2})]); + let right = encode_array(&[serde_json::json!({"x": 2})]); + let merged = merge_set_op_rows(&[left, right], SetOpKind::UnionDistinct); + assert_eq!(decode(&merged), r#"[{"x":1},{"x":2}]"#); + } + + #[test] + fn intersect_keeps_rows_present_in_every_branch() { + let left = encode_array(&[serde_json::json!({"x": 1}), serde_json::json!({"x": 2})]); + let right = encode_array(&[serde_json::json!({"x": 2}), serde_json::json!({"x": 3})]); + let merged = merge_set_op_rows(&[left, right], SetOpKind::Intersect); + assert_eq!(decode(&merged), r#"[{"x":2}]"#); + } + + #[test] + fn except_drops_rows_present_in_later_branches() { + let left = encode_array(&[serde_json::json!({"x": 1}), serde_json::json!({"x": 2})]); + let right = encode_array(&[serde_json::json!({"x": 2})]); + let merged = merge_set_op_rows(&[left, right], SetOpKind::Except); + assert_eq!(decode(&merged), r#"[{"x":1}]"#); + } +} diff --git a/nodedb/src/control/server/exchange/resolve/join_input.rs b/nodedb/src/control/server/exchange/resolve/join_input.rs index 16a50d81b..a0de9f7e9 100644 --- a/nodedb/src/control/server/exchange/resolve/join_input.rs +++ b/nodedb/src/control/server/exchange/resolve/join_input.rs @@ -15,6 +15,7 @@ use crate::control::server::exchange::gather::{ }; use super::capture::DistributedReadCapture; +use super::exchange::provider_scan_of_rows; /// Resolve a `HashJoin` input slot. /// @@ -55,16 +56,8 @@ pub(super) async fn resolve_join_input( // Response as a msgpack array — so the two shapes match. let outcome = gather_all_cores(state, tenant_id, database_id, *child, trace_id, txn_id).await?; - let provider_scan = PhysicalPlan::Query(QueryOp::ProviderScan { - provider: None, - rows: flatten_to_relational_rows(&outcome.merged_array), - filters: Vec::new(), - projection: Vec::new(), - sort_keys: Vec::new(), - limit: None, - offset: 0, - distinct: false, - }); + let provider_scan = + provider_scan_of_rows(flatten_to_relational_rows(&outcome.merged_array)); Ok(Some(Box::new(provider_scan))) } @@ -129,16 +122,7 @@ pub(super) async fn resolve_join_input( } else { outcome.merged_array }; - let provider_scan = PhysicalPlan::Query(QueryOp::ProviderScan { - provider: None, - rows: flatten_to_relational_rows(&merged), - filters: Vec::new(), - projection: Vec::new(), - sort_keys: Vec::new(), - limit: None, - offset: 0, - distinct: false, - }); + let provider_scan = provider_scan_of_rows(flatten_to_relational_rows(&merged)); Ok(Some(Box::new(provider_scan))) } @@ -231,14 +215,7 @@ pub(super) async fn gather_join_build_side( }); } - Ok(Some(Box::new(PhysicalPlan::Query(QueryOp::ProviderScan { - provider: None, - rows: flatten_to_relational_rows(&outcome.merged_array), - filters: Vec::new(), - projection: Vec::new(), - sort_keys: Vec::new(), - limit: None, - offset: 0, - distinct: false, - })))) + Ok(Some(Box::new(provider_scan_of_rows( + flatten_to_relational_rows(&outcome.merged_array), + )))) } diff --git a/nodedb/src/control/server/exchange/resolve/materialize.rs b/nodedb/src/control/server/exchange/resolve/materialize.rs index 1d2a2c417..009d8636f 100644 --- a/nodedb/src/control/server/exchange/resolve/materialize.rs +++ b/nodedb/src/control/server/exchange/resolve/materialize.rs @@ -34,6 +34,8 @@ pub(super) async fn materialize_providers( rows: _, filters, projection, + computed_columns, + window_functions, sort_keys, limit, offset, @@ -46,6 +48,8 @@ pub(super) async fn materialize_providers( rows: encoded, filters, projection, + computed_columns, + window_functions, sort_keys, limit, offset, @@ -223,6 +227,8 @@ pub(super) async fn materialize_providers( input, filters, projection, + computed_columns, + window_functions, sort_keys, limit, offset, @@ -233,6 +239,8 @@ pub(super) async fn materialize_providers( input: Box::new(input), filters, projection, + computed_columns, + window_functions, sort_keys, limit, offset, @@ -240,6 +248,16 @@ pub(super) async fn materialize_providers( })) } + // SetOp: recurse into every branch so nested catalog providers are + // filled before the set-op resolver materializes the branches. + PhysicalPlan::Query(QueryOp::SetOp { inputs, op }) => { + let mut filled = Vec::with_capacity(inputs.len()); + for input in inputs { + filled.push(Box::pin(materialize_providers(state, identity, input)).await?); + } + Ok(PhysicalPlan::Query(QueryOp::SetOp { inputs: filled, op })) + } + // All other variants: no catalog providers can be nested here — // pass through unchanged. other => Ok(other), diff --git a/nodedb/src/control/server/mod.rs b/nodedb/src/control/server/mod.rs index 06a16523d..b2653b9db 100644 --- a/nodedb/src/control/server/mod.rs +++ b/nodedb/src/control/server/mod.rs @@ -21,6 +21,7 @@ pub mod response_shape; pub mod response_translate; pub mod result_stream; pub mod session_auth; +pub mod set_op_merge; pub mod shared; pub mod shuffle; pub mod surrogate_exchange; diff --git a/nodedb/src/control/server/pgwire/handler/routing/set_ops.rs b/nodedb/src/control/server/pgwire/handler/routing/set_ops.rs index 2baaaa9a3..43d15f976 100644 --- a/nodedb/src/control/server/pgwire/handler/routing/set_ops.rs +++ b/nodedb/src/control/server/pgwire/handler/routing/set_ops.rs @@ -1,8 +1,7 @@ // SPDX-License-Identifier: BUSL-1.1 -//! Set operation payload merging: UNION DISTINCT, INTERSECT, EXCEPT. -//! -//! Operates on raw msgpack payloads — no decode/re-encode round-trip. +//! Set operation payload merging for pgwire: UNION DISTINCT, INTERSECT, +//! EXCEPT over collected per-task payloads, shaped into a pgwire response. use pgwire::api::results::{FieldFormat, Response}; use pgwire::error::PgWireResult; @@ -12,6 +11,9 @@ use nodedb_physical::physical_task::PostSetOp; use crate::control::server::response_shape::compose::{self, ShapeOutcome}; use crate::control::server::response_shape::redaction::RedactionCtx; use crate::control::server::response_shape::schema::OutputSchema; +use crate::control::server::set_op_merge::{ + SetMergeMode, dedup_union_payloads, merge_set_op_payloads, +}; use super::super::super::types::sqlstate_error; use super::super::plan::{PlanKind, multirow_payload_to_response}; @@ -49,323 +51,3 @@ pub(super) fn apply_set_ops( }, ) } - -/// Merge multiple Data Plane response payloads and deduplicate rows (UNION DISTINCT). -/// -/// Each payload is a msgpack-encoded array of rows. Deduplication is performed -/// at the binary level: each row's raw msgpack bytes serve as the canonical key, -/// eliminating the decode → JSON string → re-encode round-trip. -/// -/// Output: a single msgpack array containing all unique rows in encounter order. -fn dedup_union_payloads(payloads: &[Vec]) -> Vec { - use nodedb_query::msgpack_scan; - - let mut seen: std::collections::HashSet> = std::collections::HashSet::new(); - let mut unique_row_bytes: Vec> = Vec::new(); - - for payload in payloads { - if payload.is_empty() { - continue; - } - - let bytes = payload.as_slice(); - let first = bytes[0]; - - let (count, hdr_len) = if (0x90..=0x9f).contains(&first) { - ((first & 0x0f) as usize, 1) - } else if first == 0xdc && bytes.len() >= 3 { - (u16::from_be_bytes([bytes[1], bytes[2]]) as usize, 3) - } else if first == 0xdd && bytes.len() >= 5 { - ( - u32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]) as usize, - 5, - ) - } else { - tracing::warn!( - payload_len = bytes.len(), - "dedup_union_payloads: payload is not a msgpack array; treating as single row" - ); - let key = bytes.to_vec(); - if seen.insert(key.clone()) { - unique_row_bytes.push(key); - } - continue; - }; - - let mut pos = hdr_len; - for _ in 0..count { - if pos >= bytes.len() { - break; - } - let elem_start = pos; - match msgpack_scan::skip_value(bytes, pos) { - Some(next_pos) => { - let row_bytes = bytes[elem_start..next_pos].to_vec(); - if seen.insert(row_bytes.clone()) { - unique_row_bytes.push(row_bytes); - } - pos = next_pos; - } - None => { - tracing::warn!( - pos, - payload_len = bytes.len(), - "dedup_union_payloads: could not skip msgpack element; stopping early" - ); - break; - } - } - } - } - - let row_count = unique_row_bytes.len(); - let total_data: usize = unique_row_bytes.iter().map(|r| r.len()).sum(); - let mut out = Vec::with_capacity(total_data + 5); - write_array_header(&mut out, row_count); - for row in unique_row_bytes { - out.extend_from_slice(&row); - } - out -} - -enum SetMergeMode { - Intersect, - Except, -} - -/// Merge payloads for INTERSECT or EXCEPT set operations. -/// -/// For INTERSECT: keep rows that appear in ALL payloads. -/// For EXCEPT: keep rows from first payload that don't appear in any subsequent payload. -fn merge_set_op_payloads(payloads: &[Vec], mode: SetMergeMode) -> Vec { - use nodedb_query::msgpack_scan; - - if payloads.is_empty() { - return vec![0x90]; - } - - fn extract_rows(payload: &[u8]) -> Vec> { - if payload.is_empty() { - return Vec::new(); - } - let first = payload[0]; - let (count, hdr_len) = if (0x90..=0x9f).contains(&first) { - ((first & 0x0f) as usize, 1) - } else if first == 0xdc && payload.len() >= 3 { - (u16::from_be_bytes([payload[1], payload[2]]) as usize, 3) - } else if first == 0xdd && payload.len() >= 5 { - ( - u32::from_be_bytes([payload[1], payload[2], payload[3], payload[4]]) as usize, - 5, - ) - } else { - return vec![payload.to_vec()]; - }; - - let mut rows = Vec::with_capacity(count); - let mut pos = hdr_len; - for _ in 0..count { - if pos >= payload.len() { - break; - } - let start = pos; - match msgpack_scan::skip_value(payload, pos) { - Some(next) => { - rows.push(payload[start..next].to_vec()); - pos = next; - } - None => break, - } - } - rows - } - - fn logical_row_bytes(row: &[u8]) -> &[u8] { - msgpack_scan::extract_field(row, 0, "data") - .map(|(start, end)| &row[start..end]) - .unwrap_or(row) - } - - fn write_values_only_key(value: &[u8], out: &mut Vec) -> Option<()> { - if let Some((count, mut pos)) = msgpack_scan::map_header(value, 0) { - write_array_header(out, count); - for _ in 0..count { - pos = msgpack_scan::skip_value(value, pos)?; - let val_start = pos; - pos = msgpack_scan::skip_value(value, pos)?; - write_values_only_key(&value[val_start..pos], out)?; - } - return Some(()); - } - - if let Some((count, mut pos)) = msgpack_scan::array_header(value, 0) { - write_array_header(out, count); - for _ in 0..count { - let elem_start = pos; - pos = msgpack_scan::skip_value(value, pos)?; - write_values_only_key(&value[elem_start..pos], out)?; - } - return Some(()); - } - - out.extend_from_slice(value); - Some(()) - } - - fn extract_value_parts(row: &[u8]) -> Vec> { - let logical = logical_row_bytes(row); - - if let Some((count, mut pos)) = msgpack_scan::map_header(logical, 0) { - let mut parts = Vec::with_capacity(count); - for _ in 0..count { - pos = match msgpack_scan::skip_value(logical, pos) { - Some(next) => next, - None => return vec![logical.to_vec()], - }; - let val_start = pos; - pos = match msgpack_scan::skip_value(logical, pos) { - Some(next) => next, - None => return vec![logical.to_vec()], - }; - let mut normalized = Vec::with_capacity(pos - val_start); - if write_values_only_key(&logical[val_start..pos], &mut normalized).is_none() { - return vec![logical.to_vec()]; - } - parts.push(normalized); - } - return parts; - } - - if let Some((count, mut pos)) = msgpack_scan::array_header(logical, 0) { - let mut parts = Vec::with_capacity(count); - for _ in 0..count { - let elem_start = pos; - pos = match msgpack_scan::skip_value(logical, pos) { - Some(next) => next, - None => return vec![logical.to_vec()], - }; - let mut normalized = Vec::with_capacity(pos - elem_start); - if write_values_only_key(&logical[elem_start..pos], &mut normalized).is_none() { - return vec![logical.to_vec()]; - } - parts.push(normalized); - } - return parts; - } - - vec![logical.to_vec()] - } - - fn extract_values_key(row: &[u8]) -> Vec { - let parts = extract_value_parts(row); - let mut vals = Vec::new(); - write_array_header(&mut vals, parts.len()); - for part in parts { - vals.extend_from_slice(&part); - } - vals - } - - fn rows_match(left: &[u8], right: &[u8]) -> bool { - let left_parts = extract_value_parts(left); - let right_parts = extract_value_parts(right); - let shared_len = left_parts.len().min(right_parts.len()); - - if shared_len == 0 { - return left_parts.is_empty() && right_parts.is_empty(); - } - - left_parts[..shared_len] == right_parts[..shared_len] - && (left_parts.len() == shared_len || right_parts.len() == shared_len) - } - - let first_rows = extract_rows(&payloads[0]); - let mut result_rows: Vec> = match mode { - SetMergeMode::Intersect => { - let other_rows: Vec>> = - payloads[1..].iter().map(|p| extract_rows(p)).collect(); - first_rows - .into_iter() - .filter(|row| { - other_rows - .iter() - .all(|rows| rows.iter().any(|other| rows_match(row, other))) - }) - .map(|row| logical_row_bytes(&row).to_vec()) - .collect() - } - SetMergeMode::Except => { - let other_rows: Vec> = - payloads[1..].iter().flat_map(|p| extract_rows(p)).collect(); - first_rows - .into_iter() - .filter(|row| !other_rows.iter().any(|other| rows_match(row, other))) - .map(|row| logical_row_bytes(&row).to_vec()) - .collect() - } - }; - - let mut seen = std::collections::HashSet::new(); - result_rows.retain(|r| seen.insert(extract_values_key(r))); - - let row_count = result_rows.len(); - let total: usize = result_rows.iter().map(|r| r.len()).sum(); - let mut out = Vec::with_capacity(total + 5); - write_array_header(&mut out, row_count); - for row in result_rows { - out.extend_from_slice(&row); - } - out -} - -fn write_array_header(out: &mut Vec, count: usize) { - if count < 16 { - out.push(0x90 | count as u8); - } else if count <= u16::MAX as usize { - out.push(0xdc); - out.extend_from_slice(&(count as u16).to_be_bytes()); - } else { - out.push(0xdd); - out.extend_from_slice(&(count as u32).to_be_bytes()); - } -} - -#[cfg(test)] -mod tests { - use super::{SetMergeMode, merge_set_op_payloads}; - - fn encode_array(rows: &[serde_json::Value]) -> Vec { - nodedb_types::json_to_msgpack(&serde_json::Value::Array(rows.to_vec())).unwrap() - } - - #[test] - fn intersect_compares_wrapped_rows_by_logical_data_values() { - let left = encode_array(&[ - serde_json::json!({"id":"u1","data":{"id":"u1","name":"Alice"}}), - serde_json::json!({"id":"u2","data":{"id":"u2","name":"Bob"}}), - ]); - let right = encode_array(&[ - serde_json::json!({"id":"doc-1","data":{"user_id":"u1"}}), - serde_json::json!({"id":"doc-2","data":{"user_id":"u3"}}), - ]); - - let merged = merge_set_op_payloads(&[left, right], SetMergeMode::Intersect); - let json = crate::data::executor::response_codec::decode_payload_to_json(&merged); - - assert_eq!(json, r#"[{"id":"u1","name":"Alice"}]"#); - } - - #[test] - fn except_returns_unwrapped_logical_rows() { - let left = encode_array(&[ - serde_json::json!({"id":"u1","data":{"id":"u1"}}), - serde_json::json!({"id":"u2","data":{"id":"u2"}}), - ]); - let right = encode_array(&[serde_json::json!({"id":"doc-1","data":{"user_id":"u1"}})]); - - let merged = merge_set_op_payloads(&[left, right], SetMergeMode::Except); - let json = crate::data::executor::response_codec::decode_payload_to_json(&merged); - - assert_eq!(json, r#"[{"id":"u2"}]"#); - } -} diff --git a/nodedb/src/control/server/response_shape/redaction/query.rs b/nodedb/src/control/server/response_shape/redaction/query.rs index 0b54ecd80..fa5642f4b 100644 --- a/nodedb/src/control/server/response_shape/redaction/query.rs +++ b/nodedb/src/control/server/response_shape/redaction/query.rs @@ -238,6 +238,14 @@ fn collect_sources(plan: &PhysicalPlan, qualifier: &str, out: &mut Vec<(String, collect_sources(input, qualifier, out); return; } + // Every set-operation branch contributes rows under the same + // derived-table qualifier. + QueryOp::SetOp { inputs, .. } => { + for input in inputs { + collect_sources(input, qualifier, out); + } + return; + } QueryOp::Aggregate { collection, input, .. } diff --git a/nodedb/src/control/server/response_shape/types/plan_kind.rs b/nodedb/src/control/server/response_shape/types/plan_kind.rs index 0f68cd00d..1eb811cfb 100644 --- a/nodedb/src/control/server/response_shape/types/plan_kind.rs +++ b/nodedb/src/control/server/response_shape/types/plan_kind.rs @@ -104,6 +104,10 @@ pub fn describe_plan(plan: &PhysicalPlan) -> PlanKind { // PostProcess reshapes a multi-row subquery; its kind is the child's. PhysicalPlan::Query(QueryOp::PostProcess { input, .. }) => describe_plan(input), + // SetOp resolves to a ProviderScan of merged rows; route MultiRow so + // each row streams as its own pgwire row. + PhysicalPlan::Query(QueryOp::SetOp { .. }) => PlanKind::MultiRow, + // An insert with a projection returns real stored rows and must be decoded // and redacted, else it silently leaks unredacted rows like `Merge` did. PhysicalPlan::Kv( diff --git a/nodedb/src/control/server/set_op_merge/intersect_except.rs b/nodedb/src/control/server/set_op_merge/intersect_except.rs new file mode 100644 index 000000000..0dc9c11b0 --- /dev/null +++ b/nodedb/src/control/server/set_op_merge/intersect_except.rs @@ -0,0 +1,102 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! `INTERSECT` / `EXCEPT` merge over msgpack row payloads, comparing rows +//! by logical column values. + +use crate::control::server::payload_merge::{encode_msgpack_array, extract_msgpack_elements}; + +use super::row_key::{extract_values_key, logical_row_bytes, rows_match}; + +/// Which filter [`merge_set_op_payloads`] applies to the first payload. +pub(crate) enum SetMergeMode { + /// Keep rows present in every payload. + Intersect, + /// Keep rows of the first payload absent from every later payload. + Except, +} + +/// Merge payloads for INTERSECT or EXCEPT. +/// +/// Rows compare by value (see `row_key::rows_match`). The output holds the +/// logical row bytes (the `{id, data}` wrapper is stripped), deduplicated by +/// values key, as one msgpack array. No payloads yields an empty array. +pub(crate) fn merge_set_op_payloads(payloads: &[Vec], mode: SetMergeMode) -> Vec { + if payloads.is_empty() { + return vec![0x90]; + } + + let first_rows = extract_msgpack_elements(&payloads[0]); + let mut result_rows: Vec> = match mode { + SetMergeMode::Intersect => { + let other_rows: Vec>> = payloads[1..] + .iter() + .map(|p| extract_msgpack_elements(p)) + .collect(); + first_rows + .into_iter() + .filter(|row| { + other_rows + .iter() + .all(|rows| rows.iter().any(|other| rows_match(row, other))) + }) + .map(|row| logical_row_bytes(&row).to_vec()) + .collect() + } + SetMergeMode::Except => { + let other_rows: Vec> = payloads[1..] + .iter() + .flat_map(|p| extract_msgpack_elements(p)) + .collect(); + first_rows + .into_iter() + .filter(|row| !other_rows.iter().any(|other| rows_match(row, other))) + .map(|row| logical_row_bytes(&row).to_vec()) + .collect() + } + }; + + let mut seen = std::collections::HashSet::with_capacity(result_rows.len()); + result_rows.retain(|r| seen.insert(extract_values_key(r))); + + encode_msgpack_array(&result_rows) +} + +#[cfg(test)] +mod tests { + use super::{SetMergeMode, merge_set_op_payloads}; + + fn encode_array(rows: &[serde_json::Value]) -> Vec { + nodedb_types::json_to_msgpack(&serde_json::Value::Array(rows.to_vec())).unwrap() + } + + #[test] + fn intersect_compares_wrapped_rows_by_logical_data_values() { + let left = encode_array(&[ + serde_json::json!({"id":"u1","data":{"id":"u1","name":"Alice"}}), + serde_json::json!({"id":"u2","data":{"id":"u2","name":"Bob"}}), + ]); + let right = encode_array(&[ + serde_json::json!({"id":"doc-1","data":{"user_id":"u1"}}), + serde_json::json!({"id":"doc-2","data":{"user_id":"u3"}}), + ]); + + let merged = merge_set_op_payloads(&[left, right], SetMergeMode::Intersect); + let json = crate::data::executor::response_codec::decode_payload_to_json(&merged); + + assert_eq!(json, r#"[{"id":"u1","name":"Alice"}]"#); + } + + #[test] + fn except_returns_unwrapped_logical_rows() { + let left = encode_array(&[ + serde_json::json!({"id":"u1","data":{"id":"u1"}}), + serde_json::json!({"id":"u2","data":{"id":"u2"}}), + ]); + let right = encode_array(&[serde_json::json!({"id":"doc-1","data":{"user_id":"u1"}})]); + + let merged = merge_set_op_payloads(&[left, right], SetMergeMode::Except); + let json = crate::data::executor::response_codec::decode_payload_to_json(&merged); + + assert_eq!(json, r#"[{"id":"u2"}]"#); + } +} diff --git a/nodedb/src/control/server/set_op_merge/mod.rs b/nodedb/src/control/server/set_op_merge/mod.rs new file mode 100644 index 000000000..bb2e7716b --- /dev/null +++ b/nodedb/src/control/server/set_op_merge/mod.rs @@ -0,0 +1,13 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Protocol-neutral set-operation merging over msgpack row payloads: +//! `UNION DISTINCT`, `INTERSECT`, `EXCEPT`. Operates on raw msgpack bytes +//! with no decode/re-encode round-trip. Used by the pgwire per-task set-op +//! path and by the coordinator's `QueryOp::SetOp` resolver. + +mod intersect_except; +mod row_key; +mod union; + +pub(crate) use intersect_except::{SetMergeMode, merge_set_op_payloads}; +pub(crate) use union::dedup_union_payloads; diff --git a/nodedb/src/control/server/set_op_merge/row_key.rs b/nodedb/src/control/server/set_op_merge/row_key.rs new file mode 100644 index 000000000..06c06d37e --- /dev/null +++ b/nodedb/src/control/server/set_op_merge/row_key.rs @@ -0,0 +1,129 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Row-identity helpers for set-operation merging: the logical row inside a +//! `{id, data}` wrapper, and the values-only key that compares two rows by +//! column values regardless of column names. + +use nodedb_query::msgpack_scan; + +/// The logical row: the `data` field of an `{id, data}` storage wrapper, or +/// the row itself when it carries no wrapper. +pub(super) fn logical_row_bytes(row: &[u8]) -> &[u8] { + msgpack_scan::extract_field(row, 0, "data") + .map(|(start, end)| &row[start..end]) + .unwrap_or(row) +} + +/// Write a msgpack array header for `count` elements. +pub(super) fn write_array_header(out: &mut Vec, count: usize) { + if count < 16 { + out.push(0x90 | count as u8); + } else if count <= u16::MAX as usize { + out.push(0xdc); + out.extend_from_slice(&(count as u16).to_be_bytes()); + } else { + out.push(0xdd); + out.extend_from_slice(&(count as u32).to_be_bytes()); + } +} + +/// Append `value` to `out` with every map rewritten as an array of its +/// values, recursively, so two rows with different column names but equal +/// values compare equal. `None` when `value` is malformed msgpack. +fn write_values_only_key(value: &[u8], out: &mut Vec) -> Option<()> { + if let Some((count, mut pos)) = msgpack_scan::map_header(value, 0) { + write_array_header(out, count); + for _ in 0..count { + pos = msgpack_scan::skip_value(value, pos)?; + let val_start = pos; + pos = msgpack_scan::skip_value(value, pos)?; + write_values_only_key(&value[val_start..pos], out)?; + } + return Some(()); + } + + if let Some((count, mut pos)) = msgpack_scan::array_header(value, 0) { + write_array_header(out, count); + for _ in 0..count { + let elem_start = pos; + pos = msgpack_scan::skip_value(value, pos)?; + write_values_only_key(&value[elem_start..pos], out)?; + } + return Some(()); + } + + out.extend_from_slice(value); + Some(()) +} + +/// The logical row's column values, each normalized to a values-only key, +/// in column order. A scalar or malformed row yields itself as one part. +pub(super) fn extract_value_parts(row: &[u8]) -> Vec> { + let logical = logical_row_bytes(row); + + if let Some((count, mut pos)) = msgpack_scan::map_header(logical, 0) { + let mut parts = Vec::with_capacity(count); + for _ in 0..count { + pos = match msgpack_scan::skip_value(logical, pos) { + Some(next) => next, + None => return vec![logical.to_vec()], + }; + let val_start = pos; + pos = match msgpack_scan::skip_value(logical, pos) { + Some(next) => next, + None => return vec![logical.to_vec()], + }; + let mut normalized = Vec::with_capacity(pos - val_start); + if write_values_only_key(&logical[val_start..pos], &mut normalized).is_none() { + return vec![logical.to_vec()]; + } + parts.push(normalized); + } + return parts; + } + + if let Some((count, mut pos)) = msgpack_scan::array_header(logical, 0) { + let mut parts = Vec::with_capacity(count); + for _ in 0..count { + let elem_start = pos; + pos = match msgpack_scan::skip_value(logical, pos) { + Some(next) => next, + None => return vec![logical.to_vec()], + }; + let mut normalized = Vec::with_capacity(pos - elem_start); + if write_values_only_key(&logical[elem_start..pos], &mut normalized).is_none() { + return vec![logical.to_vec()]; + } + parts.push(normalized); + } + return parts; + } + + vec![logical.to_vec()] +} + +/// One key per row: its value parts as a msgpack array. +pub(super) fn extract_values_key(row: &[u8]) -> Vec { + let parts = extract_value_parts(row); + let mut vals = Vec::new(); + write_array_header(&mut vals, parts.len()); + for part in parts { + vals.extend_from_slice(&part); + } + vals +} + +/// Whether two rows match by value: their shared column prefix is equal and +/// one row has no columns beyond that prefix. +pub(super) fn rows_match(left: &[u8], right: &[u8]) -> bool { + let left_parts = extract_value_parts(left); + let right_parts = extract_value_parts(right); + let shared_len = left_parts.len().min(right_parts.len()); + + if shared_len == 0 { + return left_parts.is_empty() && right_parts.is_empty(); + } + + left_parts[..shared_len] == right_parts[..shared_len] + && (left_parts.len() == shared_len || right_parts.len() == shared_len) +} diff --git a/nodedb/src/control/server/set_op_merge/union.rs b/nodedb/src/control/server/set_op_merge/union.rs new file mode 100644 index 000000000..4bcb9da77 --- /dev/null +++ b/nodedb/src/control/server/set_op_merge/union.rs @@ -0,0 +1,55 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! `UNION DISTINCT` merge: every row from every payload, deduplicated on +//! its raw msgpack bytes, in encounter order. + +use crate::control::server::payload_merge::{encode_msgpack_array, extract_msgpack_elements}; + +/// Merge multiple row payloads and drop duplicate rows. +/// +/// Each payload is a msgpack array of rows. A row's raw msgpack bytes are +/// its dedup key, so no decode round-trip runs. A payload that is not a +/// msgpack array is treated as one row. The output is one msgpack array of +/// the unique rows in encounter order. +pub(crate) fn dedup_union_payloads(payloads: &[Vec]) -> Vec { + let rows: Vec> = payloads + .iter() + .flat_map(|payload| extract_msgpack_elements(payload)) + .collect(); + let mut seen: std::collections::HashSet> = + std::collections::HashSet::with_capacity(rows.len()); + let mut unique_rows: Vec> = Vec::with_capacity(rows.len()); + + for row in rows { + if seen.insert(row.clone()) { + unique_rows.push(row); + } + } + + encode_msgpack_array(&unique_rows) +} + +#[cfg(test)] +mod tests { + use super::dedup_union_payloads; + + fn encode_array(rows: &[serde_json::Value]) -> Vec { + nodedb_types::json_to_msgpack(&serde_json::Value::Array(rows.to_vec())).unwrap() + } + + #[test] + fn union_distinct_drops_byte_identical_rows_across_payloads() { + let left = encode_array(&[serde_json::json!({"x": 1}), serde_json::json!({"x": 2})]); + let right = encode_array(&[serde_json::json!({"x": 2}), serde_json::json!({"x": 3})]); + + let merged = dedup_union_payloads(&[left, right]); + let json = crate::data::executor::response_codec::decode_payload_to_json(&merged); + + assert_eq!(json, r#"[{"x":1},{"x":2},{"x":3}]"#); + } + + #[test] + fn union_distinct_of_no_payloads_is_empty_array() { + assert_eq!(dedup_union_payloads(&[]), vec![0x90]); + } +} diff --git a/nodedb/src/control/server/shared/authorization/requirements.rs b/nodedb/src/control/server/shared/authorization/requirements.rs index dc7b1e6c6..7655de50f 100644 --- a/nodedb/src/control/server/shared/authorization/requirements.rs +++ b/nodedb/src/control/server/shared/authorization/requirements.rs @@ -73,6 +73,8 @@ mod tests { rows: Vec::new(), filters: Vec::new(), projection: Vec::new(), + computed_columns: Vec::new(), + window_functions: Vec::new(), sort_keys: Vec::new(), limit: None, offset: 0, diff --git a/nodedb/src/control/server/shared/authorization/requirements/query.rs b/nodedb/src/control/server/shared/authorization/requirements/query.rs index 54ab53f8b..58f191e8a 100644 --- a/nodedb/src/control/server/shared/authorization/requirements/query.rs +++ b/nodedb/src/control/server/shared/authorization/requirements/query.rs @@ -77,6 +77,11 @@ pub(super) fn collect_query_requirements<'a>( pending.push(input); true } + // Every set-operation branch is a body whose collections are authorized. + PhysicalPlan::Query(QueryOp::SetOp { inputs, .. }) => { + pending.extend(inputs.iter()); + true + } PhysicalPlan::Query(QueryOp::ProviderScan { provider: Some(provider), .. diff --git a/nodedb/src/control/server/shared/clone_read/temporal.rs b/nodedb/src/control/server/shared/clone_read/temporal.rs index 037d55205..40c8990ca 100644 --- a/nodedb/src/control/server/shared/clone_read/temporal.rs +++ b/nodedb/src/control/server/shared/clone_read/temporal.rs @@ -19,6 +19,9 @@ pub(super) fn extract_system_as_of_ms( PhysicalPlan::Query(QueryOp::PostProcess { input, .. }) => { extract_system_as_of_ms(Some(&**input)) } + PhysicalPlan::Query(QueryOp::SetOp { inputs, .. }) => inputs + .iter() + .find_map(|input| extract_system_as_of_ms(Some(input))), // Index-only/overlay engines carry no qualifier; compose with a data-bearing collection. PhysicalPlan::Vector(_) | PhysicalPlan::Graph(_) diff --git a/nodedb/src/control/server/shared/plan_util.rs b/nodedb/src/control/server/shared/plan_util.rs index c663ea901..044bdffbf 100644 --- a/nodedb/src/control/server/shared/plan_util.rs +++ b/nodedb/src/control/server/shared/plan_util.rs @@ -111,6 +111,12 @@ pub(crate) fn extract_collection(plan: &PhysicalPlan) -> Option<&str> { PhysicalPlan::Query(QueryOp::Exchange(op)) => extract_collection(&op.child), // PostProcess: recurse into the materialized child. PhysicalPlan::Query(QueryOp::PostProcess { input, .. }) => extract_collection(input), + // SetOp: the first branch that names a collection, the same way a + // join reports its left side. Callers that need every branch walk + // the inputs themselves. + PhysicalPlan::Query(QueryOp::SetOp { inputs, .. }) => { + inputs.iter().find_map(extract_collection) + } // ProviderScan is a catalog/constant source — no user collection. PhysicalPlan::Query(QueryOp::ProviderScan { .. }) => None, // KV ops carry their own collection (sorted-index-only ops return None). diff --git a/nodedb/src/control/server/shared/write_admission/predicate/txn_buffering/classify.rs b/nodedb/src/control/server/shared/write_admission/predicate/txn_buffering/classify.rs index 29c9350d4..8c08d289d 100644 --- a/nodedb/src/control/server/shared/write_admission/predicate/txn_buffering/classify.rs +++ b/nodedb/src/control/server/shared/write_admission/predicate/txn_buffering/classify.rs @@ -284,7 +284,8 @@ pub fn plan_requires_txn_buffering(plan: &PhysicalPlan) -> bool { | QueryOp::RecursiveValue { .. } | QueryOp::LateralTopK { .. } | QueryOp::LateralLoop { .. } - | QueryOp::PostProcess { .. }, + | QueryOp::PostProcess { .. } + | QueryOp::SetOp { .. }, ) => false, // ---- Meta: control / maintenance ops — internal orchestration, never a client `task.plan`. @@ -1520,6 +1521,8 @@ mod tests { rows: Vec::new(), filters: Vec::new(), projection: Vec::new(), + computed_columns: Vec::new(), + window_functions: Vec::new(), sort_keys: Vec::new(), limit: None, offset: 0, diff --git a/nodedb/src/data/executor/dispatch/query.rs b/nodedb/src/data/executor/dispatch/query.rs index 180a112d0..46bcb76ae 100644 --- a/nodedb/src/data/executor/dispatch/query.rs +++ b/nodedb/src/data/executor/dispatch/query.rs @@ -67,10 +67,21 @@ impl CoreLoop { }, ), + QueryOp::SetOp { .. } => self.response_error( + task, + crate::bridge::envelope::ErrorCode::Internal { + detail: "SetOp must be resolved by the coordinator (materialized and \ + merged into a ProviderScan) before dispatch" + .to_string(), + }, + ), + QueryOp::ProviderScan { rows, filters, projection, + computed_columns, + window_functions, sort_keys, limit, offset, @@ -82,6 +93,8 @@ impl CoreLoop { rows_bytes: rows, filters_bytes: filters, projection, + computed_columns_bytes: computed_columns, + window_functions_bytes: window_functions, sort_keys, limit: *limit, offset: *offset, diff --git a/nodedb/src/data/executor/handlers/aggregate/exec.rs b/nodedb/src/data/executor/handlers/aggregate/exec.rs index 6629f5c7a..9ee91777f 100644 --- a/nodedb/src/data/executor/handlers/aggregate/exec.rs +++ b/nodedb/src/data/executor/handlers/aggregate/exec.rs @@ -8,7 +8,7 @@ use tracing::debug; use super::cache_key::{AggregateCacheKeyInputs, aggregate_cache_key, legacy_aggregate_pairs}; use super::rows::{apply_user_aliases_to_rows, sort_aggregated_rows}; -use crate::bridge::envelope::{ErrorCode, Response}; +use crate::bridge::envelope::{ErrorCode, Response, Status}; use crate::bridge::scan_filter::ScanFilter; use crate::data::executor::core_loop::CoreLoop; use crate::data::executor::task::ExecutionTask; @@ -58,22 +58,41 @@ impl CoreLoop { debug!(core = self.core_id, %collection, has_input = input.is_some(), group_fields = group_by.len(), aggs = aggregates.len(), "aggregate"); - // Input-sourced aggregate (catalog): the rows come from executing the - // sub-plan (a coordinator-materialized `ProviderScan`), not from a - // per-shard collection scan. Decode the sub-plan rows and aggregate - // over them using the same streaming logic, then short-circuit before - // the per-shard fast paths (cache / index-backed / columnar memtable), - // none of which apply to coordinator-local catalog data. + // Input-sourced aggregate: the rows come from executing the sub-plan + // (a coordinator-materialized `ProviderScan` over catalog rows or a + // derived-table body), not from a per-shard collection scan. Decode + // the sub-plan rows and aggregate over them using the same streaming + // logic, then short-circuit before the per-shard fast paths (cache / + // index-backed / columnar memtable), none of which apply to + // coordinator-local rows. if let Some(sub_plan) = input { let sub_response = self.execute_plan(task, sub_plan); - // Empty / undecodable payload → aggregate over zero rows, the same as - // a per-shard scan that matched nothing. Feeding an empty doc set - // through the shared path keeps behavior identical to the scan path - // rather than surfacing the sub-plan Response (which may be a - // non-row payload). - let docs = - crate::data::executor::response_codec::decode_response_to_docs(&sub_response) - .unwrap_or_default(); + // A child error (22012 from a computed column, a resolver refusal) + // fails the statement. + if sub_response.status == Status::Error { + return sub_response; + } + // An empty payload is zero rows. A non-empty payload that is not + // a MessagePack row array is an internal error, never an empty + // aggregate. + let docs = if sub_response.payload.is_empty() { + Vec::new() + } else { + match crate::data::executor::response_codec::decode_response_to_docs(&sub_response) + { + Some(docs) => docs, + None => { + return self.response_error( + task, + ErrorCode::Internal { + detail: "aggregate input rows failed to decode: payload is not \ + a MessagePack row array" + .to_string(), + }, + ); + } + } + }; return self.aggregate_over_docs( super::streaming::over_docs::AggregateOverDocsParams { task, diff --git a/nodedb/src/data/executor/handlers/mod.rs b/nodedb/src/data/executor/handlers/mod.rs index 40935615e..3e839f49a 100644 --- a/nodedb/src/data/executor/handlers/mod.rs +++ b/nodedb/src/data/executor/handlers/mod.rs @@ -40,6 +40,7 @@ pub(super) mod merge_helpers; pub(super) mod merge_orchestrated; pub mod point; pub(super) mod provider_scan; +pub(super) mod provider_scan_compute; pub mod purge; pub mod query_collection_size; pub mod reclaim; diff --git a/nodedb/src/data/executor/handlers/provider_scan.rs b/nodedb/src/data/executor/handlers/provider_scan.rs index b48f92e9d..799f3acf1 100644 --- a/nodedb/src/data/executor/handlers/provider_scan.rs +++ b/nodedb/src/data/executor/handlers/provider_scan.rs @@ -2,15 +2,20 @@ //! Executor handler for `QueryOp::ProviderScan`. //! -//! Decodes the pre-materialized msgpack row array, applies predicate filtering, -//! offset, sort, distinct deduplication, column projection, and limit — in that -//! order — then emits the resulting rows via `response_with_payload`. +//! Decodes the pre-materialized msgpack row array, applies predicate +//! filtering, window functions + computed columns, sort, distinct +//! deduplication, offset, column projection, and limit — in that order — +//! then emits the resulting rows via `response_with_payload`. Windows and +//! computed columns run before sort so `ORDER BY` can name their aliases. +//! Sort runs before offset so `ORDER BY ... OFFSET n` skips the first `n` +//! rows of the sorted set, not the decoded set. use nodedb_query::msgpack_scan; use crate::bridge::envelope::{ErrorCode, Response}; use crate::bridge::scan_filter::ScanFilter; use crate::data::executor::core_loop::CoreLoop; +use crate::data::executor::handlers::provider_scan_compute::apply_windows_and_computed; use crate::data::executor::handlers::sort_utils::sort_msgpack_rows; use crate::data::executor::msgpack_utils::write_str; use crate::data::executor::response_codec::encode_binary_rows; @@ -21,6 +26,8 @@ pub(in crate::data::executor) struct ProviderScanParams<'a> { pub rows_bytes: &'a [u8], pub filters_bytes: &'a [u8], pub projection: &'a [String], + pub computed_columns_bytes: &'a [u8], + pub window_functions_bytes: &'a [u8], pub sort_keys: &'a [nodedb_physical::physical_plan::SortKeySpec], pub limit: Option, pub offset: usize, @@ -30,8 +37,8 @@ pub(in crate::data::executor) struct ProviderScanParams<'a> { impl CoreLoop { /// Execute a `ProviderScan` plan node. /// - /// Processing order: decode rows → filter → offset → sort → distinct → - /// project → limit → emit. + /// Processing order: decode rows → filter → windows + computed columns → + /// sort → distinct → offset → project → limit → emit. pub(in crate::data::executor) fn execute_provider_scan( &mut self, task: &ExecutionTask, @@ -41,6 +48,8 @@ impl CoreLoop { rows_bytes, filters_bytes, projection, + computed_columns_bytes, + window_functions_bytes, sort_keys, limit, offset, @@ -92,16 +101,25 @@ impl CoreLoop { } } - // ── 3. Offset. ──────────────────────────────────────────────────────── - if offset > 0 { - if offset >= rows.len() { - rows.clear(); - } else { - rows.drain(..offset); - } + // ── 3. Window functions + computed columns. ───────────────────────── + // Runs before sort so `ORDER BY` can name a window or computed alias. + // Each window spec orders its own partitions, so the row order here + // does not affect window results. Skipped entirely (zero-decode + // msgpack path) when both byte slices are empty. + if !window_functions_bytes.is_empty() || !computed_columns_bytes.is_empty() { + rows = match apply_windows_and_computed( + rows, + window_functions_bytes, + computed_columns_bytes, + ) { + Ok(r) => r, + Err(e) => return self.response_error(task, e), + }; } // ── 4. Sort. ────────────────────────────────────────────────────────── + // Runs before offset: `ORDER BY ... OFFSET n` skips the first `n` rows + // of the SORTED set, not the decoded set. if !sort_keys.is_empty() && let Err(e) = sort_msgpack_rows(&mut rows, sort_keys) { @@ -124,7 +142,16 @@ impl CoreLoop { }); } - // ── 6. Project. ─────────────────────────────────────────────────────── + // ── 6. Offset. ──────────────────────────────────────────────────────── + if offset > 0 { + if offset >= rows.len() { + rows.clear(); + } else { + rows.drain(..offset); + } + } + + // ── 7. Project. ─────────────────────────────────────────────────────── let rows: Vec> = if projection.is_empty() { rows } else { @@ -133,14 +160,14 @@ impl CoreLoop { .collect() }; - // ── 7. Limit. ───────────────────────────────────────────────────────── + // ── 8. Limit. ───────────────────────────────────────────────────────── let rows = if let Some(n) = limit { rows.into_iter().take(n).collect() } else { rows }; - // ── 8. Emit. ────────────────────────────────────────────────────────── + // ── 9. Emit. ────────────────────────────────────────────────────────── let payload = encode_binary_rows(&rows); self.response_with_payload(task, payload) } diff --git a/nodedb/src/data/executor/handlers/provider_scan_compute.rs b/nodedb/src/data/executor/handlers/provider_scan_compute.rs new file mode 100644 index 000000000..5dc27c063 --- /dev/null +++ b/nodedb/src/data/executor/handlers/provider_scan_compute.rs @@ -0,0 +1,100 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Window-function and computed-column evaluation for `QueryOp::ProviderScan`. +//! +//! Runs after filter and before sort/distinct/offset/project/limit in the +//! `ProviderScan` pipeline. Each msgpack row decodes to a `serde_json::Value`, +//! windows evaluate over the full row set, computed columns evaluate +//! per-row, and the result re-encodes to msgpack. Skipped entirely when both +//! byte slices are empty, so the zero-decode msgpack path stays untouched for +//! a plain relational scan. + +use crate::bridge::envelope::ErrorCode; +use crate::bridge::expr_eval::ComputedColumn; +use crate::bridge::window_func::{WindowFuncSpec, evaluate_window_functions}; + +/// Decode a `Vec` from MessagePack, tagging decode failures with which +/// byte slice (`kind`, e.g. `"window"` or `"computed"`) failed. +fn decode_bytes<'a, T: zerompk::FromMessagePack<'a>>( + bytes: &'a [u8], + kind: &str, +) -> crate::Result { + zerompk::from_msgpack(bytes).map_err(|e| { + crate::Error::DataPlane(ErrorCode::Internal { + detail: format!("ProviderScan: malformed {kind} bytes: {e}"), + }) + }) +} + +/// Apply window functions then computed columns to `rows`, both optional and +/// independently controlled by `window_bytes` / `computed_bytes` being +/// non-empty. Returns `rows` unchanged, still msgpack-encoded, when both are +/// empty. +pub(in crate::data::executor) fn apply_windows_and_computed( + rows: Vec>, + window_bytes: &[u8], + computed_bytes: &[u8], +) -> crate::Result>> { + if window_bytes.is_empty() && computed_bytes.is_empty() { + return Ok(rows); + } + + let window_specs: Vec = if window_bytes.is_empty() { + Vec::new() + } else { + decode_bytes(window_bytes, "window")? + }; + let computed_cols: Vec = if computed_bytes.is_empty() { + Vec::new() + } else { + decode_bytes(computed_bytes, "computed")? + }; + + let mut json_rows: Vec<(String, serde_json::Value)> = Vec::with_capacity(rows.len()); + for (idx, row) in rows.iter().enumerate() { + let value = nodedb_types::value_from_msgpack(row).map_err(|e| { + crate::Error::DataPlane(ErrorCode::Internal { + detail: format!("ProviderScan: malformed row for window/computed evaluation: {e}"), + }) + })?; + json_rows.push((idx.to_string(), serde_json::Value::from(value))); + } + + if !window_specs.is_empty() { + evaluate_window_functions(&mut json_rows, &window_specs).map_err(crate::Error::from)?; + } + + for (_, row_json) in &mut json_rows { + if computed_cols.is_empty() { + continue; + } + // Every computed column evaluates against the row as it stood before + // this loop, matching `apply_projection`'s semantics: later computed + // columns never observe earlier ones' results. + let doc_val = nodedb_types::Value::from(row_json.clone()); + for cc in &computed_cols { + let already_present = matches!(row_json.get(&cc.alias), Some(v) if !v.is_null()); + if already_present { + continue; + } + let v = cc.expr.eval(&doc_val)?; + if let serde_json::Value::Object(obj) = row_json { + obj.insert(cc.alias.clone(), serde_json::Value::from(v)); + } + } + } + + let mut out = Vec::with_capacity(json_rows.len()); + for (_, row_json) in json_rows { + let bytes = nodedb_types::json_to_msgpack(&row_json).map_err(|e| { + crate::Error::DataPlane(ErrorCode::Internal { + detail: format!( + "ProviderScan: failed to re-encode row after window/computed evaluation: {e}" + ), + }) + })?; + out.push(bytes); + } + + Ok(out) +} diff --git a/nodedb/tests/inproc/cases/executor_tests/test_cross_type_join/inline_hash_join.rs b/nodedb/tests/inproc/cases/executor_tests/test_cross_type_join/inline_hash_join.rs index 8ce8e1420..5f30bf18d 100644 --- a/nodedb/tests/inproc/cases/executor_tests/test_cross_type_join/inline_hash_join.rs +++ b/nodedb/tests/inproc/cases/executor_tests/test_cross_type_join/inline_hash_join.rs @@ -171,6 +171,8 @@ fn inline_hash_join_honors_qualified_left_keys() { rows: response_codec::flatten_to_relational_rows(&left_data), filters: Vec::new(), projection: Vec::new(), + computed_columns: Vec::new(), + window_functions: Vec::new(), sort_keys: Vec::new(), limit: None, offset: 0, @@ -181,6 +183,8 @@ fn inline_hash_join_honors_qualified_left_keys() { rows: response_codec::flatten_to_relational_rows(&right_data), filters: Vec::new(), projection: Vec::new(), + computed_columns: Vec::new(), + window_functions: Vec::new(), sort_keys: Vec::new(), limit: None, offset: 0, diff --git a/nodedb/tests/inproc/cases/executor_tests/test_cross_type_join/multi_core_joins.rs b/nodedb/tests/inproc/cases/executor_tests/test_cross_type_join/multi_core_joins.rs index 0452152fc..321116d8b 100644 --- a/nodedb/tests/inproc/cases/executor_tests/test_cross_type_join/multi_core_joins.rs +++ b/nodedb/tests/inproc/cases/executor_tests/test_cross_type_join/multi_core_joins.rs @@ -126,6 +126,8 @@ fn multi_core_broadcast_inner_join() { rows: response_codec::flatten_to_relational_rows(&phase1_payload), filters: Vec::new(), projection: Vec::new(), + computed_columns: Vec::new(), + window_functions: Vec::new(), sort_keys: Vec::new(), limit: None, offset: 0, @@ -268,6 +270,8 @@ fn multi_core_broadcast_left_join() { rows: response_codec::flatten_to_relational_rows(&phase1_payload), filters: Vec::new(), projection: Vec::new(), + computed_columns: Vec::new(), + window_functions: Vec::new(), sort_keys: Vec::new(), limit: None, offset: 0, @@ -450,6 +454,8 @@ fn multi_core_broadcast_merge_simulation() { rows: response_codec::flatten_to_relational_rows(&data), filters: Vec::new(), projection: Vec::new(), + computed_columns: Vec::new(), + window_functions: Vec::new(), sort_keys: Vec::new(), limit: None, offset: 0, diff --git a/nodedb/tests/wire/cases/pgwire_extended_query.rs b/nodedb/tests/wire/cases/pgwire_extended_query.rs index 612ee388c..1c309641e 100644 --- a/nodedb/tests/wire/cases/pgwire_extended_query.rs +++ b/nodedb/tests/wire/cases/pgwire_extended_query.rs @@ -105,10 +105,10 @@ async fn extended_query_constant_and_param_projection() { rows[0].len() ); - // x may decode as any integer-compatible type; compare via text. - let x_text: String = rows[0].get::<_, String>("x"); + // A constant integer is typed as int8 in the row description. + let x: i64 = rows[0].get("x"); let y: &str = rows[0].get("y"); - assert_eq!(x_text, "1"); + assert_eq!(x, 1); assert_eq!(y, "hi"); } @@ -136,9 +136,9 @@ async fn extended_query_pure_constant_projection() { rows[0].len() ); - let x_text: String = rows[0].get::<_, String>("x"); + let x: i64 = rows[0].get("x"); let y: &str = rows[0].get("y"); - assert_eq!(x_text, "1"); + assert_eq!(x, 1); assert_eq!(y, "hi"); } diff --git a/nodedb/tests/wire/cases/pgwire_extended_query_engines2.rs b/nodedb/tests/wire/cases/pgwire_extended_query_engines2.rs index 38537449e..3f03ae3c7 100644 --- a/nodedb/tests/wire/cases/pgwire_extended_query_engines2.rs +++ b/nodedb/tests/wire/cases/pgwire_extended_query_engines2.rs @@ -228,8 +228,8 @@ async fn extended_query_array_engine_smoke_and_const_stmt() { .await .expect("constant execute after Array DDL"); assert_eq!(const_rows.len(), 1, "constant projection must return 1 row"); - let x_text: String = const_rows[0].get::<_, String>(0); - assert_eq!(x_text, "1"); + let x: i64 = const_rows[0].get(0); + assert_eq!(x, 1); } // ── Cross-engine: parameter error cases ────────────────────────────────────── diff --git a/nodedb/tests/wire/cases/sql_subquery_from.rs b/nodedb/tests/wire/cases/sql_subquery_from.rs index 4b516bd4c..ebd420987 100644 --- a/nodedb/tests/wire/cases/sql_subquery_from.rs +++ b/nodedb/tests/wire/cases/sql_subquery_from.rs @@ -101,3 +101,281 @@ async fn derived_group_by_in_from_is_supported() { assert_eq!(totals.get("b"), Some(&7.0), "category b total should be 7"); assert_eq!(totals.get("c"), Some(&5.0), "category c total should be 5"); } + +/// A computed column over a constant derived table must evaluate. The +/// constant body lowers to a provider row, and the outer projection must +/// run over that row instead of resolving `x * 2` by name to NULL. +#[tokio::test] +async fn computed_column_over_constant_derived_table_evaluates() { + let srv = TestServer::start().await; + + let rows = srv + .query_rows("SELECT x * 2 AS doubled FROM (SELECT 1 AS x) AS s") + .await + .expect("computed column over a constant derived table must plan"); + + assert_eq!(rows, vec![vec!["2".to_string()]], "got {rows:?}"); +} + +/// A computed column over a derived table whose body is an aggregate must +/// evaluate over the aggregate's output row. +#[tokio::test] +async fn computed_column_over_aggregate_derived_table_evaluates() { + let srv = TestServer::start().await; + create_items(&srv).await; + + let rows = srv + .query_rows("SELECT total * 2 AS doubled FROM (SELECT SUM(qty) AS total FROM items) AS s") + .await + .expect("computed column over an aggregate derived table must plan"); + + assert_eq!(rows, vec![vec!["30".to_string()]], "got {rows:?}"); +} + +/// A computed column over a grouped derived table must evaluate per output +/// row and keep every projected column. +#[tokio::test] +async fn computed_column_over_grouped_derived_table_evaluates() { + let srv = TestServer::start().await; + create_items(&srv).await; + + let rows = srv + .query_rows( + "SELECT category, total * 2 AS doubled \ + FROM (SELECT category, SUM(qty) AS total FROM items GROUP BY category) AS agg", + ) + .await + .expect("computed column over a grouped derived table must plan"); + + let mut got: Vec<(String, String)> = + rows.iter().map(|r| (r[0].clone(), r[1].clone())).collect(); + got.sort(); + assert_eq!( + got, + vec![ + ("a".to_string(), "6".to_string()), + ("b".to_string(), "14".to_string()), + ("c".to_string(), "10".to_string()), + ], + "got {rows:?}" + ); +} + +/// A computed column over a UNION ALL derived table must evaluate per row. +#[tokio::test] +async fn computed_column_over_union_derived_table_evaluates() { + let srv = TestServer::start().await; + + let rows = srv + .query_rows("SELECT x * 2 AS doubled FROM (SELECT 1 AS x UNION ALL SELECT 2 AS x) AS s") + .await + .expect("computed column over a UNION ALL derived table must plan"); + + let mut got: Vec = rows.iter().map(|r| r[0].clone()).collect(); + got.sort(); + assert_eq!(got, vec!["2".to_string(), "4".to_string()], "got {rows:?}"); +} + +/// Division by zero in the projection over a constant derived table must +/// raise `22012`, never fold to a NULL row. +#[tokio::test] +async fn projection_division_by_zero_over_constant_derived_table_errors_22012() { + let srv = TestServer::start().await; + + srv.expect_error("SELECT x / 0 FROM (SELECT 1 AS x) AS s", "22012") + .await; +} + +/// Division by zero in an aggregate argument over a constant derived table +/// must raise `22012`, never fold to a NULL aggregate. +#[tokio::test] +async fn aggregate_argument_division_by_zero_over_constant_derived_table_errors_22012() { + let srv = TestServer::start().await; + + srv.expect_error("SELECT SUM(x / 0) FROM (SELECT 1 AS x) AS s", "22012") + .await; +} + +/// Division by zero in a GROUP BY key over a constant derived table must +/// raise `22012`, never return an empty result with a missing column. +#[tokio::test] +async fn group_by_key_division_by_zero_over_constant_derived_table_errors_22012() { + let srv = TestServer::start().await; + + srv.expect_error( + "SELECT x, COUNT(*) FROM (SELECT 1 AS x) AS s GROUP BY x / 0", + "22012", + ) + .await; +} + +/// Division by zero in a window PARTITION BY over a constant derived table +/// must raise `22012`, never fold to a NULL window value. +#[tokio::test] +async fn window_partition_division_by_zero_over_constant_derived_table_errors_22012() { + let srv = TestServer::start().await; + + srv.expect_error( + "SELECT SUM(x) OVER (PARTITION BY x / 0) FROM (SELECT 1 AS x) AS s", + "22012", + ) + .await; +} + +/// Division by zero in the projection over an aggregate derived table must +/// raise `22012`. +#[tokio::test] +async fn projection_division_by_zero_over_aggregate_derived_table_errors_22012() { + let srv = TestServer::start().await; + create_items(&srv).await; + + srv.expect_error( + "SELECT total / 0 FROM (SELECT SUM(qty) AS total FROM items) AS s", + "22012", + ) + .await; +} + +/// A grouped query over a constant derived table must keep every projected +/// column in the result, not drop the non-aggregate column. +#[tokio::test] +async fn group_by_over_constant_derived_table_keeps_projected_columns() { + let srv = TestServer::start().await; + + let rows = srv + .query_rows("SELECT x, COUNT(*) AS n FROM (SELECT 1 AS x) AS s GROUP BY x") + .await + .expect("GROUP BY over a constant derived table must plan"); + + assert_eq!( + rows, + vec![vec!["1".to_string(), "1".to_string()]], + "got {rows:?}" + ); +} + +/// Control: the same projection over a derived table that scans a +/// collection raises `22012`. The derived table itself is not the trigger. +#[tokio::test] +async fn projection_division_by_zero_over_scan_derived_table_errors_22012() { + let srv = TestServer::start().await; + create_items(&srv).await; + + srv.expect_error( + "SELECT x / 0 FROM (SELECT qty AS x FROM items) AS s", + "22012", + ) + .await; +} + +/// A computed column that references an inner alias (`qty AS x`) over a +/// scanning derived table must resolve through the alias. Merging the outer +/// projection onto the inner scan must not discard the inner rename. +#[tokio::test] +async fn computed_column_over_aliased_scan_derived_table_evaluates() { + let srv = TestServer::start().await; + create_items(&srv).await; + + let rows = srv + .query_rows( + "SELECT x * 2 AS doubled FROM (SELECT qty AS x FROM items WHERE id = 'i3') AS s", + ) + .await + .expect("computed column over an aliased scan derived table must plan"); + + assert_eq!(rows, vec![vec!["6".to_string()]], "got {rows:?}"); +} + +/// A computed column over a grouped derived table with an outer ORDER BY +/// (which routes through the row post-processor) must evaluate per row. +#[tokio::test] +async fn computed_column_over_grouped_derived_table_with_order_by_evaluates() { + let srv = TestServer::start().await; + create_items(&srv).await; + + let rows = srv + .query_rows( + "SELECT category, total * 2 AS doubled \ + FROM (SELECT category, SUM(qty) AS total FROM items GROUP BY category) AS agg \ + ORDER BY category", + ) + .await + .expect("computed column over a grouped derived table with ORDER BY must plan"); + + assert_eq!( + rows, + vec![ + vec!["a".to_string(), "6".to_string()], + vec!["b".to_string(), "14".to_string()], + vec!["c".to_string(), "10".to_string()], + ], + "got {rows:?}" + ); +} + +/// An aggregate over a grouped derived table (aggregate of aggregates) must +/// run over the inner group rows, not over an empty collection. +#[tokio::test] +async fn aggregate_over_grouped_derived_table_evaluates() { + let srv = TestServer::start().await; + create_items(&srv).await; + + let rows = srv + .query_rows( + "SELECT SUM(total) AS grand, COUNT(*) AS groups \ + FROM (SELECT category, SUM(qty) AS total FROM items GROUP BY category) AS agg", + ) + .await + .expect("aggregate over a grouped derived table must plan"); + + // SUM renders as a float text today; compare numerically. + assert_eq!(rows.len(), 1, "got {rows:?}"); + let grand: f64 = rows[0][0].parse().expect("grand total must be numeric"); + assert_eq!(grand, 15.0, "got {rows:?}"); + assert_eq!(rows[0][1], "3", "got {rows:?}"); +} + +/// An aggregate over a UNION ALL derived table must run over the union rows. +#[tokio::test] +async fn aggregate_over_union_derived_table_evaluates() { + let srv = TestServer::start().await; + + let rows = srv + .query_rows("SELECT SUM(x) AS total FROM (SELECT 1 AS x UNION ALL SELECT 2 AS x) AS s") + .await + .expect("aggregate over a UNION ALL derived table must plan"); + + // SUM renders as a float text today; compare numerically. + assert_eq!(rows.len(), 1, "got {rows:?}"); + let total: f64 = rows[0][0].parse().expect("total must be numeric"); + assert_eq!(total, 3.0, "got {rows:?}"); +} + +/// A window function over a grouped derived table must rank the inner group +/// rows. +#[tokio::test] +async fn window_over_grouped_derived_table_evaluates() { + let srv = TestServer::start().await; + create_items(&srv).await; + + let rows = srv + .query_rows( + "SELECT category, RANK() OVER (ORDER BY total DESC) AS rnk \ + FROM (SELECT category, SUM(qty) AS total FROM items GROUP BY category) AS agg", + ) + .await + .expect("window function over a grouped derived table must plan"); + + let mut got: Vec<(String, String)> = + rows.iter().map(|r| (r[0].clone(), r[1].clone())).collect(); + got.sort(); + assert_eq!( + got, + vec![ + ("a".to_string(), "3".to_string()), + ("b".to_string(), "1".to_string()), + ("c".to_string(), "2".to_string()), + ], + "got {rows:?}" + ); +}