diff --git a/nodedb-cluster-tests/tests/transport_security.rs b/nodedb-cluster-tests/tests/transport_security.rs index 4995455e2..ec64f2774 100644 --- a/nodedb-cluster-tests/tests/transport_security.rs +++ b/nodedb-cluster-tests/tests/transport_security.rs @@ -38,12 +38,12 @@ fn insecure_counter_guard() -> &'static Mutex<()> { LOCK.get_or_init(|| Mutex::new(())) } +use nodedb_cluster::topology::{ClusterTopology, NodeInfo, NodeState}; use nodedb_cluster::transport::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer}; use nodedb_cluster::{ MacKey, NexarTransport, RaftRpcHandler, TlsCredentials, TransportCredentials, generate_node_credentials, insecure_transport_count, spki_pin_from_cert_der, }; -use nodedb_cluster::topology::{ClusterTopology, NodeInfo, NodeState}; use nodedb_raft::message::{AppendEntriesRequest, AppendEntriesResponse, RequestVoteResponse}; use nodedb_raft::transport::RaftTransport; diff --git a/nodedb-physical/src/physical_plan/mod.rs b/nodedb-physical/src/physical_plan/mod.rs index 914bd28f5..5d037f6ff 100644 --- a/nodedb-physical/src/physical_plan/mod.rs +++ b/nodedb-physical/src/physical_plan/mod.rs @@ -324,6 +324,8 @@ impl PhysicalPlan { | PhysicalPlan::Graph(GraphOp::WccSuperstep(_)) => None, // Exchange: recurse into the child plan to extract the collection. PhysicalPlan::Query(QueryOp::Exchange(op)) => op.child.collection(), + // PostProcess: recurse into the materialized input plan. + PhysicalPlan::Query(QueryOp::PostProcess { input, .. }) => input.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 → None). diff --git a/nodedb-physical/src/physical_plan/query.rs b/nodedb-physical/src/physical_plan/query.rs index ae8422cfa..ffe42a9ec 100644 --- a/nodedb-physical/src/physical_plan/query.rs +++ b/nodedb-physical/src/physical_plan/query.rs @@ -92,6 +92,46 @@ pub enum QueryOp { distinct: bool, }, + /// Relational post-processing over a materialized child plan. + /// + /// Coordinator-resolved, exactly like [`QueryOp::Exchange`]: at resolve time + /// the coordinator gathers `input`'s rows, flattens them to the relational + /// row shape, and rewrites this node into a [`QueryOp::ProviderScan`] + /// carrying the same relational parameters — which applies + /// filter→offset→sort→distinct→project→limit on a single core. A Data-Plane + /// core must NEVER see this node (defensive error if it does). + /// + /// Lowered from `SqlPlan::Subquery`: an outer `ORDER BY` / `OFFSET` / + /// `DISTINCT` / post-reorder `LIMIT` over a subquery/derived-table body + /// whose leaf plan could not absorb those constraints. Constraints the leaf + /// DID absorb (a `WHERE` pushed into a search engine, an unordered `LIMIT` + /// folded into `top_k`) are applied by `input` and not repeated here. + PostProcess { + /// The subquery body to materialize. Wrapped in `Exchange{Gather}` by + /// the converter when the body is a sharded source, so the gather runs + /// exactly once over the full union before post-processing. + input: Box, + /// Serialized `Vec` (MessagePack). Empty = no predicate. + #[serde(default)] + filters: Vec, + /// Output column names to keep. Empty = emit all columns. + #[serde(default)] + projection: Vec, + /// Sort keys: `(column_name, ascending)`. Empty = unordered. + #[serde(default)] + sort_keys: Vec<(String, bool)>, + /// Maximum rows to emit after offset/sort/distinct/projection. + /// `None` = unlimited. + #[serde(default)] + limit: Option, + /// Number of rows to skip before applying limit. + #[serde(default)] + offset: usize, + /// SQL DISTINCT: deduplicate on the projected row. + #[serde(default)] + distinct: bool, + }, + /// Aggregate: GROUP BY + aggregate functions. Aggregate { collection: String, diff --git a/nodedb-physical/src/physical_plan/routing.rs b/nodedb-physical/src/physical_plan/routing.rs index 57bfd9fb9..d9376e7ca 100644 --- a/nodedb-physical/src/physical_plan/routing.rs +++ b/nodedb-physical/src/physical_plan/routing.rs @@ -75,6 +75,13 @@ pub fn plan_contains_cluster_partitioned_leaf(plan: &PhysicalPlan) -> bool { plan_contains_cluster_partitioned_leaf(child) } + // Recurse through PostProcess (its materialized input is the real + // plan). Normally resolved to a `ProviderScan` before routing, but a + // conservative recursion keeps routing correct if one survives. + PhysicalPlan::Query(QueryOp::PostProcess { input, .. }) => { + plan_contains_cluster_partitioned_leaf(input) + } + // Recurse through lateral outer plans. PhysicalPlan::Query(QueryOp::LateralTopK { outer_plan, .. }) | PhysicalPlan::Query(QueryOp::LateralLoop { outer_plan, .. }) => { diff --git a/nodedb-sql/src/parser/preprocess/pipeline.rs b/nodedb-sql/src/parser/preprocess/pipeline.rs index d368407dc..f2354fdcc 100644 --- a/nodedb-sql/src/parser/preprocess/pipeline.rs +++ b/nodedb-sql/src/parser/preprocess/pipeline.rs @@ -13,7 +13,9 @@ use super::function_args::rewrite_object_literal_args; use super::lex::{first_sql_word, has_brace_outside_literals, has_operator_outside_literals}; use super::object_literal_stmt::try_rewrite_object_literal; -use super::search_vector::try_rewrite_search_using_vector; +use super::search_vector::{ + try_rewrite_nested_search_using_vector, try_rewrite_search_using_vector, +}; use super::temporal::extract as extract_temporal; use super::vector_ops::{ rewrite_arrow_distance, rewrite_cosine_distance, rewrite_neg_inner_product, @@ -56,10 +58,16 @@ pub fn preprocess(sql: &str) -> Result, SqlError> { // `SELECT * FROM ORDER BY vector_distance(...) LIMIT k` before any // first-word dispatch — the rewritten form re-enters the rest of the // pipeline as a plain SELECT. + // + // The same clause in subquery position — `FROM (SEARCH ...) s` — is + // rewritten in place so a k-NN result composes with joins and filters. let (temporal_sql, search_vector_rewritten) = match try_rewrite_search_using_vector(&temporal_sql) { Some(rewritten) => (rewritten, true), - None => (temporal_sql, false), + None => match try_rewrite_nested_search_using_vector(&temporal_sql) { + Some(rewritten) => (rewritten, true), + None => (temporal_sql, false), + }, }; let first_word = first_sql_word(&temporal_sql) diff --git a/nodedb-sql/src/parser/preprocess/search_vector.rs b/nodedb-sql/src/parser/preprocess/search_vector.rs index 351278fdd..5bffc2bae 100644 --- a/nodedb-sql/src/parser/preprocess/search_vector.rs +++ b/nodedb-sql/src/parser/preprocess/search_vector.rs @@ -6,6 +6,12 @@ //! `` may be omitted (and the third arg becomes the limit). When the //! collection has a single declared vector column the planner resolves the //! field; otherwise `vector_distance` rejects the call with a typed error. +//! +//! The same clause is also accepted in subquery position — `FROM (SEARCH ...) +//! s`, `IN (SELECT id FROM (SEARCH ...) s)` — so a k-NN result composes with +//! joins and relational filters inside one statement. + +use super::lex::{find_ascii_case_insensitive, find_operator_positions}; const SEARCH_KEYWORD: &str = "SEARCH"; const USING_KEYWORD: &str = "USING"; @@ -13,43 +19,100 @@ const VECTOR_KEYWORD: &str = "VECTOR"; pub fn try_rewrite_search_using_vector(sql: &str) -> Option { let trimmed = sql.trim_end_matches(|c: char| c == ';' || c.is_whitespace()); - let upper = trimmed.to_uppercase(); - let stripped = upper.trim_start(); - if !stripped.starts_with(SEARCH_KEYWORD) { - return None; - } let leading = trimmed.len() - trimmed.trim_start().len(); - let after_search = trimmed[leading + SEARCH_KEYWORD.len()..].trim_start(); - if after_search.is_empty() { + let (select, consumed) = parse_search_clause(&trimmed[leading..])?; + // A top-level `SEARCH` clause owns the whole statement; anything past the + // closing paren is not part of the DSL form and must not be rewritten. + if !trimmed[leading + consumed..].trim().is_empty() { return None; } - let (collection, rest) = take_identifier(after_search)?; - let rest = rest.trim_start(); - let rest_upper = rest.to_uppercase(); - if !rest_upper.starts_with(USING_KEYWORD) { - return None; + let trailing = &sql[trimmed.len()..]; + Some(format!("{select}{trailing}")) +} + +/// Rewrite every `(SEARCH USING VECTOR(...))` occurrence that sits in +/// subquery position into `(SELECT ...)`, leaving the rest of `sql` untouched. +/// +/// Returns `None` when no occurrence was rewritten. +pub fn try_rewrite_nested_search_using_vector(sql: &str) -> Option { + // Cheap gate: every other pass in the pipeline fronts its scan with a + // substring check, and this one runs on every statement that is not itself + // a top-level SEARCH. + find_ascii_case_insensitive(sql, SEARCH_KEYWORD)?; + + let mut out = String::new(); + let mut copied = 0usize; + for open in find_operator_positions(sql, "(") { + if open < copied { + continue; + } + let after_open = open + 1; + let inner = &sql[after_open..]; + let leading = inner.len() - inner.trim_start().len(); + let Some((select, consumed)) = parse_search_clause(&inner[leading..]) else { + continue; + }; + out.push_str(&sql[copied..after_open]); + out.push_str(&select); + copied = after_open + leading + consumed; } - let after_using = rest[USING_KEYWORD.len()..].trim_start(); - let after_using_upper = after_using.to_uppercase(); - if !after_using_upper.starts_with(VECTOR_KEYWORD) { + if copied == 0 { return None; } - let after_vec = after_using[VECTOR_KEYWORD.len()..].trim_start(); - let body = strip_parentheses(after_vec)?; + out.push_str(&sql[copied..]); + Some(out) +} + +/// Parse a `SEARCH USING VECTOR(...)` clause at the start of `input`. +/// +/// Returns the canonical `SELECT` and the number of bytes consumed (through +/// the clause's closing paren). +fn parse_search_clause(input: &str) -> Option<(String, usize)> { + let after_search = keyword_rest(input, SEARCH_KEYWORD)?; + let (collection, rest) = take_identifier(after_search.trim_start())?; + let after_using = keyword_rest(rest, USING_KEYWORD)?; + let after_vector = keyword_rest(after_using, VECTOR_KEYWORD)?; + let (body, consumed_args) = take_parenthesized(after_vector)?; let (field, vector_expr, limit) = split_vector_args(body)?; - let trailing = sql[leading + (trimmed.len() - leading)..].to_string(); let order_by = match field { Some(name) => format!("vector_distance({name}, {vector_expr})"), None => format!("vector_distance({vector_expr})"), }; - Some(format!( - "SELECT * FROM {collection} ORDER BY {order_by} LIMIT {limit}{trailing}" - )) + let select = format!("SELECT * FROM {collection} ORDER BY {order_by} LIMIT {limit}"); + let consumed = input.len() - after_vector.len() + consumed_args; + Some((select, consumed)) } +/// Match `keyword` at the start of `input` (case-insensitive, whole word) and +/// return the remainder with leading whitespace stripped. +fn keyword_rest<'a>(input: &'a str, keyword: &str) -> Option<&'a str> { + let trimmed = input.trim_start(); + // Compare bytes: slicing `trimmed` at `keyword.len()` would panic when the + // input starts with a multi-byte character. An ASCII match guarantees the + // offset is a char boundary, so the slices below are safe. + let bytes = trimmed.as_bytes(); + if bytes.len() < keyword.len() + || !bytes[..keyword.len()].eq_ignore_ascii_case(keyword.as_bytes()) + { + return None; + } + let rest = &trimmed[keyword.len()..]; + if rest.chars().next().is_some_and(is_ident_char) { + return None; + } + Some(rest.trim_start()) +} + +/// Take a collection name: a bare identifier, or a double-quoted one. The +/// quoted form is returned with its quotes intact so the rewritten `SELECT` +/// keeps the exact spelling the user wrote — a name that needs quoting still +/// needs it after the rewrite. fn take_identifier(input: &str) -> Option<(&str, &str)> { + if input.starts_with('"') { + return take_quoted_identifier(input); + } let end = input .char_indices() .find(|(_, c)| !is_ident_char(*c)) @@ -61,17 +124,58 @@ fn take_identifier(input: &str) -> Option<(&str, &str)> { Some((&input[..end], &input[end..])) } +/// Take a `"quoted name"`, honouring the doubled-quote escape (`""`). Returns +/// the quoted span verbatim and the remainder. +fn take_quoted_identifier(input: &str) -> Option<(&str, &str)> { + let mut chars = input.char_indices().skip(1); + while let Some((offset, c)) = chars.next() { + if c != '"' { + continue; + } + // A doubled quote is an escaped quote, not the end of the identifier. + if input[offset + 1..].starts_with('"') { + chars.next(); + continue; + } + let end = offset + 1; + // An empty name ("") is not a usable collection reference. + if end == 2 { + return None; + } + return Some((&input[..end], &input[end..])); + } + None +} + fn is_ident_char(c: char) -> bool { c.is_ascii_alphanumeric() || c == '_' } -fn strip_parentheses(input: &str) -> Option<&str> { - let trimmed = input.trim(); - let bytes = trimmed.as_bytes(); - if bytes.first() != Some(&b'(') || bytes.last() != Some(&b')') { +/// Take a parenthesized group at the start of `input`, returning its body and +/// the number of bytes consumed (through the matching close paren). Parens +/// inside string literals do not affect nesting depth. +fn take_parenthesized(input: &str) -> Option<(&str, usize)> { + if !input.starts_with('(') { return None; } - Some(trimmed[1..trimmed.len() - 1].trim()) + let mut depth = 0i32; + let mut in_single = false; + let mut in_double = false; + for (offset, c) in input.char_indices() { + match c { + '\'' if !in_double => in_single = !in_single, + '"' if !in_single => in_double = !in_double, + '(' if !in_single && !in_double => depth += 1, + ')' if !in_single && !in_double => { + depth -= 1; + if depth == 0 { + return Some((input[1..offset].trim(), offset + 1)); + } + } + _ => {} + } + } + None } fn split_vector_args(body: &str) -> Option<(Option, String, String)> { @@ -157,6 +261,135 @@ mod tests { assert!(try_rewrite_search_using_vector("SEARCH c USING FUSION(ARRAY[0.5])").is_none()); } + #[test] + fn rewrites_search_in_derived_table_position() { + let out = try_rewrite_nested_search_using_vector( + "SELECT * FROM (SEARCH docs USING VECTOR(emb, ARRAY[0.1], 2)) s WHERE s.id = 'a1'", + ) + .unwrap(); + assert_eq!( + out, + "SELECT * FROM (SELECT * FROM docs ORDER BY vector_distance(emb, ARRAY[0.1]) LIMIT 2) s WHERE s.id = 'a1'" + ); + } + + #[test] + fn rewrites_every_nested_occurrence() { + let out = try_rewrite_nested_search_using_vector( + "SELECT a.id FROM (SEARCH docs USING VECTOR(emb, ARRAY[0.1], 2)) a \ + JOIN (SEARCH docs USING VECTOR(emb, ARRAY[0.9], 3)) b ON a.id = b.id", + ) + .unwrap(); + assert_eq!(out.matches("SELECT * FROM docs ORDER BY").count(), 2); + assert!(!out.to_uppercase().contains("SEARCH")); + } + + #[test] + fn nested_rewrite_ignores_search_inside_string_literal() { + assert!( + try_rewrite_nested_search_using_vector( + "SELECT * FROM docs WHERE title = '(SEARCH docs USING VECTOR(emb, ARRAY[0.1], 2))'", + ) + .is_none() + ); + } + + #[test] + fn rewrites_a_quoted_collection_name() { + let out = try_rewrite_search_using_vector( + "SEARCH \"MixedCase\" USING VECTOR(embedding, ARRAY[0.1, 0.3], 2)", + ) + .unwrap(); + assert_eq!( + out, + "SELECT * FROM \"MixedCase\" ORDER BY vector_distance(embedding, ARRAY[0.1, 0.3]) LIMIT 2", + "the quotes must survive the rewrite, or a name that needs them stops resolving" + ); + } + + #[test] + fn rewrites_a_quoted_collection_name_in_subquery_position() { + let out = try_rewrite_nested_search_using_vector( + "SELECT id FROM (SEARCH \"MixedCase\" USING VECTOR(embedding, ARRAY[0.1], 2)) s", + ) + .unwrap(); + assert_eq!( + out, + "SELECT id FROM (SELECT * FROM \"MixedCase\" ORDER BY vector_distance(embedding, ARRAY[0.1]) LIMIT 2) s" + ); + } + + #[test] + fn quoted_collection_name_keeps_its_escaped_quotes() { + let out = try_rewrite_search_using_vector( + "SEARCH \"od\"\"d\" USING VECTOR(embedding, ARRAY[0.1], 1)", + ) + .unwrap(); + assert!( + out.starts_with("SELECT * FROM \"od\"\"d\" ORDER BY"), + "got: {out}" + ); + } + + #[test] + fn rejects_an_unterminated_or_empty_quoted_collection_name() { + assert!( + try_rewrite_search_using_vector("SEARCH \"unterminated USING VECTOR(e, ARRAY[0.1], 1)") + .is_none() + ); + assert!( + try_rewrite_search_using_vector("SEARCH \"\" USING VECTOR(e, ARRAY[0.1], 1)").is_none() + ); + } + + #[test] + fn nested_rewrite_survives_multibyte_text_after_an_open_paren() { + // The clause scan runs at every `(`. Matching the keyword on bytes + // keeps a multi-byte literal from slicing mid-character. + assert!( + try_rewrite_nested_search_using_vector( + "INSERT INTO t (name) VALUES ('日本語') /* SEARCH */", + ) + .is_none() + ); + assert!( + try_rewrite_nested_search_using_vector( + "SELECT * FROM t WHERE tag IN ('🙂😀', 'SEARCH')" + ) + .is_none() + ); + } + + #[test] + fn nested_rewrite_returns_none_for_plain_subquery() { + assert!( + try_rewrite_nested_search_using_vector("SELECT * FROM (SELECT * FROM docs) s") + .is_none() + ); + } + + #[test] + fn nested_rewrite_returns_none_for_fusion_subquery() { + assert!( + try_rewrite_nested_search_using_vector( + "SELECT * FROM (SEARCH docs USING FUSION(ARRAY[0.5])) s" + ) + .is_none() + ); + } + + #[test] + fn top_level_rewrite_rejects_trailing_clause() { + // Only the DSL form is a statement; trailing SQL means the caller wrote + // something else and must get a parse error, not a silent rewrite. + assert!( + try_rewrite_search_using_vector( + "SEARCH t USING VECTOR(emb, ARRAY[1.0], 3) WHERE x = 1" + ) + .is_none() + ); + } + #[test] fn handles_trailing_semicolon() { let out = diff --git a/nodedb-sql/src/planner/select/entry.rs b/nodedb-sql/src/planner/select/entry.rs index 3ef28fa0a..939be7e04 100644 --- a/nodedb-sql/src/planner/select/entry.rs +++ b/nodedb-sql/src/planner/select/entry.rs @@ -356,6 +356,16 @@ fn apply_limit(mut plan: SqlPlan, limit_clause: &Option) -> Sq } }; + // The LIMIT belongs to the query reading the CTE, not to the CTE body, so + // it lands on the outer plan — without this a derived table like + // `FROM (...) s LIMIT n` comes back unbounded. + if let SqlPlan::Cte { definitions, outer } = plan { + return SqlPlan::Cte { + definitions, + outer: Box::new(apply_limit(*outer, limit_clause)), + }; + } + match plan { SqlPlan::Scan { ref mut limit, diff --git a/nodedb-sql/src/types/plan/cacheability.rs b/nodedb-sql/src/types/plan/cacheability.rs index 723992c08..5f5e1f24c 100644 --- a/nodedb-sql/src/types/plan/cacheability.rs +++ b/nodedb-sql/src/types/plan/cacheability.rs @@ -76,6 +76,7 @@ impl SqlPlan { .fold(outer.cache_eligibility(), |eligibility, (_, plan)| { eligibility.combine(plan.cache_eligibility()) }), + Self::Subquery { input, .. } => input.cache_eligibility(), Self::LateralTopK { outer, .. } => outer.cache_eligibility(), Self::LateralLoop { outer, inner, .. } => { outer.cache_eligibility().combine(inner.cache_eligibility()) diff --git a/nodedb-sql/src/types/plan/variants.rs b/nodedb-sql/src/types/plan/variants.rs index 03d7cf1d1..b3bb32062 100644 --- a/nodedb-sql/src/types/plan/variants.rs +++ b/nodedb-sql/src/types/plan/variants.rs @@ -478,6 +478,40 @@ pub enum SqlPlan { outer: Box, }, + /// Relational post-processing over a subquery/derived-table body whose leaf + /// plan cannot absorb the outer query's constraints. + /// + /// Produced by CTE / derived-table inlining when the referenced body is not + /// a plain `Scan` (which carries its own filters/sort/limit/offset/distinct) + /// and the outer reference adds constraints the body has no slot for — an + /// `ORDER BY`, `OFFSET`, `DISTINCT`, or a `LIMIT` that must apply *after* a + /// reorder. Without this node those constraints were silently dropped, + /// returning the body's unordered/unoffset/undeduplicated rows. + /// + /// The Data Plane materializes `input`'s rows, then applies, in this order: + /// filter → offset → sort → distinct → project → limit — matching + /// `QueryOp::ProviderScan` semantics (which this lowers onto). Constraints + /// the body CAN absorb (e.g. an unordered `LIMIT` folded into a vector + /// search's `top_k`, or a `WHERE` pushed into the engine as a post-filter) + /// are applied at the leaf during inlining and are NOT repeated here. + Subquery { + /// The subquery body whose rows are post-processed. + input: Box, + /// Predicates applied to the materialized rows (outer `WHERE` that the + /// body leaf did not absorb). + filters: Vec, + /// Outer projection (target list). Empty = inherit the body's columns. + projection: Vec, + /// Outer `ORDER BY` keys applied over the materialized rows. + sort_keys: Vec, + /// Outer `OFFSET` (0 = none). + offset: usize, + /// Outer `DISTINCT`. + distinct: bool, + /// Outer `LIMIT` (`None` = unbounded). + limit: Option, + }, + // ── Array (ND sparse) ───────────────────────────────────── /// `CREATE ARRAY DIMS (...) ATTRS (...) TILE_EXTENTS (...)`. /// AST is engine-agnostic — the Origin converter builds the typed diff --git a/nodedb-sql/src/visitor/plan_visitor/args.rs b/nodedb-sql/src/visitor/plan_visitor/args.rs index b1b316e9d..6cf0dce63 100644 --- a/nodedb-sql/src/visitor/plan_visitor/args.rs +++ b/nodedb-sql/src/visitor/plan_visitor/args.rs @@ -32,6 +32,17 @@ pub struct ScanVisitArgs<'a> { pub temporal: &'a TemporalScope, } +/// Parameters for [`super::trait_def::PlanVisitor::subquery`]. +pub struct SubqueryVisitArgs<'a> { + pub input: &'a SqlPlan, + pub filters: &'a [Filter], + pub projection: &'a [Projection], + pub sort_keys: &'a [SortKey], + pub offset: usize, + pub distinct: bool, + pub limit: Option, +} + /// Parameters for [`super::trait_def::PlanVisitor::document_index_lookup`]. pub struct DocumentIndexLookupVisitArgs<'a> { pub collection: &'a str, diff --git a/nodedb-sql/src/visitor/plan_visitor/dispatch_rest.rs b/nodedb-sql/src/visitor/plan_visitor/dispatch_rest.rs index d89a7f1c9..2572ae135 100644 --- a/nodedb-sql/src/visitor/plan_visitor/dispatch_rest.rs +++ b/nodedb-sql/src/visitor/plan_visitor/dispatch_rest.rs @@ -22,6 +22,23 @@ pub(super) fn dispatch_rest( SqlPlan::Intersect { left, right, all } => visitor.intersect(left, right, *all), SqlPlan::Except { left, right, all } => visitor.except(left, right, *all), SqlPlan::Cte { definitions, outer } => visitor.cte(definitions, outer), + SqlPlan::Subquery { + input, + filters, + projection, + sort_keys, + offset, + distinct, + limit, + } => visitor.subquery(super::args::SubqueryVisitArgs { + input, + filters, + projection, + sort_keys, + offset: *offset, + distinct: *distinct, + limit: *limit, + }), SqlPlan::CreateArray { name, dims, diff --git a/nodedb-sql/src/visitor/plan_visitor/trait_def.rs b/nodedb-sql/src/visitor/plan_visitor/trait_def.rs index fb6866776..88427f310 100644 --- a/nodedb-sql/src/visitor/plan_visitor/trait_def.rs +++ b/nodedb-sql/src/visitor/plan_visitor/trait_def.rs @@ -9,8 +9,8 @@ use super::args::{ AggregateVisitArgs, CreateArrayVisitArgs, DocumentIndexLookupVisitArgs, HybridSearchTripleVisitArgs, HybridSearchVisitArgs, InsertVisitArgs, JoinVisitArgs, LateralLoopVisitArgs, LateralTopKVisitArgs, MergeVisitArgs, RecursiveScanVisitArgs, - RecursiveValueVisitArgs, ScanVisitArgs, SpatialScanVisitArgs, TimeseriesScanVisitArgs, - UpdateFromVisitArgs, UpsertVisitArgs, VectorSearchVisitArgs, + RecursiveValueVisitArgs, ScanVisitArgs, SpatialScanVisitArgs, SubqueryVisitArgs, + TimeseriesScanVisitArgs, UpdateFromVisitArgs, UpsertVisitArgs, VectorSearchVisitArgs, }; use crate::fts_types::FtsQuery; use crate::temporal::TemporalScope; @@ -234,6 +234,9 @@ pub trait PlanVisitor { outer: &SqlPlan, ) -> Result; + /// Handle [`SqlPlan::Subquery`]. + fn subquery(&mut self, args: SubqueryVisitArgs<'_>) -> Result; + /// Handle [`SqlPlan::CreateArray`]. fn create_array(&mut self, args: CreateArrayVisitArgs<'_>) -> Result; diff --git a/nodedb/src/control/exec_receiver/support.rs b/nodedb/src/control/exec_receiver/support.rs index 0133810eb..0b4d51767 100644 --- a/nodedb/src/control/exec_receiver/support.rs +++ b/nodedb/src/control/exec_receiver/support.rs @@ -64,6 +64,10 @@ pub(super) fn plan_contains_exchange(plan: &PhysicalPlan) -> bool { .any(|child| plan_contains_exchange(child)), QueryOp::LateralTopK { outer_plan, .. } => plan_contains_exchange(outer_plan), QueryOp::LateralLoop { outer_plan, .. } => plan_contains_exchange(outer_plan), + // PostProcess wraps a materialized child that, before coordinator + // resolution, still carries an `Exchange{Gather}` — recurse so an + // unresolved PostProcess is correctly flagged as Exchange-bearing. + QueryOp::PostProcess { input, .. } => plan_contains_exchange(input), // 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 a32bc8c57..6ad0d4d2b 100644 --- a/nodedb/src/control/gateway/version_set.rs +++ b/nodedb/src/control/gateway/version_set.rs @@ -396,6 +396,12 @@ pub fn touched_collections(plan: &PhysicalPlan) -> Vec { out.extend(child_collections); } + // PostProcess: recurse into the materialized child for the + // collections it reads. + PostProcess { input, .. } => { + out.extend(touched_collections(input)); + } + // ProviderScan is a catalog/constant source — no user collection. ProviderScan { .. } => {} diff --git a/nodedb/src/control/planner/rls_injection.rs b/nodedb/src/control/planner/rls_injection.rs index 1618fcf40..5126f2e8e 100644 --- a/nodedb/src/control/planner/rls_injection.rs +++ b/nodedb/src/control/planner/rls_injection.rs @@ -193,6 +193,15 @@ fn inject_rls_for_plan( inject_rls_for_plan(tenant_id, child, rls_store, auth)?; } + // ── PostProcess: recurse into the materialized child ── + // + // The post-processor wraps a subquery body (a sharded scan under + // `Exchange{Gather}`); without this arm the catch-all swallows it and + // the inner scan never receives its RLS filter — an RLS bypass. + PhysicalPlan::Query(QueryOp::PostProcess { input, .. }) => { + inject_rls_for_plan(tenant_id, input, rls_store, auth)?; + } + // ── LateralTopK / LateralLoop: recurse into outer_plan; also inject // RLS into the inner_filters that the executor applies per outer row ── // diff --git a/nodedb/src/control/planner/sql_plan_convert/expr.rs b/nodedb/src/control/planner/sql_plan_convert/expr.rs index ab7130716..04ebe25a2 100644 --- a/nodedb/src/control/planner/sql_plan_convert/expr.rs +++ b/nodedb/src/control/planner/sql_plan_convert/expr.rs @@ -274,6 +274,12 @@ pub(super) fn convert_sort_keys(keys: &[SortKey]) -> Vec<(String, bool)> { } /// 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. @@ -337,8 +343,74 @@ pub(super) fn inline_cte(plan: &SqlPlan, cte_name: &str, cte_plan: &SqlPlan) -> window_functions: inner_w.clone(), temporal: *inner_t, } - } else { + } 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, + } } } } @@ -421,7 +493,215 @@ pub(super) fn inline_cte(plan: &SqlPlan, cte_name: &str, cte_plan: &SqlPlan) -> limit: *limit, }, + // 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/output_schema.rs b/nodedb/src/control/planner/sql_plan_convert/output_schema.rs index e6433b96c..3a852febf 100644 --- a/nodedb/src/control/planner/sql_plan_convert/output_schema.rs +++ b/nodedb/src/control/planner/sql_plan_convert/output_schema.rs @@ -419,6 +419,21 @@ pub fn build_output_schema( SqlPlan::Cte { outer, .. } => { build_output_schema(std::slice::from_ref(outer.as_ref()), catalog, database_id) } + // A post-processor's projected shape is its outer projection; an empty + // projection (SELECT *) inherits the body's columns. (Synthesized during + // conversion, so this is normally unreached — the schema is derived from + // the pre-inline `Cte` above — but the arm keeps derivation correct if a + // `Subquery` is ever schema-derived directly.) + SqlPlan::Subquery { + input, projection, .. + } => { + if projection.is_empty() { + build_output_schema(std::slice::from_ref(input.as_ref()), catalog, database_id) + } else { + let types = HashMap::new(); + schema_from_projection(projection, &types, &[]) + } + } SqlPlan::LateralTopK { projection, .. } | SqlPlan::LateralLoop { projection, .. } => { // No single source collection spans both outer and inner rows; // default every projected field to `Text` rather than picking 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 9733b0d52..d9d91235d 100644 --- a/nodedb/src/control/planner/sql_plan_convert/set_ops.rs +++ b/nodedb/src/control/planner/sql_plan_convert/set_ops.rs @@ -2,7 +2,7 @@ //! Set operations and miscellaneous plan conversions (UNION, INTERSECT, EXCEPT, CTE, etc.). -use nodedb_sql::types::{Projection, SqlPlan, SqlValue}; +use nodedb_sql::types::{Projection, SortKey, SqlExpr, SqlPlan, SqlValue}; use crate::bridge::envelope::PhysicalPlan; use crate::types::{TenantId, VShardId}; @@ -216,6 +216,125 @@ pub(super) fn convert_cte( convert_one(&resolved, tenant_id, ctx) } +/// Lower `SqlPlan::Subquery` — relational post-processing over a subquery body +/// 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. +pub(super) fn convert_subquery( + args: nodedb_sql::SubqueryVisitArgs<'_>, + tenant_id: TenantId, + ctx: &ConvertContext, +) -> crate::Result> { + let nodedb_sql::SubqueryVisitArgs { + input, + filters, + projection, + 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; + + // 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 }, + })); + } + + Ok(vec![PhysicalTask { + tenant_id, + // Coordinator-local: resolved to a `ProviderScan` over the gathered + // rows (empty collection, like a constant result), dispatched once. + vshard_id: VShardId::from_collection_in_database(ctx.database_id, ""), + database_id: ctx.database_id, + plan: PhysicalPlan::Query(QueryOp::PostProcess { + input: Box::new(child), + filters: super::filter::serialize_filters(filters)?, + projection: lower_subquery_projection(projection)?, + sort_keys: lower_subquery_sort_keys(sort_keys)?, + limit, + offset, + distinct, + }), + post_set_op: PostSetOp::None, + txn_id: None, + }]) +} + +/// Lower outer projection items to plain column names for the relational tail. +/// +/// 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 projection has no slot in the row-post-processing tail — it is +/// projected in an outer SELECT, not here. +fn lower_subquery_projection(projection: &[Projection]) -> crate::Result> { + let mut names = Vec::with_capacity(projection.len()); + for p in projection { + match p { + Projection::Column(qname) => { + names.push(qname.rsplit('.').next().unwrap_or(qname).to_string()); + } + Projection::Star | Projection::QualifiedStar(_) => return Ok(Vec::new()), + Projection::Computed { .. } => { + return Err(crate::Error::PlanError { + detail: "a computed projection over an ORDER BY / OFFSET / DISTINCT subquery \ + is not supported; select the base columns in the subquery and \ + compute them in an outer SELECT" + .into(), + }); + } + } + } + Ok(names) +} + +/// Lower outer ORDER BY keys to `(column, ascending)` pairs. Only column +/// references are representable in the row-post-processing tail; a computed +/// ORDER BY expression must be projected in the subquery first. +fn lower_subquery_sort_keys(keys: &[SortKey]) -> crate::Result> { + keys.iter() + .map(|k| match &k.expr { + SqlExpr::Column { name, .. } => Ok(( + name.rsplit('.').next().unwrap_or(name).to_string(), + k.ascending, + )), + _ => Err(crate::Error::PlanError { + detail: "ORDER BY over a subquery supports only column references; project a \ + computed ORDER BY expression in the subquery first" + .into(), + }), + }) + .collect() +} + #[cfg(test)] mod tests { use super::*; diff --git a/nodedb/src/control/planner/sql_plan_convert/visitor/arms_set_ops.rs b/nodedb/src/control/planner/sql_plan_convert/visitor/arms_set_ops.rs index ec9382933..746c1b4eb 100644 --- a/nodedb/src/control/planner/sql_plan_convert/visitor/arms_set_ops.rs +++ b/nodedb/src/control/planner/sql_plan_convert/visitor/arms_set_ops.rs @@ -72,6 +72,13 @@ macro_rules! impl_set_ops_arms_for_convert_visitor { ) -> crate::Result> { super::super::set_ops::convert_cte(definitions, outer, self.tenant_id, self.ctx) } + + fn subquery( + &mut self, + args: nodedb_sql::SubqueryVisitArgs<'_>, + ) -> crate::Result> { + super::super::set_ops::convert_subquery(args, self.tenant_id, self.ctx) + } }; } diff --git a/nodedb/src/control/security/identity/plan_permission.rs b/nodedb/src/control/security/identity/plan_permission.rs index 6a130d2a5..874d40ca0 100644 --- a/nodedb/src/control/security/identity/plan_permission.rs +++ b/nodedb/src/control/security/identity/plan_permission.rs @@ -85,6 +85,11 @@ pub fn required_permission(plan: &crate::bridge::envelope::PhysicalPlan) -> Perm // it only redistributes rows produced by the wrapped plan. PhysicalPlan::Query(QueryOp::Exchange(op)) => required_permission(&op.child), + // PostProcess only reshapes rows produced by its child (sort / offset / + // distinct / limit / projection); its required permission is the + // child's — recurse rather than assume Read. + PhysicalPlan::Query(QueryOp::PostProcess { input, .. }) => required_permission(input), + PhysicalPlan::Text( TextOp::Search { .. } | TextOp::BM25ScoreScan { .. } diff --git a/nodedb/src/control/server/exchange/resolve/exchange.rs b/nodedb/src/control/server/exchange/resolve/exchange.rs index 9aa110954..4f3fa66bf 100644 --- a/nodedb/src/control/server/exchange/resolve/exchange.rs +++ b/nodedb/src/control/server/exchange/resolve/exchange.rs @@ -435,6 +435,116 @@ async fn resolve_exchange( }))) } + // PostProcess: 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. + PhysicalPlan::Query(QueryOp::PostProcess { + input, + filters, + projection, + sort_keys, + limit, + offset, + distinct, + }) => { + // 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 { + PhysicalPlan::Query(QueryOp::Exchange(ExchangeOp { + child, + mode: ExchangeMode::Gather { as_aggregate }, + })) => (*child, as_aggregate), + other => (other, false), + }; + + // Resolve any Exchange nested inside the child first (e.g. a + // `HashJoin` build-side `Broadcast`) so the plan gathered below is + // self-contained — no Exchange may reach a Data-Plane core. + let child = match Box::pin(resolve_exchange( + state, + database_id, + tenant_id, + child, + trace_id, + txn_id, + captures, + )) + .await? + { + Resolved::Plan(p) => p, + // The unwrapped body is not itself a root Gather / stream; + // surface these defensively without dropping post-processing. + Resolved::Gathered(resp, wms, caps) => { + return Ok(Resolved::Gathered(resp, wms, caps)); + } + Resolved::Stream(s) => return Ok(Resolved::Stream(s)), + }; + + // A vector search emits hit rows (`{id, distance, doc_id, body}`) + // rather than flat storage rows; they need the hit-specific flatten + // (merge `body` columns, surface `distance` / `_surrogate`) so the + // tail can sort / distinct / project by any document column. + let is_vector_hits = matches!( + &child, + PhysicalPlan::Vector( + nodedb_physical::physical_plan::VectorOp::Search { .. } + | nodedb_physical::physical_plan::VectorOp::MultiSearch { .. } + ) + ); + + // Record the child's single base collection in the in-transaction + // read-set at its own observed read-version (mirrors the root + // Gather arm). Autocommit reads skip the catalog lookup. + let probe_collection: Option = if txn_id.is_some() { + child.collection().map(str::to_owned) + } else { + None + }; + + let outcome: GatherOutcome = + 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(state, database_id, tenant_id, &coll)? + { + captures.push(DistributedReadCapture { + scan_plan, + read_version_lsn: outcome.read_version_lsn, + }); + } + + let merged = if as_aggregate { + finalize_aggregate(&outcome.merged_array) + } else { + outcome.merged_array + }; + + // Flatten to the flat relational row shape the `ProviderScan` tail + // consumes: vector hits merge their document `body`; storage + // `{id, data}` rows unwrap their `data`; computed rows pass through. + let rows = if is_vector_hits { + crate::data::executor::response_codec::flatten_vector_hits_to_relational_rows( + &merged, + ) + } else { + crate::data::executor::response_codec::flatten_to_relational_rows(&merged) + }; + Ok(Resolved::Plan(PhysicalPlan::Query(QueryOp::ProviderScan { + provider: None, + rows, + filters, + projection, + sort_keys, + limit, + offset, + distinct, + }))) + } + // All other plan variants: pass through unchanged. other => Ok(Resolved::Plan(other)), } diff --git a/nodedb/src/control/server/exchange/resolve/materialize.rs b/nodedb/src/control/server/exchange/resolve/materialize.rs index ca5afe9f1..4024ddf76 100644 --- a/nodedb/src/control/server/exchange/resolve/materialize.rs +++ b/nodedb/src/control/server/exchange/resolve/materialize.rs @@ -209,6 +209,29 @@ pub(super) async fn materialize_providers( })) } + // PostProcess: recurse into the materialized child so any catalog + // providers nested in the subquery body are filled before resolution. + PhysicalPlan::Query(QueryOp::PostProcess { + input, + filters, + projection, + sort_keys, + limit, + offset, + distinct, + }) => { + let input = Box::pin(materialize_providers(state, identity, *input)).await?; + Ok(PhysicalPlan::Query(QueryOp::PostProcess { + input: Box::new(input), + filters, + projection, + sort_keys, + limit, + offset, + distinct, + })) + } + // All other variants: no catalog providers can be nested here — // pass through unchanged. other => Ok(other), diff --git a/nodedb/src/control/server/response_shape/types.rs b/nodedb/src/control/server/response_shape/types.rs index e0103e673..c07c7b249 100644 --- a/nodedb/src/control/server/response_shape/types.rs +++ b/nodedb/src/control/server/response_shape/types.rs @@ -101,6 +101,11 @@ pub fn describe_plan(plan: &PhysicalPlan) -> PlanKind { // Recurse into the child to determine the plan kind. PhysicalPlan::Query(QueryOp::Exchange(op)) => describe_plan(&op.child), + // PostProcess reshapes a multi-row subquery result; its kind is the + // child's. (Unresolved at this point; resolved to a ProviderScan = + // MultiRow before dispatch.) + PhysicalPlan::Query(QueryOp::PostProcess { input, .. }) => describe_plan(input), + // DML operations that return affected row count. PhysicalPlan::Document(DocumentOp::PointPut { .. }) | PhysicalPlan::Document(DocumentOp::BatchInsert { .. }) diff --git a/nodedb/src/control/server/shared/authorization/requirements/query.rs b/nodedb/src/control/server/shared/authorization/requirements/query.rs index 96223a47a..d55b9b98e 100644 --- a/nodedb/src/control/server/shared/authorization/requirements/query.rs +++ b/nodedb/src/control/server/shared/authorization/requirements/query.rs @@ -70,6 +70,13 @@ pub(super) fn collect_query_requirements<'a>( pending.push(&exchange.child); true } + // Recurse into the post-processor's materialized child so the wrapped + // subquery body's collection is authorized; without this the catch-all + // returns `false` and the body would go unchecked. + PhysicalPlan::Query(QueryOp::PostProcess { input, .. }) => { + pending.push(input); + true + } PhysicalPlan::Query(QueryOp::ProviderScan { provider: Some(provider), .. diff --git a/nodedb/src/control/server/shared/plan_util.rs b/nodedb/src/control/server/shared/plan_util.rs index c9d4a0bdd..a78d00323 100644 --- a/nodedb/src/control/server/shared/plan_util.rs +++ b/nodedb/src/control/server/shared/plan_util.rs @@ -97,6 +97,9 @@ pub(crate) fn extract_collection(plan: &PhysicalPlan) -> Option<&str> { | PhysicalPlan::Graph(GraphOp::WccSuperstep(_)) => None, // Exchange: recurse into the child plan to extract the collection. PhysicalPlan::Query(QueryOp::Exchange(op)) => extract_collection(&op.child), + // PostProcess: recurse into the materialized child (twin of + // `PhysicalPlan::collection`). + PhysicalPlan::Query(QueryOp::PostProcess { input, .. }) => extract_collection(input), // 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.rs b/nodedb/src/control/server/shared/write_admission/predicate/txn_buffering.rs index 8bbcbfb6c..227b75113 100644 --- a/nodedb/src/control/server/shared/write_admission/predicate/txn_buffering.rs +++ b/nodedb/src/control/server/shared/write_admission/predicate/txn_buffering.rs @@ -384,7 +384,8 @@ pub fn plan_requires_txn_buffering(plan: &PhysicalPlan) -> bool { | QueryOp::RecursiveScan { .. } | QueryOp::RecursiveValue { .. } | QueryOp::LateralTopK { .. } - | QueryOp::LateralLoop { .. }, + | QueryOp::LateralLoop { .. } + | QueryOp::PostProcess { .. }, ) => false, // ---- Meta: control / maintenance / internal-mechanism ops. diff --git a/nodedb/src/data/executor/dispatch/query.rs b/nodedb/src/data/executor/dispatch/query.rs index 9cf4a8e5d..56f8884c8 100644 --- a/nodedb/src/data/executor/dispatch/query.rs +++ b/nodedb/src/data/executor/dispatch/query.rs @@ -58,6 +58,15 @@ impl CoreLoop { }, ), + QueryOp::PostProcess { .. } => self.response_error( + task, + crate::bridge::envelope::ErrorCode::Internal { + detail: "PostProcess must be resolved by the coordinator (gathered and \ + lowered to a ProviderScan) before dispatch" + .to_string(), + }, + ), + QueryOp::ProviderScan { rows, filters, diff --git a/nodedb/src/data/executor/response_codec/mod.rs b/nodedb/src/data/executor/response_codec/mod.rs index 123e25414..79d068052 100644 --- a/nodedb/src/data/executor/response_codec/mod.rs +++ b/nodedb/src/data/executor/response_codec/mod.rs @@ -41,4 +41,6 @@ pub(in crate::data::executor) use hits::{ NeighborEntry, NeighborMultiEntry, SubgraphEdge, VectorSearchHit, }; pub(crate) use raw::{decode_raw_scan_to_docs, encode_raw_document_rows}; -pub use raw::{encode_binary_rows, flatten_to_relational_rows}; +pub use raw::{ + encode_binary_rows, flatten_to_relational_rows, flatten_vector_hits_to_relational_rows, +}; diff --git a/nodedb/src/data/executor/response_codec/raw.rs b/nodedb/src/data/executor/response_codec/raw.rs index f2a08750d..598d45bb7 100644 --- a/nodedb/src/data/executor/response_codec/raw.rs +++ b/nodedb/src/data/executor/response_codec/raw.rs @@ -135,6 +135,77 @@ pub fn flatten_to_relational_rows(bytes: &[u8]) -> Vec { encode_binary_rows(&flat) } +/// Flatten a gathered array of vector-search hits into flat relational rows for +/// the post-processing tail (`QueryOp::PostProcess` → `ProviderScan`). +/// +/// A hit is `{id: , distance, doc_id?, body?: }`. +/// This mirrors the client-facing vector response translator +/// (`control::server::response_translate::vector`) but emits MessagePack rows +/// instead of JSON: the document `body`'s columns become top-level (so ORDER BY +/// / DISTINCT / projection can reference any document column), `distance` is +/// surfaced, `_surrogate` carries the internal id, and the document's own `id` +/// column wins over the surrogate. Rows without a body (e.g. a +/// `skip_payload_fetch` hit) still surface `id` / `distance` / `_surrogate`. +/// +/// RLS is already enforced by the Data Plane (`VectorOp::Search.rls_filters`, +/// injected into the inner search before dispatch), so these gathered hits are +/// post-RLS. +pub fn flatten_vector_hits_to_relational_rows(bytes: &[u8]) -> Vec { + use nodedb_types::Value; + + #[derive(zerompk::FromMessagePack)] + #[msgpack(map)] + struct Hit { + id: u32, + distance: f32, + doc_id: Option, + body: Option>, + } + + let hits: Vec = match zerompk::from_msgpack(bytes) { + Ok(h) => h, + // Not a hit array (already flat, or a non-row payload): leave as-is. + Err(_) => return bytes.to_vec(), + }; + + let rows: Vec> = hits + .into_iter() + .filter_map(|h| { + // Base columns come from the document body (its own `id` wins over + // the internal surrogate). The body is *bare* msgpack (the storage + // wire shape), so decode/encode with the native `Value` codec — the + // derived `zerompk` codec is tagged and would corrupt the row. + let mut fields: std::collections::HashMap = match h + .body + .as_deref() + .and_then(|b| nodedb_types::value_from_msgpack(b).ok()) + { + Some(Value::Object(map)) => map, + _ => std::collections::HashMap::new(), + }; + fields + .entry("distance".to_string()) + .or_insert(Value::Float(h.distance as f64)); + if !fields.contains_key("id") { + match h.doc_id { + Some(pk) => { + fields.insert("id".to_string(), Value::String(pk)); + } + None => { + fields.insert("id".to_string(), Value::Integer(h.id as i64)); + } + } + } + fields + .entry("_surrogate".to_string()) + .or_insert(Value::Integer(h.id as i64)); + nodedb_types::value_to_msgpack(&Value::Object(fields)).ok() + }) + .collect(); + + encode_binary_rows(&rows) +} + /// Encode a list of pre-built binary msgpack rows into a single msgpack array. /// /// Each row is already a valid msgpack value (typically a map). This just @@ -161,3 +232,84 @@ pub(super) fn msgpack_write_array_header(buf: &mut Vec, len: usize) { buf.extend_from_slice(&(len as u32).to_be_bytes()); } } + +#[cfg(test)] +mod tests { + use super::*; + use nodedb_types::Value; + use std::collections::HashMap; + + /// A vector hit in the Data-Plane wire shape: a *bare* msgpack map with + /// the same field names `VectorSearchHit` emits (`#[msgpack(map)]`). + #[derive(zerompk::ToMessagePack)] + #[msgpack(map)] + struct TestHit { + id: u32, + distance: f32, + doc_id: Option, + body: Option>, + } + + fn bare_doc(pairs: Vec<(&str, Value)>) -> Vec { + let map = pairs.into_iter().map(|(k, v)| (k.to_string(), v)).collect(); + nodedb_types::value_to_msgpack(&Value::Object(map)).unwrap() + } + + /// Decode the flattened output (a bare msgpack array of bare row maps) into + /// per-row column maps. + fn decode_rows(out: &[u8]) -> Vec> { + match nodedb_types::value_from_msgpack(out) { + Ok(Value::Array(rows)) => rows + .into_iter() + .map(|r| match r { + Value::Object(m) => m, + other => panic!("expected row map, got {other:?}"), + }) + .collect(), + other => panic!("expected row array, got {other:?}"), + } + } + + #[test] + fn flatten_vector_hits_merges_body_and_surfaces_metadata() { + let body = bare_doc(vec![ + ("id", Value::String("r0".into())), + ("tag", Value::String("keep".into())), + ]); + let input = zerompk::to_msgpack_vec(&vec![TestHit { + id: 7, + distance: 0.5, + doc_id: None, + body: Some(body), + }]) + .unwrap(); + + let rows = decode_rows(&flatten_vector_hits_to_relational_rows(&input)); + assert_eq!(rows.len(), 1); + let row = &rows[0]; + // Document's own `id` wins over the internal surrogate. + assert_eq!(row.get("id"), Some(&Value::String("r0".into()))); + // Payload column surfaced to top level. + assert_eq!(row.get("tag"), Some(&Value::String("keep".into()))); + // Search metadata preserved. + assert_eq!(row.get("distance"), Some(&Value::Float(0.5f32 as f64))); + assert_eq!(row.get("_surrogate"), Some(&Value::Integer(7))); + } + + #[test] + fn flatten_vector_hits_without_body_falls_back_to_surrogate_id() { + let input = zerompk::to_msgpack_vec(&vec![TestHit { + id: 42, + distance: 1.25, + doc_id: None, + body: None, + }]) + .unwrap(); + + let rows = decode_rows(&flatten_vector_hits_to_relational_rows(&input)); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].get("id"), Some(&Value::Integer(42))); + assert_eq!(rows[0].get("_surrogate"), Some(&Value::Integer(42))); + assert_eq!(rows[0].get("distance"), Some(&Value::Float(1.25))); + } +} diff --git a/nodedb/tests/sql_search_subquery_composition.rs b/nodedb/tests/sql_search_subquery_composition.rs new file mode 100644 index 000000000..35fba7c35 --- /dev/null +++ b/nodedb/tests/sql_search_subquery_composition.rs @@ -0,0 +1,268 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Integration coverage for `SEARCH ... USING VECTOR(...)` in subquery +//! position. +//! +//! A `SEARCH` result is a relation, so it has to compose: usable as a derived +//! table, joinable, and reachable from an `IN (...)` predicate. Without that, +//! every hybrid vector-plus-relational query needs two round trips and the +//! relational filter can only be applied after the k-NN cut, which silently +//! shrinks the result set a caller asked for. + +mod common; + +use common::pgwire_harness::TestServer; + +/// Fixture rows, nearest-first for the query vector used below: `r0` (exact +/// match), `r1`, then `r2`. `tag` splits them so a filter over the k-NN result +/// can be told apart from a filter applied before the cut. +async fn create_vector_collection(server: &TestServer, name: &str) { + server + .exec(&format!("CREATE COLLECTION {name}")) + .await + .unwrap(); + server + .exec(&format!( + "CREATE VECTOR INDEX idx_{name}_emb ON {name} (embedding) METRIC cosine DIM 4" + )) + .await + .unwrap(); + for (id, tag, v) in [ + ("r0", "keep", [0.10f32, 0.20, 0.30, 0.40]), + ("r1", "drop", [0.11, 0.21, 0.31, 0.41]), + ("r2", "keep", [0.90, 0.80, 0.70, 0.60]), + ] { + server + .exec(&format!( + "INSERT INTO {name} (id, tag, embedding) VALUES \ + ('{id}', '{tag}', ARRAY[{},{},{},{}])", + v[0], v[1], v[2], v[3] + )) + .await + .unwrap(); + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn search_is_usable_as_a_derived_table() { + let server = TestServer::start().await; + create_vector_collection(&server, "vec_derived").await; + + let rows = server + .query_text( + "SELECT id FROM \ + (SEARCH vec_derived USING VECTOR(embedding, ARRAY[0.1, 0.2, 0.3, 0.4], 2)) s", + ) + .await + .unwrap(); + assert_eq!( + rows, + vec!["r0".to_string(), "r1".to_string()], + "SEARCH in FROM position must yield the same k rows, in distance order, as the top-level form" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn outer_limit_narrows_the_search_result() { + let server = TestServer::start().await; + create_vector_collection(&server, "vec_outer_limit").await; + + let rows = server + .query_text( + "SELECT id FROM \ + (SEARCH vec_outer_limit USING VECTOR(embedding, ARRAY[0.1, 0.2, 0.3, 0.4], 3)) s \ + LIMIT 1", + ) + .await + .unwrap(); + assert_eq!( + rows, + vec!["r0".to_string()], + "an outer LIMIT must cut the k-NN result, got: {rows:?}" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn relational_filter_applies_over_search_results() { + let server = TestServer::start().await; + create_vector_collection(&server, "vec_filtered").await; + + // The predicate reaches the engine as a search filter, so the k-NN cut is + // taken over matching rows: k = 2 yields the two nearest rows tagged + // 'keep' (r0, r2), not the tagged subset of the two nearest (r0 alone). + let rows = server + .query_text( + "SELECT id FROM \ + (SEARCH vec_filtered USING VECTOR(embedding, ARRAY[0.1, 0.2, 0.3, 0.4], 2)) s \ + WHERE s.tag = 'keep'", + ) + .await + .unwrap(); + assert_eq!( + rows, + vec!["r0".to_string(), "r2".to_string()], + "the WHERE clause must run inside the engine, and k counts matching rows, got: {rows:?}" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn search_result_projects_a_single_column() { + let server = TestServer::start().await; + create_vector_collection(&server, "vec_projected").await; + + let rows = server + .query_rows( + "SELECT s.id FROM \ + (SEARCH vec_projected USING VECTOR(embedding, ARRAY[0.1, 0.2, 0.3, 0.4], 2)) s", + ) + .await + .unwrap(); + assert_eq!(rows.len(), 2, "got: {rows:?}"); + for row in &rows { + assert_eq!( + row.len(), + 1, + "the outer projection must narrow the SEARCH output, got: {row:?}" + ); + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn search_accepts_a_quoted_collection_name() { + let server = TestServer::start().await; + create_vector_collection(&server, "vec_quoted").await; + + let bare = server + .query_text("SEARCH vec_quoted USING VECTOR(embedding, ARRAY[0.1, 0.2, 0.3, 0.4], 2)") + .await + .unwrap(); + let quoted = server + .query_text("SEARCH \"vec_quoted\" USING VECTOR(embedding, ARRAY[0.1, 0.2, 0.3, 0.4], 2)") + .await + .unwrap(); + assert_eq!( + quoted, bare, + "a quoted collection name must search the same collection as the bare form" + ); + + let quoted_subquery = server + .query_text( + "SELECT id FROM \ + (SEARCH \"vec_quoted\" USING VECTOR(embedding, ARRAY[0.1, 0.2, 0.3, 0.4], 2)) s", + ) + .await + .unwrap(); + assert_eq!(quoted_subquery, vec!["r0".to_string(), "r1".to_string()]); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn malformed_search_subquery_is_rejected() { + let server = TestServer::start().await; + create_vector_collection(&server, "vec_bad_args").await; + + // One argument is neither the two-arg (vector, k) nor the three-arg + // (field, vector, k) form — it must not be rewritten into a SELECT. + server + .expect_error( + "SELECT * FROM (SEARCH vec_bad_args USING VECTOR(ARRAY[0.1, 0.2, 0.3, 0.4])) s", + "parse error", + ) + .await; +} + +// ── Post-processing over the k-NN result (QueryOp::PostProcess) ────────────── +// +// An outer ORDER BY / OFFSET / DISTINCT (and a LIMIT that must apply after a +// reorder) cannot be absorbed by the vector-search leaf. Before the +// post-processing operator these were silently dropped, returning the raw +// distance-ordered k rows. These tests pin the corrected behaviour end to end. + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn outer_order_by_payload_column_reorders_search_result() { + let server = TestServer::start().await; + create_vector_collection(&server, "vec_order_payload").await; + + // The three nearest are r0, r1, r2 (distance order). An outer ORDER BY on a + // *document* column (`id`) must re-sort them — proving payload columns are + // materialized from the hit body, not left nested. + let rows = server + .query_text( + "SELECT id FROM \ + (SEARCH vec_order_payload USING VECTOR(embedding, ARRAY[0.1, 0.2, 0.3, 0.4], 3)) s \ + ORDER BY s.id DESC", + ) + .await + .unwrap(); + assert_eq!( + rows, + vec!["r2".to_string(), "r1".to_string(), "r0".to_string()], + "outer ORDER BY on a payload column must reorder the k-NN result, got: {rows:?}" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn outer_offset_skips_leading_search_rows() { + let server = TestServer::start().await; + create_vector_collection(&server, "vec_offset").await; + + // No explicit ORDER BY: rows stay in distance order (r0, r1, r2); OFFSET 1 + // drops the nearest. + let rows = server + .query_text( + "SELECT id FROM \ + (SEARCH vec_offset USING VECTOR(embedding, ARRAY[0.1, 0.2, 0.3, 0.4], 3)) s \ + OFFSET 1", + ) + .await + .unwrap(); + assert_eq!( + rows, + vec!["r1".to_string(), "r2".to_string()], + "outer OFFSET must skip leading rows of the k-NN result, got: {rows:?}" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn outer_distinct_dedups_search_result() { + let server = TestServer::start().await; + create_vector_collection(&server, "vec_distinct").await; + + // Tags across the three nearest are keep / drop / keep — DISTINCT on the + // projected column must collapse to two rows. + let mut rows = server + .query_text( + "SELECT DISTINCT tag FROM \ + (SEARCH vec_distinct USING VECTOR(embedding, ARRAY[0.1, 0.2, 0.3, 0.4], 3)) s", + ) + .await + .unwrap(); + rows.sort(); + assert_eq!( + rows, + vec!["drop".to_string(), "keep".to_string()], + "outer DISTINCT must dedup the projected column, got: {rows:?}" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn outer_order_by_distance_then_limit_takes_farthest() { + let server = TestServer::start().await; + create_vector_collection(&server, "vec_order_distance").await; + + // ORDER BY distance DESC reverses the k-NN order, so LIMIT 1 takes the + // FARTHEST of the three nearest (r2) — a LIMIT that must apply after the + // reorder, not fold into the search top_k. + let rows = server + .query_text( + "SELECT id FROM \ + (SEARCH vec_order_distance USING VECTOR(embedding, ARRAY[0.1, 0.2, 0.3, 0.4], 3)) s \ + ORDER BY s.distance DESC LIMIT 1", + ) + .await + .unwrap(); + assert_eq!( + rows, + vec!["r2".to_string()], + "LIMIT after an outer ORDER BY must cut the reordered rows, got: {rows:?}" + ); +}