hawdb-executor 0.15.0

use super::*;
use hawdb_plan_cypher::{GraphMatchStep, LogicalPlan, ProjectionExpression};

// The oracle lowers the generic operators directly, independently of optimizer fast paths.
fn lower(plan: LogicalPlan) -> PhysicalPlan {
    match plan {
        LogicalPlan::GraphMatch { program, input } => PhysicalPlan::GraphMatchExec {
            program,
            input: input.map(|input| Box::new(lower(*input))),
        },
        LogicalPlan::Project { items, input } => PhysicalPlan::ProjectExec {
            items,
            input: Box::new(lower(*input)),
        },
        LogicalPlan::Aggregate {
            group_keys,
            items,
            input,
        } => PhysicalPlan::AggregateExec {
            group_keys,
            items,
            input: Box::new(lower(*input)),
        },
        LogicalPlan::Filter { predicate, input } => PhysicalPlan::FilterExec {
            predicate,
            input: Box::new(lower(*input)),
        },
        LogicalPlan::Sort { items, input } => PhysicalPlan::SortExec {
            items,
            input: Box::new(lower(*input)),
        },
        LogicalPlan::Limit {
            offset,
            limit,
            input,
        } => PhysicalPlan::LimitExec {
            offset,
            limit,
            input: Box::new(lower(*input)),
        },
        LogicalPlan::Distinct { input } => PhysicalPlan::DistinctExec {
            input: Box::new(lower(*input)),
        },
        _ => panic!("unexpected generic plan: {plan:?}"),
    }
}

fn execute(query: &str) -> Result<Vec<Binding>> {
    let logical = hawdb_plan_cypher::plan_pipeline_query(query, &BTreeMap::new())?;
    execute_logical(logical)
}

fn execute_logical(logical: LogicalPlan) -> Result<Vec<Binding>> {
    let plan = lower(logical);
    with_context(None, |context| {
        let mut rows = Vec::new();
        execute_binding_batches(&plan, context, ExecutionLimit::unlimited(), &mut |batch| {
            rows.extend(batch);
            Ok(BatchControl::Continue)
        })?;
        Ok(rows)
    })
}

#[test]
fn repeated_with_preserves_aliases_graph_identity_and_scalar_shadowing() {
    let rows = execute("MATCH (n:Memory) WITH n AS renamed, n.id AS score WITH renamed, score AS n MATCH (renamed)-[:MENTIONS]->(m:Memory) RETURN id(renamed) AS identity, n, m.id AS target").unwrap();
    assert_eq!(rows.len(), 1);
    assert_eq!(
        rows[0].values,
        BTreeMap::from([
            ("identity".into(), Value::Int(0)),
            ("n".into(), Value::Int(1)),
            ("target".into(), Value::Int(2))
        ])
    );
    for query in [
        "MATCH (n:Memory) WITH n.id AS value RETURN n",
        "MATCH (n:Memory) WITH n.id AS n MATCH (n) RETURN n",
    ] {
        assert!(execute(query).is_err(), "{query}");
    }
}

#[test]
fn grouped_nodes_are_restored_before_later_matches() {
    let rows = execute("MATCH (n:Memory) WITH n AS anchor, COUNT(n) AS seen OPTIONAL MATCH (anchor)-[:MENTIONS]->(m:Memory) WITH anchor, seen, COUNT(m) AS links RETURN anchor.id AS id, seen, links ORDER BY id").unwrap();
    assert_eq!(
        rows.iter()
            .map(|row| row.values["links"].clone())
            .collect::<Vec<_>>(),
        vec![Value::Int(1), Value::Int(0)]
    );
    assert_eq!(rows[0].values["id"], Value::Int(1));
    assert_eq!(rows[1].values["id"], Value::Int(2));
    assert!(rows.iter().all(|row| row.values["seen"] == Value::Int(1)));
}

#[test]
fn optional_where_covers_the_complete_pattern_and_preserves_null_entities() {
    let rows = execute("MATCH (n:Memory) OPTIONAL MATCH (n)-[r:MENTIONS]->(m:Memory) WHERE m.id = 99 RETURN n.id AS id, m, id(m) AS identity, type(r) AS kind ORDER BY id").unwrap();
    assert_eq!(rows.len(), 2);
    for row in rows {
        for key in ["m", "identity", "kind"] {
            assert_eq!(row.values[key], Value::Null);
        }
    }
    let rows = execute("MATCH (n:Memory) OPTIONAL MATCH (n)-[:MENTIONS]->(m:Memory), (x:Missing) RETURN n.id AS id, m, x ORDER BY id").unwrap();
    assert_eq!(rows.len(), 2);
    assert!(rows
        .iter()
        .all(|row| row.values["m"] == Value::Null && row.values["x"] == Value::Null));
}

#[test]
fn optional_node_correlation_is_not_a_query_template() {
    let rows = execute("MATCH (left:Memory) OPTIONAL MATCH (right:Memory) WHERE right.id = left.id WITH left AS kept, COUNT(right) AS same OPTIONAL MATCH (kept)-[:MENTIONS]->(linked:Memory) WITH kept, same, COUNT(linked) AS linked_count OPTIONAL MATCH (kept)-[:ABSENT]->(missing:Memory) RETURN kept.id AS id, same, linked_count, COUNT(missing) AS absent ORDER BY id").unwrap();
    assert_eq!(rows.len(), 2);
    assert!(rows
        .iter()
        .all(|row| row.values["same"] == Value::Int(1) && row.values["absent"] == Value::Int(0)));
    assert_eq!(rows[0].values["linked_count"], Value::Int(1));
    assert_eq!(rows[1].values["linked_count"], Value::Int(0));
}

#[test]
fn scalar_columns_aggregate_without_stale_graph_bindings() {
    let rows = execute("MATCH (n:Memory) WITH n.id AS n RETURN COUNT(n) + COUNT(DISTINCT n) AS total, COLLECT(n) AS values").unwrap();
    assert_eq!(rows[0].values["total"], Value::Int(4));
    assert_eq!(
        rows[0].values["values"],
        Value::List(vec![Value::Int(1), Value::Int(2)])
    );
    let rows = execute("MATCH (n:Missing) RETURN COUNT(n) + 1 AS total").unwrap();
    assert_eq!(rows[0].values["total"], Value::Int(1));
}

#[test]
fn relationship_aliases_are_restored_after_grouping() {
    let rows = execute("MATCH (n:Memory)-[r:MENTIONS]->(m:Memory) WITH r AS edge, COUNT(m) AS seen RETURN id(edge) AS id, type(edge) AS kind, seen").unwrap();
    assert_eq!(rows[0].values["id"], Value::Int(0));
    assert_eq!(rows[0].values["kind"], Value::String("MENTIONS".into()));
    assert_eq!(rows[0].values["seen"], Value::Int(1));
}

#[test]
fn optional_null_import_never_matches_real_node_zero() {
    let rows = execute("MATCH (n:Memory) OPTIONAL MATCH (n)-[:ABSENT]->(m:Memory) WITH m AS kept MATCH (kept) RETURN kept").unwrap();
    assert!(rows.is_empty());
    let rows = execute("MATCH (n:Memory) OPTIONAL MATCH (n)-[:ABSENT]->(m:Memory) WITH m AS kept OPTIONAL MATCH (kept)-[:MENTIONS]->(target:Memory) RETURN COUNT(target) AS count").unwrap();
    assert_eq!(rows[0].values["count"], Value::Int(0));
}

fn unsupported_bounded_match_is_independent_of_input_and_relationship_type(
    configure: impl Fn(&mut GraphMatchStep),
) {
    for rel_type in ["ABSENT", "MENTIONS"] {
        for query in [
            format!("MATCH (n:Memory)-[:{rel_type}*0..1]->(m:Memory) RETURN m"),
            format!("MATCH (seed:Missing) WITH seed AS n MATCH (n)-[:{rel_type}*0..1]->(m:Memory) RETURN m"),
        ] {
            let mut plan =
                lower(hawdb_plan_cypher::plan_pipeline_query(&query, &BTreeMap::new()).unwrap());
            let PhysicalPlan::ProjectExec { input, .. } = &mut plan else {
                panic!("expected projection");
            };
            let PhysicalPlan::GraphMatchExec { program, .. } = input.as_mut() else {
                panic!("expected generic MATCH");
            };
            configure(
                program
                    .steps
                    .iter_mut()
                    .find(|step| matches!(step, GraphMatchStep::Expand { .. }))
                    .expect("expected expansion"),
            );
            with_context(None, |context| {
                let empty = store::ReadFixture::default();
                for (graph, store) in [
                    ("populated", context.store),
                    ("empty", &empty as &dyn crate::store::GraphExecutionRead),
                ] {
                    let context = BatchReadContext { store, ..context };
                    let result = execute_binding_batches(
                        &plan,
                        context,
                        ExecutionLimit::unlimited(),
                        &mut |_| panic!("unsupported bounded MATCH emitted a row"),
                    );
                    assert!(
                        matches!(result, Err(HawDBError::Execution(ref message))
                            if message.contains("bounded MATCH expansion requires an outgoing typed pattern without relationship bindings")),
                        "graph={graph}, query={query}: {result:?}"
                    );
                }
            });
        }
    }
}

#[test]
fn bounded_match_direction_rejection_is_independent_of_relationship_type() {
    for unsupported in [
        RelationshipDirection::Incoming,
        RelationshipDirection::Undirected,
    ] {
        unsupported_bounded_match_is_independent_of_input_and_relationship_type(|step| {
            let GraphMatchStep::Expand { direction, .. } = step else {
                panic!("expected expansion");
            };
            *direction = unsupported;
        });
    }
}

#[test]
fn bounded_match_property_rejection_is_independent_of_relationship_type() {
    unsupported_bounded_match_is_independent_of_input_and_relationship_type(|step| {
        let GraphMatchStep::Expand { properties, .. } = step else {
            panic!("expected expansion");
        };
        properties.insert("weight".into(), Value::Int(10));
    });
}

#[test]
fn bounded_match_relationship_binding_rejection_is_independent_of_relationship_type() {
    unsupported_bounded_match_is_independent_of_input_and_relationship_type(|step| {
        let GraphMatchStep::Expand { relationship, .. } = step else {
            panic!("expected expansion");
        };
        *relationship = Some("edge".into());
    });
}

#[test]
fn bounded_match_untyped_rejection_is_independent_of_input() {
    unsupported_bounded_match_is_independent_of_input_and_relationship_type(|step| {
        let GraphMatchStep::Expand { rel_type, .. } = step else {
            panic!("expected expansion");
        };
        rel_type.clear();
    });
}

#[test]
fn valid_bounded_match_preserves_empty_input() {
    for rel_type in ["ABSENT", "MENTIONS"] {
        let query = format!(
            "MATCH (seed:Missing) WITH seed AS n MATCH (n)-[:{rel_type}*0..1]->(m:Memory) RETURN m"
        );
        assert!(execute(&query).unwrap().is_empty());
        let plan = lower(
            hawdb_plan_cypher::plan_pipeline_query(
                &format!("MATCH (n:Memory)-[:{rel_type}*0..1]->(m:Memory) RETURN m"),
                &BTreeMap::new(),
            )
            .unwrap(),
        );
        with_context(None, |context| {
            let empty = store::ReadFixture::default();
            let context = BatchReadContext {
                store: &empty,
                ..context
            };
            execute_binding_batches(&plan, context, ExecutionLimit::unlimited(), &mut |_| {
                panic!("empty bounded MATCH emitted a row")
            })
            .unwrap();
        });
    }
}

#[test]
fn relationship_uniqueness_is_local_to_each_match_clause() {
    let rows = execute(
        "MATCH (a:Memory)-[:MENTIONS]-(b:Memory)-[:MENTIONS]-(c:Memory) RETURN COUNT(a) AS count",
    )
    .unwrap();
    assert_eq!(rows[0].values["count"], Value::Int(0));
    let rows = execute("MATCH (a:Memory)-[:MENTIONS]-(b:Memory) MATCH (b)-[:MENTIONS]-(c:Memory) RETURN COUNT(a) AS count").unwrap();
    assert_eq!(rows[0].values["count"], Value::Int(2));
}

#[test]
fn scalar_property_predicates_do_not_read_shadowed_nodes() {
    let parameters = BTreeMap::from([(
        "replacement".into(),
        Value::Map(BTreeMap::from([("id".into(), Value::Int(99))])),
    )]);
    let plan = lower(
        hawdb_plan_cypher::plan_pipeline_query(
            "MATCH (n:Memory) WITH $replacement AS n WHERE n.id = 99 RETURN n.id AS id",
            &parameters,
        )
        .unwrap(),
    );
    let rows = with_context(None, |context| {
        let mut rows = Vec::new();
        execute_binding_batches(&plan, context, ExecutionLimit::unlimited(), &mut |batch| {
            rows.extend(batch);
            Ok(BatchControl::Continue)
        })
        .unwrap();
        rows
    });
    assert_eq!(rows.len(), 2);
    assert!(rows.iter().all(|row| row.values["id"] == Value::Int(99)));
}

#[test]
fn match_propagates_errors_cancellation_and_downstream_stop() {
    let query = "MATCH (n:Memory) OPTIONAL MATCH (m:Memory) WHERE lower(m.id) = 'x' RETURN m";
    let error = execute(query).unwrap_err();
    assert!(error
        .to_string()
        .contains("LOWER expression requires a string"));
    let plan = lower(
        hawdb_plan_cypher::plan_pipeline_query("MATCH (n:Memory) RETURN n", &BTreeMap::new())
            .unwrap(),
    );
    let token = hawdb_core::RuntimeCancellationToken::new();
    token.cancel();
    let task = RuntimeTaskContext::without_deadline(token);
    with_context(Some(&task), |context| {
        let result =
            execute_binding_batches(&plan, context, ExecutionLimit::unlimited(), &mut |_| {
                panic!("cancelled query emitted")
            });
        assert!(result.unwrap_err().to_string().contains("cancelled"));
    });
    with_context(None, |context| {
        let memory = ExecutionMemoryConfig {
            batch_rows: std::num::NonZeroUsize::new(1).unwrap(),
            ..context.memory.clone()
        };
        let ledger = QueryMemoryLedger::new(memory.query_memory_bytes);
        let context = BatchReadContext {
            memory: &memory,
            memory_ledger: &ledger,
            ..context
        };
        for fail in [false, true] {
            let mut calls = 0;
            let result =
                execute_binding_batches(&plan, context, ExecutionLimit::unlimited(), &mut |_| {
                    calls += 1;
                    if fail {
                        Err(HawDBError::Execution("consumer failure".into()))
                    } else {
                        Ok(BatchControl::Stop)
                    }
                });
            assert_eq!(calls, 1);
            assert_eq!(result.is_err(), fail);
            assert_eq!(ledger.snapshot().used_bytes, 0);
        }
    });
}

#[test]
fn match_respects_query_memory_admission() {
    let plan = lower(
        hawdb_plan_cypher::plan_pipeline_query("MATCH (n:Memory) RETURN n", &BTreeMap::new())
            .unwrap(),
    );
    with_context(None, |context| {
        let memory = ExecutionMemoryConfig {
            query_memory_bytes: std::num::NonZeroUsize::new(1).unwrap(),
            ..context.memory.clone()
        };
        let ledger = QueryMemoryLedger::new(memory.query_memory_bytes);
        let context = BatchReadContext {
            memory: &memory,
            memory_ledger: &ledger,
            ..context
        };
        let result =
            execute_binding_batches(&plan, context, ExecutionLimit::unlimited(), &mut |_| {
                panic!("over-budget query emitted")
            });
        assert!(result.is_err());
        assert_eq!(ledger.snapshot().used_bytes, 0);
    });
}

#[test]
fn visibility_filters_inside_optional_matching_and_survives_with() {
    let query = "MATCH (n:Memory) WITH n AS anchor OPTIONAL MATCH (anchor)-[:MENTIONS]->(m:Memory) RETURN anchor.id AS id, m";
    let logical = hawdb_plan_cypher::plan_pipeline_query(query, &BTreeMap::new()).unwrap();
    let visible = hawdb_plan_cypher::apply_node_visibility_predicates(logical, &|variable| {
        Predicate::PropertyIn {
            variable: variable.to_string(),
            property: "id".to_string(),
            values: vec![Value::Int(1)],
        }
    });
    let rows = execute_logical(visible).unwrap();
    assert_eq!(rows.len(), 1);
    assert_eq!(rows[0].values["id"], Value::Int(1));
    assert_eq!(rows[0].values["m"], Value::Null);
}

#[test]
fn ordering_can_read_unreturned_input_columns_without_exposing_them() {
    for query in [
        "MATCH (n:Memory) RETURN n.id + 1 AS result ORDER BY n.id DESC",
        "MATCH (n:Memory) WITH n.id + 1 AS score WITH score AS result ORDER BY score DESC RETURN result",
    ] {
        let rows = execute(query).unwrap();
        assert_eq!(rows.iter().map(|row| row.values["result"].clone()).collect::<Vec<_>>(), vec![Value::Int(3), Value::Int(2)]);
        assert!(rows.iter().all(|row| row.values.len() == 1));
    }
    let rows = execute("MATCH (n:Memory) RETURN n.id AS id, COUNT(n) AS count ORDER BY CASE WHEN n.id = 1 THEN 3 ELSE 0 END DESC").unwrap();
    assert_eq!(rows[0].values["id"], Value::Int(1));
    assert!(rows
        .iter()
        .all(|row| row.values.len() == 2 && row.values["count"] == Value::Int(1)));
    assert!(execute("MATCH (n:Memory) RETURN DISTINCT 1 AS result ORDER BY n.id").is_err());
}

#[test]
fn repeated_default_projection_names_preserve_each_result_value() {
    let rows =
        execute("MATCH (n:Memory) RETURN COALESCE(n.id, 0), COALESCE(n.absent, 0) ORDER BY n.id")
            .unwrap();
    assert_eq!(rows.len(), 2);
    for (index, row) in rows.iter().enumerate() {
        assert_eq!(row.values.len(), 2);
        assert_eq!(row.values["coalesce"], Value::Int(index as i64 + 1));
        assert_eq!(row.values["coalesce#2"], Value::Int(0));
    }
}

#[test]
fn match_expansion_checks_cancellation_before_rejecting_target_labels() {
    use crate::traversal::{visit_one_hop_relationships_with_context, OneHopRelationshipSpec};
    use hawdb_storage::{adjacency::AdjacencyDirection, RelRecord};
    with_context(None, |context| {
        let token = hawdb_core::RuntimeCancellationToken::new();
        let task = RuntimeTaskContext::without_deadline(token.clone());
        let mut relationships: Vec<RelRecord> = Vec::new();
        context
            .store
            .visit_adjacent_relationships_owned(
                NodeId(0),
                None,
                AdjacencyDirection::Outgoing,
                &mut |edge| {
                    relationships.push(edge);
                    Ok(ScanControl::Continue)
                },
            )
            .unwrap();
        let store = super::store::ReadFixture {
            nodes: vec![context.store.node_owned(NodeId(1)).unwrap().unwrap()],
            relationships,
            adjacency_cancellation: Some(token),
            ..super::store::ReadFixture::default()
        };
        let result = visit_one_hop_relationships_with_context(
            &store,
            OneHopRelationshipSpec {
                source: NodeId(0),
                rel_type_id: None,
                target_label_ids: Some(&[]),
                rel_properties: &BTreeMap::new(),
                relationship_scan_filter: None,
                direction: RelationshipDirection::Outgoing,
            },
            crate::store::AdjacencyReadMemory {
                budget_bytes: 4096,
                account: None,
            },
            context.observer,
            Some(&task),
            &mut |_, _| panic!("a filtered target must never emit"),
        );
        assert!(result.unwrap_err().to_string().contains("cancelled"));
    });
}

#[test]
fn consecutive_optional_clauses_preserve_cartesian_multiplicity() {
    with_context(None, |context| {
        let mut catalog = Catalog::default();
        let label = catalog.get_or_create_label("Item");
        let outgoing = catalog.get_or_create_rel_type("OUT");
        let incoming = catalog.get_or_create_rel_type("IN");
        let store = super::store::ReadFixture {
            nodes: (0..6)
                .map(|id| NodeRecord {
                    id: NodeId(id),
                    labels: [label].into_iter().collect(),
                    properties: BTreeMap::from([("id".into(), Value::Int(id as i64))]),
                })
                .collect(),
            relationships: (1..6)
                .map(|id| hawdb_storage::RelRecord {
                    id: hawdb_storage::RelId(id),
                    source: NodeId(if id <= 2 { 0 } else { id }),
                    target: NodeId(if id <= 2 { id } else { 0 }),
                    rel_type: if id <= 2 { outgoing } else { incoming },
                    properties: BTreeMap::new(),
                })
                .collect(),
            ..super::store::ReadFixture::default()
        };
        let context = BatchReadContext {
            catalog: &catalog,
            store: &store,
            ..context
        };
        let run = |plan: &PhysicalPlan| {
            let mut rows = Vec::new();
            execute_binding_batches(plan, context, ExecutionLimit::unlimited(), &mut |batch| {
                rows.extend(batch);
                Ok(BatchControl::Continue)
            })
            .unwrap();
            assert_eq!(rows.len(), 1);
            rows[0].values["total"].clone()
        };
        for (distinct, expected) in [("", 12), ("DISTINCT ", 5)] {
            let query = format!("MATCH (n:Item {{id: 0}}) OPTIONAL MATCH (n)-[left:OUT]->() OPTIONAL MATCH (n)<-[right:IN]-() RETURN COUNT({distinct}left) + COUNT({distinct}right) AS total");
            let plan =
                lower(hawdb_plan_cypher::plan_pipeline_query(&query, &BTreeMap::new()).unwrap());
            assert_eq!(run(&plan), Value::Int(expected));
        }
    });
}

#[test]
fn optional_degree_preserves_nullable_graph_match_rows_without_traversing_zero() {
    for (source, expected) in [("m", [0, 0]), ("n", [1, 0])] {
        let input = lower(
            hawdb_plan_cypher::plan_pipeline_query(
                "MATCH (n:Memory) OPTIONAL MATCH (n)-[:ABSENT]->(m:Memory) RETURN n.id AS id, m",
                &BTreeMap::new(),
            )
            .unwrap(),
        );
        let plan = PhysicalPlan::OptionalDegreeExec {
            source_variable: source.into(),
            rel_type: "MENTIONS".into(),
            rel_properties: BTreeMap::new(),
            direction: hawdb_core::RelationshipDirection::Outgoing,
            target_label: "Memory".into(),
            target_properties: BTreeMap::new(),
            alias: "links".into(),
            input: Box::new(input),
        };
        let mut rows = with_context(None, |context| {
            let mut rows = Vec::new();
            execute_binding_batches(&plan, context, ExecutionLimit::unlimited(), &mut |batch| {
                rows.extend(batch);
                Ok(BatchControl::Continue)
            })?;
            Ok::<_, HawDBError>(rows)
        })
        .unwrap();
        rows.sort_by_key(|row| row.values["id"].clone());
        assert_eq!(rows.len(), 2);
        for (index, row) in rows.iter().enumerate() {
            assert_eq!(row.values["id"], Value::Int(index as i64 + 1));
            assert_eq!(row.values["m"], Value::Null);
            assert!(!row.nodes.contains_key("m"));
            assert_eq!(row.values["links"], Value::Int(expected[index]));
        }
    }
}

#[test]
fn native_expand_preserves_bound_relationship_identity() {
    with_context(None, |context| {
        let rel_type = context.catalog.rel_type_id("MENTIONS").unwrap();
        let store = store::ReadFixture {
            nodes: (0..2)
                .map(|id| context.store.node_owned(NodeId(id)).unwrap().unwrap())
                .collect(),
            relationships: (0..2)
                .map(|id| hawdb_storage::RelRecord {
                    id: hawdb_storage::RelId(id),
                    source: NodeId(0),
                    target: NodeId(1),
                    rel_type,
                    properties: BTreeMap::new(),
                })
                .collect(),
            ..store::ReadFixture::default()
        };
        let context = BatchReadContext {
            store: &store,
            ..context
        };
        for (rel_type, label, optional) in [
            ("MENTIONS", "Memory", true),
            ("MENTIONS", "Missing", true),
            ("ABSENT", "Memory", true),
            ("MENTIONS", "Missing", false),
        ] {
            let input = lower(
                hawdb_plan_cypher::plan_pipeline_query(
                    "MATCH (n:Memory)-[edge:MENTIONS]->(old:Memory) RETURN n, edge",
                    &BTreeMap::new(),
                )
                .unwrap(),
            );
            let plan = PhysicalPlan::AdjacencyExpandExec {
                source_variable: "n".into(),
                source_label: "Memory".into(),
                rel_variable: Some("edge".into()),
                rel_type: rel_type.into(),
                rel_properties: BTreeMap::new(),
                direction: hawdb_core::RelationshipDirection::Outgoing,
                target_variable: "candidate".into(),
                target_label: label.into(),
                min_hops: 1,
                max_hops: 1,
                optional,
                graph_budget: None,
                input: Box::new(input),
            };
            let mut rows = Vec::new();
            execute_binding_batches(&plan, context, ExecutionLimit::unlimited(), &mut |batch| {
                rows.extend(batch);
                Ok(BatchControl::Continue)
            })
            .unwrap();
            if !optional {
                assert!(rows.is_empty());
                continue;
            }
            assert_eq!(rows.len(), 2, "{rel_type}, {label}");
            rows.sort_by_key(|row| row.relationships["edge"].id);
            for (id, row) in rows.iter().enumerate() {
                assert_eq!(
                    row.relationships["edge"].id,
                    hawdb_storage::RelId(id as u64)
                );
                if rel_type == "MENTIONS" && label == "Memory" {
                    assert_eq!(row.nodes["candidate"].id, NodeId(1));
                } else {
                    assert_eq!(row.values["candidate"], Value::Null);
                    assert!(!row.nodes.contains_key("candidate"));
                }
            }
        }
    });
}

// These exercise the actual query owners. The fixture cancels at the source
// allocation boundary, after each caller's entry checkpoint but before copying.
fn live_adjacency_guard_plan(owner: usize) -> PhysicalPlan {
    let input = || PhysicalPlan::SeqNodeScan {
        variable: "n".into(),
        label: "Memory".into(),
    };
    let bound_input = || {
        lower(
            hawdb_plan_cypher::plan_pipeline_query(
                "MATCH (n:Memory), (m:Memory) RETURN n, m",
                &BTreeMap::new(),
            )
            .unwrap(),
        )
    };
    match owner {
        0 => PhysicalPlan::OptionalDegreeExec {
            source_variable: "n".into(),
            rel_type: "MENTIONS".into(),
            rel_properties: BTreeMap::new(),
            direction: RelationshipDirection::Outgoing,
            target_label: "Memory".into(),
            target_properties: BTreeMap::new(),
            alias: "degree".into(),
            input: Box::new(input()),
        },
        1 | 2 => PhysicalPlan::AdjacencyExpandExec {
            source_variable: "n".into(),
            source_label: "Memory".into(),
            rel_variable: (owner == 1).then(|| "r".into()),
            rel_type: "MENTIONS".into(),
            rel_properties: BTreeMap::new(),
            direction: RelationshipDirection::Outgoing,
            target_variable: "m".into(),
            target_label: "Memory".into(),
            min_hops: 1,
            max_hops: if owner == 2 { 2 } else { 1 },
            optional: false,
            graph_budget: None,
            input: Box::new(input()),
        },
        3 => PhysicalPlan::FilterExec {
            predicate: Predicate::RelationshipExists {
                variable: "n".into(),
                rel_type: "MENTIONS".into(),
                direction: RelationshipDirection::Outgoing,
                target_label: "Memory".into(),
            },
            input: Box::new(input()),
        },
        4 => PhysicalPlan::FilterExec {
            predicate: Predicate::BoundRelationshipExists {
                source_variable: "n".into(),
                rel_type: "MENTIONS".into(),
                direction: RelationshipDirection::Outgoing,
                target_variable: "m".into(),
            },
            input: Box::new(bound_input()),
        },
        5 => PhysicalPlan::AdjacencyExistsExec {
            source_variable: "n".into(),
            rel_type: "MENTIONS".into(),
            direction: RelationshipDirection::Outgoing,
            target_variable: "m".into(),
            input: Box::new(bound_input()),
        },
        _ => unreachable!(),
    }
}

fn assert_live_adjacency_cancel_before_copy(owner: usize) {
    with_context(None, |context| {
        let token = hawdb_core::RuntimeCancellationToken::new();
        let task = RuntimeTaskContext::without_deadline(token.clone());
        let store = store::ReadFixture {
            nodes: (0..2)
                .map(|id| context.store.node_owned(NodeId(id)).unwrap().unwrap())
                .collect(),
            relationships: vec![hawdb_storage::RelRecord {
                id: hawdb_storage::RelId(0),
                source: NodeId(0),
                target: NodeId(1),
                rel_type: context.catalog.rel_type_id("MENTIONS").unwrap(),
                properties: BTreeMap::from([("body".into(), Value::String("payload".repeat(256)))]),
            }],
            adjacency_cancellation: Some(token),
            ..store::ReadFixture::default()
        };
        let context = BatchReadContext {
            store: &store,
            task_context: Some(&task),
            ..context
        };
        let result = execute_binding_batches(
            &live_adjacency_guard_plan(owner),
            context,
            ExecutionLimit::unlimited(),
            &mut |_| panic!("cancelled adjacency emitted a row"),
        );
        assert!(result.unwrap_err().to_string().contains("cancelled"));
        assert_eq!(
            store.relationship_copies.get(),
            0,
            "query owner {owner} copied after cancellation"
        );
    });
}

#[test]
fn optional_degree_live_adjacency_cancels_before_copy() {
    assert_live_adjacency_cancel_before_copy(0);
}
#[test]
fn one_hop_expand_live_adjacency_cancels_before_copy() {
    assert_live_adjacency_cancel_before_copy(1);
}
#[test]
fn bounded_expand_live_adjacency_cancels_before_copy() {
    assert_live_adjacency_cancel_before_copy(2);
}
#[test]
fn relationship_predicate_live_adjacency_cancels_before_copy() {
    assert_live_adjacency_cancel_before_copy(3);
}
#[test]
fn bound_predicate_live_adjacency_cancels_before_copy() {
    assert_live_adjacency_cancel_before_copy(4);
}
#[test]
fn adjacency_exists_live_adjacency_cancels_before_copy() {
    assert_live_adjacency_cancel_before_copy(5);
}

fn assert_exists_source_budget_before_copy(owner: usize) {
    with_context(None, |context| {
        let store = store::ReadFixture {
            nodes: (0..2)
                .map(|id| context.store.node_owned(NodeId(id)).unwrap().unwrap())
                .collect(),
            relationships: vec![hawdb_storage::RelRecord {
                id: hawdb_storage::RelId(0),
                source: NodeId(0),
                target: NodeId(1),
                rel_type: context.catalog.rel_type_id("MENTIONS").unwrap(),
                properties: BTreeMap::from([("body".into(), Value::String("x".repeat(65536)))]),
            }],
            ..store::ReadFixture::default()
        };
        let mut memory = context.memory.clone();
        memory.query_memory_bytes = std::num::NonZeroUsize::new(32768).unwrap();
        let ledger = QueryMemoryLedger::new(memory.query_memory_bytes);
        let context = BatchReadContext {
            store: &store,
            memory: &memory,
            memory_ledger: &ledger,
            ..context
        };
        let emitted = Cell::new(0);
        let result = execute_binding_batches(
            &live_adjacency_guard_plan(owner),
            context,
            ExecutionLimit::unlimited(),
            &mut |batch| {
                emitted.set(emitted.get() + batch.len());
                Ok(BatchControl::Continue)
            },
        );
        assert!(result.unwrap_err().to_string().contains("exceeding"));
        assert_eq!(emitted.get(), 0);
        assert_eq!(store.relationship_copies.get(), 0);
    });
}
#[test]
fn bound_predicate_exists_source_budget_before_copy() {
    assert_exists_source_budget_before_copy(4);
}
#[test]
fn adjacency_exists_source_budget_before_copy() {
    assert_exists_source_budget_before_copy(5);
}

#[test]
fn optional_seed_null_hop_uses_final_batch_row_size() {
    use hawdb_core::graph_rag::{
        MissingScoringFeature, ScoreFeature, ScoringCombination, ScoringProgram, ScoringSpec,
        ScoringTerm,
    };
    use hawdb_core::{ResourceBudgetKind, ResourceBudgetScope};
    use hawdb_plan_cypher::ScoringSeedGraphInput;
    use std::num::NonZeroUsize;
    with_context(None, |context| {
        let prefix = PhysicalPlan::AdjacencyExpandExec {
            source_variable: "seed".into(),
            source_label: "Memory".into(),
            rel_variable: None,
            rel_type: "Unknown".into(),
            rel_properties: BTreeMap::new(),
            direction: RelationshipDirection::Outgoing,
            target_variable: "candidate".into(),
            target_label: "Memory".into(),
            min_hops: 1,
            max_hops: 1,
            optional: true,
            graph_budget: None,
            input: Box::new(PhysicalPlan::ProjectExec {
                items: vec![Projection {
                    name: "pad".into(),
                    expression: ProjectionExpression::Literal(Value::String("x".repeat(16384))),
                }],
                input: Box::new(PhysicalPlan::GraphSeedScan {
                    query_parameter: "text".into(),
                    label: "Memory".into(),
                    variable: "seed".into(),
                    score_column: "score".into(),
                    top_k: 1,
                    node_visibility_predicate: None,
                }),
            }),
        };
        let seed_input = ScoringSeedGraphInput::new_graph("seed", "candidate").unwrap();
        seed_input.validate_plan(&prefix).unwrap();
        let declared = PhysicalPlan::ScoringProgramExec {
            score_column: "score".into(),
            seed_graph_input: Some(seed_input),
            program: ScoringProgram::new(
                ScoringCombination::WeightedSum,
                MissingScoringFeature::Reject,
                ScoringSpec {
                    terms: vec![ScoringTerm {
                        feature: ScoreFeature::GraphSeedScore,
                        weight: 1.0,
                    }],
                    decay: vec![],
                },
            )
            .unwrap(),
            reference_time_millis: 0,
            limit: 1,
            input: Box::new(prefix),
        };
        let PhysicalPlan::ScoringProgramExec { input, .. } = &declared else {
            unreachable!()
        };
        let parameters = BTreeMap::from([("text".into(), Value::String("1".into()))]);
        // Execute the actual certified producer/project/OPTIONAL prefix with
        // its declared scorer observer. Outer scoring scratch is a separate cap.
        let observer = QueryExecutionObserver::new(&declared);
        let mut expected = Vec::new();
        execute_binding_batches(
            input,
            BatchReadContext {
                parameters: &parameters,
                observer: &observer,
                ..context
            },
            ExecutionLimit::unlimited(),
            &mut |rows| {
                expected.extend(rows);
                Ok(BatchControl::Continue)
            },
        )
        .unwrap();
        assert_eq!(expected.len(), 1);
        let expected = expected.pop().unwrap();
        assert_eq!(expected.values["candidate"], Value::Null);
        assert_eq!(expected.values["\0hawdb.scoring.hops"], Value::Null);
        let bytes = crate::binding::binding_memory_bytes(&expected);
        for cap in [bytes, bytes - 1] {
            let memory = ExecutionMemoryConfig {
                batch_payload_bytes: NonZeroUsize::new(cap).unwrap(),
                ..context.memory.clone()
            };
            let ledger = QueryMemoryLedger::new(memory.query_memory_bytes);
            let observer = QueryExecutionObserver::new(&declared);
            let mut actual = Vec::new();
            let result = execute_binding_batches(
                input,
                BatchReadContext {
                    parameters: &parameters,
                    observer: &observer,
                    memory: &memory,
                    memory_ledger: &ledger,
                    ..context
                },
                ExecutionLimit::unlimited(),
                &mut |rows| {
                    actual.extend(rows);
                    Ok(BatchControl::Continue)
                },
            );
            if cap == bytes {
                result.expect("final Null-hop row must fit its exact batch cap");
                assert_eq!(actual, vec![expected.clone()]);
            } else {
                let error = result.expect_err("one byte below final row must refuse");
                let HawDBError::ResourceBudgetExceeded(cause) = error else {
                    panic!("expected a typed final-row budget refusal, got {error}");
                };
                assert_eq!(cause.resource, ResourceBudgetKind::MemoryBytes);
                assert_eq!(cause.scope, ResourceBudgetScope::Operator);
                assert_eq!(cause.owner, "AdjacencyExpandExec output");
                assert_eq!(cause.memory_class, Some("pipeline_batch"));
                assert_eq!(cause.requested, bytes as u64);
                assert_eq!(cause.used, 0);
                assert_eq!(cause.budget, cap as u64);
                assert_eq!(
                    cause.budget_option,
                    "ExecutionMemoryConfig.batch_payload_bytes"
                );
                assert!(actual.is_empty());
            }
            assert_eq!(ledger.snapshot().used_bytes, 0);
        }
    });
}