fix(engine): stream unranked scans and hydrate return-only columns after the limit

changelog.d/query-scans-stream-and-limits-hydrate.fixed.md

@@ -0,0 +1,11 @@
+- A read that projects a wide column no longer fails with `resource limit
+  exceeded for query_memory_bytes` once the type's projected columns exceed the
+  150 MiB query pool, even when it returns one row: every unranked node scan now
+  streams bounded batches instead of holding the whole type, so an unordered
+  `limit` stops reading after its first batches and a top-k holds only its kept
+  rows. Under a `limit` over a type with more than four rows per kept row, the
+  columns only the output reads are fetched by row address for the kept rows
+  alone; `explain` shows a `HydrateColumns` node and the
+  `late_materialization` pass. See [execution][hydrate-execution].
+
+[hydrate-execution]: ../docs/dev/execution.md

crates/omnigraph-gqt-core/src/plan.rs

@@ -21,6 +21,7 @@
 //! scan Passage as $p: runtime filter text
 //! scan Passage as $p: no runtime filter
 //! filter reads [a.state, b.state]
+//! hydrate $d: columns [body, title]
 //! pass projection_pushdown
 //! not pass aggregate_pushdown
 //! ```
@@ -40,7 +41,7 @@ use serde_json::Value;
 
 use crate::report::Row;
 
-const FORMS: &str = "forms: `scan <Type>[ as $var]: columns [a, b]`, `scan <Type>[ as $var]: not columns [a, b]`, `scan <Type>[ as $var]: filter reads [v.a]`, `scan <Type>[ as $var]: no filter`, `scan <Type>[ as $var]: access id_lookup`, `hash join $var[ ran <hash_join|id_lookup>]`, `scan <Type>[ as $var]: ranked <nearest|bm25>[ fetch <n>][ nprobes <n>]`, `scan <Type>[ as $var]: runtime filter <column>`, `scan <Type>[ as $var]: no runtime filter`, `contains join $h.x contains $n.y`, `cross join $h.x contains $n.y`, `expand $src <Edge> $dst: mode <csr|indexed_scan>[ ran <csr|indexed_scan>]`, `filter reads [a.x, b.y]`, `sort tiebreak [$a.@id, $e.@type]`, `rank fuse row tiebreak [$e.@type, $e.@id]`, `expand $a $b: selection alternation [Knows out, Likes in]`, `sort no tiebreak`, `aggregate <column>: <func>(<Type>) <accumulator> <overflow> -> <Type>`, `block aggregate <gq>: <func>(<Type>) <accumulator> <overflow> -> <Type>`, `result columns [<name>: <Type>, ...]`, `type <gq>: <Type>`, `cast <gq>: <Type> -> <Type>`, `no cast <gq>`, `pass <name>`, `not pass <name>`";
+const FORMS: &str = "forms: `scan <Type>[ as $var]: columns [a, b]`, `scan <Type>[ as $var]: not columns [a, b]`, `scan <Type>[ as $var]: filter reads [v.a]`, `scan <Type>[ as $var]: no filter`, `scan <Type>[ as $var]: access id_lookup`, `hash join $var[ ran <hash_join|id_lookup>]`, `scan <Type>[ as $var]: ranked <nearest|bm25>[ fetch <n>][ nprobes <n>]`, `scan <Type>[ as $var]: runtime filter <column>`, `scan <Type>[ as $var]: no runtime filter`, `contains join $h.x contains $n.y`, `cross join $h.x contains $n.y`, `expand $src <Edge> $dst: mode <csr|indexed_scan>[ ran <csr|indexed_scan>]`, `filter reads [a.x, b.y]`, `sort tiebreak [$a.@id, $e.@type]`, `rank fuse row tiebreak [$e.@type, $e.@id]`, `expand $a $b: selection alternation [Knows out, Likes in]`, `sort no tiebreak`, `aggregate <column>: <func>(<Type>) <accumulator> <overflow> -> <Type>`, `block aggregate <gq>: <func>(<Type>) <accumulator> <overflow> -> <Type>`, `result columns [<name>: <Type>, ...]`, `type <gq>: <Type>`, `cast <gq>: <Type> -> <Type>`, `no cast <gq>`, `hydrate $var: columns [a, b]`, `pass <name>`, `not pass <name>`";
 
 const ID_LOOKUP: &str = "id_lookup";
 const JOIN_SIDES: [&str; 2] = ["hash_join", "id_lookup"];
@@ -171,6 +172,12 @@ pub(crate) enum PlanLine {
     Sort {
         tiebreak: Vec<String>,
     },
+    /// A physical `HydrateColumns` fetches exactly `columns` of `$binding` by
+    /// row address for the rows that reached the output.
+    Hydrate {
+        binding: String,
+        columns: Vec<String>,
+    },
     /// The optimizer pass `name` fired, or did not when `negated`.
     Pass {
         name: String,
@@ -372,6 +379,26 @@ pub(crate) fn parse_plan_body(body: &[(usize, &str)]) -> Result<Vec<PlanLine>, S
             lines.push(PlanLine::Sort { tiebreak });
             continue;
         }
+        if let Some(rest) = line.strip_prefix("hydrate ") {
+            let (binding, claim) = rest
+                .split_once(':')
+                .ok_or_else(|| refused("claims `hydrate $var: columns [a, b]`"))?;
+            let binding = binding
+                .trim()
+                .strip_prefix('$')
+                .filter(|binding| identifier(binding))
+                .ok_or_else(|| refused("names the hydrated binding with `$`"))?;
+            let columns = claim
+                .trim()
+                .strip_prefix("columns")
+                .and_then(column_list)
+                .ok_or_else(|| refused("claims `hydrate $var: columns [a, b]`"))?;
+            lines.push(PlanLine::Hydrate {
+                binding: binding.to_string(),
+                columns,
+            });
+            continue;
+        }
         if let Some(rest) = line.strip_prefix("hash join ") {
             let rest = rest.trim_start();
             let (binding, claim) = rest.split_once(' ').unwrap_or((rest, ""));
@@ -868,6 +895,8 @@ struct PlannedNodes {
     sorts: Vec<Result<Vec<String>, String>>,
     fusions: Vec<Result<Vec<String>, String>>,
     selections: Vec<(String, String, Value)>,
+    /// `(binding, sorted columns)` of every `HydrateColumns` binding.
+    hydrations: Vec<(String, Vec<String>)>,
 }
 
 #[derive(Debug)]
@@ -1089,6 +1118,31 @@ fn planned_physical(node: &Value, out: &mut PlannedNodes) {
     };
     let id = node.get("id").and_then(Value::as_u64);
     match node.get("node").and_then(Value::as_str) {
+        Some("HydrateColumns") => {
+            for binding in node
+                .get("bindings")
+                .and_then(Value::as_array)
+                .into_iter()
+                .flatten()
+            {
+                let columns = binding
+                    .get("columns")
+                    .and_then(Value::as_array)
+                    .into_iter()
+                    .flatten()
+                    .filter_map(Value::as_str)
+                    .map(str::to_string)
+                    .collect();
+                out.hydrations.push((
+                    binding
+                        .get("binding")
+                        .and_then(Value::as_str)
+                        .unwrap_or_default()
+                        .to_string(),
+                    sorted_set(columns),
+                ));
+            }
+        }
         Some("Aggregate") => match node.get("aggregates").and_then(Value::as_array) {
             Some(aggregates) => {
                 out.aggregates
@@ -1470,6 +1524,21 @@ pub(crate) fn plan_mismatch(
                     ));
                 }
             }
+            PlanLine::Hydrate { binding, columns } => {
+                let want = sorted_set(columns.clone());
+                let found: Vec<&Vec<String>> = nodes
+                    .hydrations
+                    .iter()
+                    .filter(|(hydrated, _)| hydrated == binding)
+                    .map(|(_, columns)| columns)
+                    .collect();
+                if !found.contains(&&want) {
+                    return Some(format!(
+                        "expect plan: no `HydrateColumns` fetches {want:?} of `${binding}`; the plan hydrates {:?}",
+                        nodes.hydrations
+                    ));
+                }
+            }
             PlanLine::Sort { tiebreak } => {
                 if nodes.sorts.is_empty() {
                     return Some("expect plan: the physical plan has no `Sort`".to_string());
@@ -2587,6 +2656,36 @@ mod tests {
         }
     }
 
+    /// A `hydrate` line reads the physical `HydrateColumns`: the exact set
+    /// of one binding's hydrated columns, in any order.
+    #[test]
+    fn hydrate_claims_read_the_physical_hydration() {
+        let explain = json!({
+            "physical_plan": {"node": "HydrateColumns", "id": 4, "bindings": [
+                {"binding": "d", "table": "node:Doc", "columns": ["title", "body"]}
+            ], "inputs": [{"node": "Page", "id": 3, "inputs": []}]},
+            "passes": ["resolve", "late_materialization"],
+        });
+        let lines = parse_plan_body(&[(0, "hydrate $d: columns [body, title]")]).unwrap();
+        assert_eq!(check(&lines, &explain), None);
+        for (claim, want) in [
+            ("hydrate $d: columns [body]", "no `HydrateColumns` fetches"),
+            ("hydrate $e: columns [body, title]", "of `$e`"),
+        ] {
+            let lines = parse_plan_body(&[(0, claim)]).unwrap();
+            let mismatch = check(&lines, &explain).unwrap();
+            assert!(mismatch.contains(want), "{claim}: {mismatch}");
+        }
+        for refused in [
+            "hydrate d: columns [body]",
+            "hydrate $d columns [body]",
+            "hydrate $d: [body]",
+            "hydrate $d: columns []",
+        ] {
+            assert!(parse_plan_body(&[(0, refused)]).is_err(), "{refused}");
+        }
+    }
+
     /// A `contains join` line reads the physical `ContainsJoin`'s columns
     /// and a `runtime filter` claim the marked scan's column; the collected
     /// side is unmarked, and a `cross join` line the plain product's filters.

crates/omnigraph-gqt/README.md

@@ -366,6 +366,7 @@ hash join $e
 expand $d Knows $e: mode indexed_scan
 filter reads [d.rank, e.rank]
 sort tiebreak [$d, $e]
+hydrate $d: columns [body]
 pass projection_pushdown
 not pass aggregate_pushdown
 aggregate total: sum(I64?) exact_integer round_to_nearest -> F64?
@@ -383,8 +384,11 @@ keys a physical `Sort` appends after user order keys. `$a` abbreviates `$a.@id`;
 `sort no tiebreak` requires an empty list. `rank fuse row tiebreak [...]` checks
 the exact downstream keys of `RankFuse`, and `rank fuse no row tiebreak` requires
 none. Dropping a type key or swapping key order fails these assertions.
-`pass <name>` states that a named optimizer pass fired, `not pass <name>` that
-it did not. Projection/read lists are sets; identity keys and selection members
+`hydrate $d: columns [..]` states that a physical `HydrateColumns` fetches
+exactly those columns of `$d` by row address above the root limit (pass
+`late_materialization`); the `scan` lines keep reading the logical plan, whose
+projection still lists them. `pass <name>` states that a named optimizer pass
+fired, `not pass <name>` that it did not. Projection/read lists are sets; identity keys and selection members
 are ordered lists. A mismatch prints the whole explain document.
 `aggregate <column>: <func>(<Type>) <accumulator> <overflow> -> <Type>`
 checks the named output's aggregate function, input type, accumulator,

crates/omnigraph-gqt/cases/v2/planner/hydrate_return_columns_above_the_limit.gqt

@@ -0,0 +1,180 @@
+# issue: none
+# notes: Twelve docs; d07 and d10 have no body. d00 and d01 both cite d10, d02 cites
+# notes: d11, d03 cites d05. Pass 6 fetches a return column only the output reads by
+# notes: row address above the root limit when its binding's scan estimates more than
+# notes: four rows per kept row. Steps 1-3 hydrate (12 > 4 x 2; 12 > 4 x 1 for $b,
+# notes: whose scan step 3's key lookup on $a does not bound). Steps 4-6 do not: a
+# notes: sorted alias, 12 is not above 4 x 3, a key lookup reads one row. Step 2
+# notes: reaches $b through a hash join, step 3 by id lookup: both destination paths
+# notes: carry the row address. Every rows step is compared with engine v1.
+
+--- runner
+timeout_ms: 10000
+environments:
+  - target: omnigraph-engine
+    storage: local-filesystem
+
+--- schema
+node Doc {
+    slug: String @key
+    rank: I32
+    title: String
+    body: String?
+    embedding: Vector(2)
+}
+
+edge Cites: Doc -> Doc
+
+--- seed
+{"type":"Doc","data":{"slug":"d00","rank":0,"title":"t-00","embedding":[0.0,0.0],"body":"body 00"}}
+{"type":"Doc","data":{"slug":"d01","rank":1,"title":"t-01","embedding":[1.0,0.0],"body":"body 01"}}
+{"type":"Doc","data":{"slug":"d02","rank":2,"title":"t-02","embedding":[2.0,0.0],"body":"body 02"}}
+{"type":"Doc","data":{"slug":"d03","rank":3,"title":"t-03","embedding":[3.0,0.0],"body":"body 03"}}
+{"type":"Doc","data":{"slug":"d04","rank":4,"title":"t-04","embedding":[4.0,0.0],"body":"body 04"}}
+{"type":"Doc","data":{"slug":"d05","rank":5,"title":"t-05","embedding":[5.0,0.0],"body":"body 05"}}
+{"type":"Doc","data":{"slug":"d06","rank":6,"title":"t-06","embedding":[6.0,0.0],"body":"body 06"}}
+{"type":"Doc","data":{"slug":"d07","rank":7,"title":"t-07","embedding":[7.0,0.0]}}
+{"type":"Doc","data":{"slug":"d08","rank":8,"title":"t-08","embedding":[8.0,0.0],"body":"body 08"}}
+{"type":"Doc","data":{"slug":"d09","rank":9,"title":"t-09","embedding":[9.0,0.0],"body":"body 09"}}
+{"type":"Doc","data":{"slug":"d10","rank":10,"title":"t-10","embedding":[10.0,0.0]}}
+{"type":"Doc","data":{"slug":"d11","rank":11,"title":"t-11","embedding":[11.0,0.0],"body":"body 11"}}
+{"edge":"Cites","from":"d00","to":"d10","data":{}}
+{"edge":"Cites","from":"d01","to":"d10","data":{}}
+{"edge":"Cites","from":"d02","to":"d11","data":{}}
+{"edge":"Cites","from":"d03","to":"d05","data":{}}
+
+--- query
+query top_two() {
+    match { $d: Doc }
+    return { $d.slug, $d.body, $d.title }
+    order { $d.rank desc }
+    limit 2
+}
+
+--- expect ordered
+{"d.body":"body 11","d.slug":"d11","d.title":"t-11"}
+{"d.slug":"d10","d.title":"t-10"}
+
+--- expect shape
+d.slug: String
+d.body: String?
+d.title: String
+
+--- expect plan
+pass late_materialization
+hydrate $d: columns [body, title]
+scan Doc as $d: columns [__id, slug, rank, body, title]
+
+--- expect same as v1
+
+--- query
+query cited_twice() {
+    match {
+        $a: Doc
+        $a cites $b
+    }
+    return { $a.slug, $b.slug, $b.body, $b.embedding }
+    order { $a.slug }
+    limit 2
+}
+
+--- expect ordered
+{"a.slug":"d00","b.embedding":[10,0],"b.slug":"d10"}
+{"a.slug":"d01","b.embedding":[10,0],"b.slug":"d10"}
+
+--- expect shape
+a.slug: String
+b.slug: String
+b.body: String?
+b.embedding: Vector(2)
+
+--- expect plan
+pass late_materialization
+hydrate $b: columns [body, embedding]
+hash join $b
+
+--- expect same as v1
+
+--- query
+query first_citation() {
+    match {
+        $a: Doc { slug: "d02" }
+        $a cites $b
+    }
+    return { $b.body }
+    limit 1
+}
+
+--- expect unordered
+{"b.body":"body 11"}
+
+--- expect shape
+b.body: String?
+
+--- expect plan
+pass late_materialization
+hydrate $b: columns [body]
+scan Doc as $b: access id_lookup
+
+--- expect same as v1
+
+--- query
+query by_title() {
+    match { $d: Doc }
+    return { $d.slug, $d.title as t }
+    order { t desc }
+    limit 2
+}
+
+--- expect ordered
+{"d.slug":"d11","t":"t-11"}
+{"d.slug":"d10","t":"t-10"}
+
+--- expect shape
+d.slug: String
+t: String
+
+--- expect plan
+not pass late_materialization
+
+--- expect same as v1
+
+--- query
+query lowest_three() {
+    match { $d: Doc }
+    return { $d.slug, $d.body }
+    order { $d.rank }
+    limit 3
+}
+
+--- expect ordered
+{"d.body":"body 00","d.slug":"d00"}
+{"d.body":"body 01","d.slug":"d01"}
+{"d.body":"body 02","d.slug":"d02"}
+
+--- expect shape
+d.slug: String
+d.body: String?
+
+--- expect plan
+not pass late_materialization
+
+--- expect same as v1
+
+--- query
+query one_body() {
+    match { $d: Doc { slug: "d03" } }
+    return { $d.body }
+    limit 1
+}
+
+--- expect unordered
+{"d.body":"body 03"}
+
+--- expect shape
+d.body: String?
+
+--- expect plan
+not pass late_materialization
+
+--- expect same as v1