linkrs-query 0.1.0

use super::*;

use crate::planning::plan::core::node_id_generator::next_node_id;
use crate::planning::plan::logical::logical_nodes::access::LogicalScanVerticesNode;
use crate::planning::plan::logical::logical_nodes::control_flow::{
    LogicalLoopNode, LogicalSelectNode,
};
use crate::planning::plan::logical::logical_nodes::flatten::LogicalFlattenNode;
use crate::planning::plan::logical::logical_nodes::graph_ops::LogicalAssignNode;
use crate::planning::plan::logical::logical_nodes::join::LogicalSemiJoinNode;
use linkrs_core::types::expr::contextual::ContextualExpression;
use linkrs_core::types::expr::expression_context::ExpressionAnalysisContext;
use linkrs_core::types::expr::ExpressionMeta;
use std::collections::HashMap;
use std::sync::Arc;

fn test_id(n: u64) -> ExpressionId {
    ExpressionId::new(n)
}

fn scan() -> LogicalNodeEnum {
    LogicalNodeEnum::ScanVertices(LogicalScanVerticesNode {
        id: next_node_id(),
        space_id: 1,
        space_name: "test".to_string(),
        tag: Some("p".to_string()),
        expression: None,
        limit: None,
        projected_properties: vec![],
        index_hint: None,
        estimated_cardinality: None,
        output_var: None,
        col_names: vec!["a.name".to_string()],
        column_types: vec![],
    })
}

fn scan_with_expr() -> (LogicalNodeEnum, Arc<ExpressionAnalysisContext>) {
    let raw_ctx = ExpressionAnalysisContext::new();
    let ctx = Arc::new(raw_ctx);
    let expr = linkrs_core::Expression::Variable("a".to_string());
    let meta = ExpressionMeta::new(expr);
    let id = ctx.register_expression(meta);
    let ctx_expr = ContextualExpression::new(id, ctx.clone());
    let node = LogicalNodeEnum::ScanVertices(LogicalScanVerticesNode {
        id: next_node_id(),
        space_id: 1,
        space_name: "test".to_string(),
        tag: Some("p".to_string()),
        expression: Some(ctx_expr),
        limit: None,
        projected_properties: vec![],
        index_hint: None,
        estimated_cardinality: None,
        output_var: None,
        col_names: vec!["a.name".to_string()],
        column_types: vec![],
    });
    (node, ctx)
}

#[test]
fn scan_schema_is_flat() {
    let mut n = scan();
    let s = n.compute_factorized_schema(&[]).unwrap();
    assert!(s.is_flat_schema());
    assert_eq!(s.num_groups(), 1);
    let flat = n.compute_flat_schema(&[]).unwrap();
    assert!(flat.is_flat_schema());
}

#[test]
fn scan_schema_with_real_expr() {
    let (mut n, _ctx) = scan_with_expr();
    let s = n.compute_factorized_schema(&[]).unwrap();
    assert!(s.is_flat_schema());
    assert_eq!(s.num_groups(), 1);
}

#[test]
fn project_dependency() {
    let raw_ctx = ExpressionAnalysisContext::new();
    let ctx = Arc::new(raw_ctx);
    let expr = linkrs_core::Expression::Variable("a".to_string());
    let meta = ExpressionMeta::new(expr);
    let id_a = ctx.register_expression(meta);
    let ctx_a = ContextualExpression::new(id_a.clone(), ctx.clone());
    let mut scan_schema = FactorizedSchema::new();
    let g0 = scan_schema.create_flat_group(false);
    scan_schema
        .insert_to_group_and_scope(id_a.clone(), g0)
        .unwrap();
    scan_schema
        .insert_to_group_and_scope_with_name(test_id(999), Some("extra".to_string()), g0)
        .unwrap();

    let col_expr = ctx_a.clone();
    let yield_col = linkrs_core::YieldColumn {
        expression: col_expr,
        alias: "a2".to_string(),
    };
    let mut proj = LogicalNodeEnum::Project(
        crate::planning::plan::logical::logical_nodes::operation::LogicalProjectNode {
            id: next_node_id(),
            input: Some(Box::new(scan())),
            columns: vec![yield_col],
            subqueries: Vec::new(),
            has_folded_expressions: false,
            output_var: None,
            col_names: vec!["a2".to_string()],
            column_types: vec![],
        },
    );
    let out = proj.compute_factorized_schema(&[scan_schema]).unwrap();
    assert!(out.is_expression_in_scope(&id_a) || out.get_group_pos_by_name("a2").is_some());
}

#[test]
fn filter_flatten() {
    let raw_ctx = ExpressionAnalysisContext::new();
    let ctx = Arc::new(raw_ctx);
    let expr = linkrs_core::Expression::Variable("b".to_string());
    let meta = ExpressionMeta::new(expr);
    let id_b = ctx.register_expression(meta);
    let ctx_b = ContextualExpression::new(id_b.clone(), ctx.clone());
    let mut child_schema = FactorizedSchema::new();
    let g0 = child_schema.create_flat_group(false);
    let g1 = child_schema.create_group();
    child_schema
        .insert_to_group_and_scope(test_id(100), g0)
        .unwrap();
    child_schema
        .insert_to_group_and_scope(id_b.clone(), g1)
        .unwrap();
    let filter = LogicalNodeEnum::Filter(
        crate::planning::plan::logical::logical_nodes::operation::LogicalFilterNode {
            id: next_node_id(),
            input: Some(Box::new(scan())),
            condition: ctx_b,
            output_var: None,
            col_names: vec![],
            column_types: vec![],
        },
    );
    let mut filter = filter;
    let out = filter.compute_factorized_schema(&[child_schema]).unwrap();
    // Single unflat group stays factorized under FlattenAllButOne.
    out.validate_at_most_one_unflat().unwrap();
    assert_eq!(out.unflat_group_pos(), Some(g1));
}

#[test]
fn aggregate_keys() {
    let raw_ctx = ExpressionAnalysisContext::new();
    let ctx = Arc::new(raw_ctx);
    let expr = linkrs_core::Expression::Variable("a".to_string());
    let meta = ExpressionMeta::new(expr);
    let id = ctx.register_expression(meta);
    let ctx_a = ContextualExpression::new(id.clone(), ctx.clone());
    let mut child_schema = FactorizedSchema::new();
    let g0 = child_schema.create_flat_group(false);
    child_schema
        .insert_to_group_and_scope(id.clone(), g0)
        .unwrap();
    let mut agg = LogicalNodeEnum::Aggregate(
        crate::planning::plan::logical::logical_nodes::operation::LogicalAggregateNode {
            id: next_node_id(),
            input: Some(Box::new(scan())),
            group_key_exprs: vec![ctx_a],
            aggregation_functions: vec![],
            aggregation_args: vec![],
            aggregation_distinct: vec![],
            aggregation_filters: vec![],
            grouping_sets: vec![],
            output_var: None,
            col_names: vec![],
            column_types: vec![],
        },
    );
    let out = agg.compute_factorized_schema(&[child_schema]).unwrap();
    assert!(out.is_flat_schema());
    assert!(out.num_groups() >= 1);
}

#[test]
fn flatten_schema() {
    let mut scan_n = scan();
    let scan_schema = scan_n.compute_factorized_schema(&[]).unwrap();
    let mut g = scan_schema.clone();
    let pos = g.create_group();
    g.insert_to_group_and_scope(test_id(555), pos).unwrap();
    assert!(!g.get_group(pos).expect("g").is_flat());
    let mut flatten = LogicalNodeEnum::Flatten(LogicalFlattenNode::new(pos, scan()));
    let out = flatten.compute_factorized_schema(&[g]).unwrap();
    assert!(out.is_flat_schema());
}

#[test]
fn flatten_already_flat_is_noop() {
    let mut scan_n = scan();
    let schema = scan_n.compute_factorized_schema(&[]).unwrap();
    assert!(schema.is_flat_schema());
    let mut flatten = LogicalNodeEnum::Flatten(LogicalFlattenNode::new(0, scan()));
    let out = flatten
        .compute_factorized_schema(std::slice::from_ref(&schema))
        .unwrap();
    assert!(out.is_flat_schema());
    assert_eq!(out.num_groups(), schema.num_groups());
}

#[test]
fn flatten_out_of_range_reports() {
    let mut scan_n = scan();
    let schema = scan_n.compute_factorized_schema(&[]).unwrap();
    let mut flatten = LogicalNodeEnum::Flatten(LogicalFlattenNode::new(99, scan()));
    let err = flatten
        .compute_factorized_schema(&[schema])
        .expect_err("out-of-range flatten must fail");
    assert_eq!(
        err,
        crate::planning::plan::factorization::FactorizationError::GroupPosOutOfRange(99)
    );
    assert!(err.to_string().contains("out of range"));
}

#[test]
fn join_merges() {
    let mut left = scan();
    let ls = left.compute_factorized_schema(&[]).unwrap();
    let mut right = scan();
    let mut rs = right.compute_factorized_schema(&[]).unwrap();
    let pos = rs.create_group();
    rs.insert_to_group_and_scope(test_id(777), pos).unwrap();
    let mut join = LogicalNodeEnum::InnerJoin(
        crate::planning::plan::logical::logical_nodes::join::LogicalInnerJoinNode {
            id: next_node_id(),
            left: Box::new(scan()),
            right: Box::new(scan()),
            hash_keys: vec![],
            probe_keys: vec![],
            recommended_algorithm: None,
            output_var: None,
            col_names: vec![],
            column_types: vec![],
        },
    );
    let out = join.compute_factorized_schema(&[ls, rs]).unwrap();
    out.validate_at_most_one_unflat().unwrap();
}

#[test]
fn fulltext_leaf_is_flat_single_group() {
    use crate::parser::ast::fulltext::FulltextQueryExpr;
    use crate::planning::plan::logical::logical_nodes::search::LogicalFulltextSearchNode;
    let mut node = LogicalNodeEnum::FulltextSearch(LogicalFulltextSearchNode {
        id: next_node_id(),
        index_name: "idx".to_string(),
        query: FulltextQueryExpr::Simple("test".to_string()),
        yield_clause: None,
        where_clause: None,
        order_clause: None,
        limit: None,
        offset: None,
        space_id: 1,
        tag_name: "person".to_string(),
        field_name: "name".to_string(),
        output_var: None,
        col_names: vec!["a".to_string()],
        column_types: vec![],
    });
    let schema = node.compute_factorized_schema(&[]).unwrap();
    assert_eq!(schema.num_groups(), 1);
    assert!(schema.is_flat_schema());
}

#[test]
fn unwind_passthrough() {
    let mut child_schema = FactorizedSchema::new();
    let g0 = child_schema.create_flat_group(false);
    let g1 = child_schema.create_group();
    child_schema
        .insert_to_group_and_scope(test_id(1), g0)
        .unwrap();
    child_schema
        .insert_to_group_and_scope(test_id(2), g1)
        .unwrap();
    let mut unwind = LogicalNodeEnum::Unwind(
        crate::planning::plan::logical::logical_nodes::graph_ops::LogicalUnwindNode {
            id: next_node_id(),
            input: Some(Box::new(scan())),
            alias: "x".to_string(),
            list_expression: {
                let raw_ctx = ExpressionAnalysisContext::new();
                let ctx = Arc::new(raw_ctx);
                let meta =
                    ExpressionMeta::new(linkrs_core::Expression::Variable("list".to_string()));
                let id = ctx.register_expression(meta);
                ContextualExpression::new(id, ctx)
            },
            output_var: None,
            col_names: vec!["x".to_string()],
            column_types: vec![],
        },
    );
    let out = unwind
        .compute_factorized_schema(&[child_schema.clone()])
        .unwrap();
    // Unresolved list flattens conservatively, then builds a fresh
    // unflat group for the alias (baseline恒建新组).
    out.validate_at_most_one_unflat().unwrap();
    assert!(out.has_unflat_group());
    assert_eq!(out.get_group_pos_by_name_opt("x"), out.unflat_group_pos());
}

#[test]
fn get_neighbors_chain_keeps_one_unflat() {
    let mut base = FactorizedSchema::new();
    let g0 = base.create_flat_group(false);
    base.insert_to_group_and_scope(test_id(1), g0).unwrap();
    let mut prev = base;
    for _ in 0..3 {
        let mut node = LogicalNodeEnum::GetNeighbors(
            crate::planning::plan::logical::logical_nodes::access::LogicalGetNeighborsNode {
                id: next_node_id(),
                space_id: 1,
                tag: None,
                src_vids: "1".to_string(),
                edge_types: vec!["knows".to_string()],
                direction: "OUT".to_string(),
                edge_props: vec![],
                tag_props: vec![],
                expression: None,
                dedup: false,
                limit: None,
                projected_properties: vec![],
                index_hint: None,
                estimated_cardinality: None,
                output_var: None,
                col_names: vec!["b".to_string()],
                column_types: vec![],
                deps: vec![scan()],
            },
        );
        let next = node.compute_factorized_schema(&[prev.clone()]).unwrap();
        next.validate_at_most_one_unflat().unwrap();
        prev = next;
    }
    prev.validate_at_most_one_unflat().unwrap();
    assert_eq!(prev.groups().iter().filter(|g| !g.is_flat()).count(), 1);
}

fn unwind_node(alias: &str, list: linkrs_core::Expression) -> (LogicalNodeEnum, ExpressionId) {
    let raw_ctx = ExpressionAnalysisContext::new();
    let ctx = Arc::new(raw_ctx);
    let id = ctx.register_expression(ExpressionMeta::new(list));
    let list_expression = ContextualExpression::new(id.clone(), ctx);
    let node = LogicalNodeEnum::Unwind(
        crate::planning::plan::logical::logical_nodes::graph_ops::LogicalUnwindNode {
            id: next_node_id(),
            input: Some(Box::new(scan())),
            alias: alias.to_string(),
            list_expression,
            output_var: None,
            col_names: vec![alias.to_string()],
            column_types: vec![],
        },
    );
    (node, id)
}

#[test]
fn unwind_list_literal_creates_unflat_group() {
    let child_schema = FactorizedSchema::new();
    let list = linkrs_core::Expression::List(vec![linkrs_core::Expression::Literal(
        linkrs_core::value::Value::Int(1),
    )]);
    let (mut node, list_id) = unwind_node("x", list);
    let out = node.compute_factorized_schema(&[child_schema]).unwrap();
    out.validate_at_most_one_unflat().unwrap();
    assert!(out.has_unflat_group());
    assert_eq!(out.get_group_pos(&list_id), out.unflat_group_pos());
    assert_eq!(out.get_group_pos_by_name("x"), out.unflat_group_pos());
}

#[test]
fn unwind_tracked_column_flattens_first() {
    let mut child_schema = FactorizedSchema::new();
    let g0 = child_schema.create_flat_group(false);
    let g1 = child_schema.create_group();
    child_schema
        .insert_to_group_and_scope(test_id(1), g0)
        .unwrap();
    let raw_ctx = ExpressionAnalysisContext::new();
    let ctx = Arc::new(raw_ctx);
    let list_id = ctx.register_expression(ExpressionMeta::new(linkrs_core::Expression::Variable(
        "items".to_string(),
    )));
    child_schema
        .insert_to_group_and_scope(list_id.clone(), g1)
        .unwrap();
    assert!(child_schema.has_unflat_group());
    let list_expression = ContextualExpression::new(list_id.clone(), ctx);
    let mut node = LogicalNodeEnum::Unwind(
        crate::planning::plan::logical::logical_nodes::graph_ops::LogicalUnwindNode {
            id: next_node_id(),
            input: Some(Box::new(scan())),
            alias: "item".to_string(),
            list_expression,
            output_var: None,
            col_names: vec!["item".to_string()],
            column_types: vec![],
        },
    );
    let out = node.compute_factorized_schema(&[child_schema]).unwrap();
    out.validate_at_most_one_unflat().unwrap();
    // Baseline builds a fresh unflat group for the unwind output: the
    // input list group is flattened in place, the alias lives in a new group.
    assert!(out.has_unflat_group());
    assert!(out.is_expression_in_scope(&list_id));
    assert_eq!(out.get_group_pos(&list_id), Some(g1));
    assert!(out.get_group(g1).is_some_and(|g| g.is_flat()));
    assert_eq!(out.get_group_pos_by_name("item"), out.unflat_group_pos());
    assert_ne!(
        out.get_group_pos_by_name("item"),
        out.get_group_pos(&list_id)
    );
}

#[test]
fn get_neighbors_registers_output_expression() {
    let raw_ctx = ExpressionAnalysisContext::new();
    let ctx = Arc::new(raw_ctx);
    let out_id = ctx.register_expression(ExpressionMeta::new(linkrs_core::Expression::Variable(
        "b".to_string(),
    )));
    let out_expr = ContextualExpression::new(out_id.clone(), ctx);
    let child_schema = FactorizedSchema::new();
    let mut node = LogicalNodeEnum::GetNeighbors(
        crate::planning::plan::logical::logical_nodes::access::LogicalGetNeighborsNode {
            id: next_node_id(),
            space_id: 1,
            tag: None,
            src_vids: "1".to_string(),
            edge_types: vec!["knows".to_string()],
            direction: "OUT".to_string(),
            edge_props: vec![],
            tag_props: vec![],
            expression: Some(out_expr),
            dedup: false,
            limit: None,
            projected_properties: vec![],
            index_hint: None,
            estimated_cardinality: None,
            output_var: None,
            col_names: vec!["b".to_string()],
            column_types: vec![],
            deps: vec![scan()],
        },
    );
    let out = node.compute_factorized_schema(&[child_schema]).unwrap();
    out.validate_at_most_one_unflat().unwrap();
    assert!(out.has_unflat_group());
    assert_eq!(out.get_group_pos(&out_id), out.unflat_group_pos());
}

#[test]
fn get_neighbors_registers_col_names_without_expression() {
    let child_schema = FactorizedSchema::new();
    let mut node = LogicalNodeEnum::GetNeighbors(
        crate::planning::plan::logical::logical_nodes::access::LogicalGetNeighborsNode {
            id: next_node_id(),
            space_id: 1,
            tag: None,
            src_vids: "1".to_string(),
            edge_types: vec!["knows".to_string()],
            direction: "OUT".to_string(),
            edge_props: vec![],
            tag_props: vec![],
            expression: None,
            dedup: false,
            limit: None,
            projected_properties: vec![],
            index_hint: None,
            estimated_cardinality: None,
            output_var: Some("b".to_string()),
            col_names: vec!["b".to_string()],
            column_types: vec![],
            deps: vec![scan()],
        },
    );
    let out = node.compute_factorized_schema(&[child_schema]).unwrap();
    out.validate_at_most_one_unflat().unwrap();
    assert!(out.has_unflat_group());
    let pos = out.unflat_group_pos().expect("unflat");
    assert_eq!(out.get_group_pos_by_name_opt("b"), Some(pos));
}

#[test]
fn traverse_registers_vertex_and_edge_alias() {
    let child_schema = FactorizedSchema::new();
    let mut node = LogicalNodeEnum::Traverse(
        crate::planning::plan::logical::logical_nodes::traversal::LogicalTraverseNode {
            id: next_node_id(),
            space_id: 1,
            start_vids: "1".to_string(),
            end_vids: None,
            edge_types: vec!["knows".to_string()],
            direction: linkrs_core::types::graph_schema::EdgeDirection::Out,
            min_steps: 1,
            max_steps: 2,
            edge_alias: Some("e".to_string()),
            vertex_alias: Some("b".to_string()),
            e_filter: None,
            v_filter: None,
            first_step_filter: None,
            path_semantic: None,
            dst_tag: None,
            input: None,
            output_var: None,
            col_names: vec![],
            column_types: vec![],
        },
    );
    let out = node.compute_factorized_schema(&[child_schema]).unwrap();
    out.validate_at_most_one_unflat().unwrap();
    let pos = out.unflat_group_pos().expect("unflat");
    assert_eq!(out.get_group_pos_by_name_opt("b"), Some(pos));
    assert_eq!(out.get_group_pos_by_name_opt("e"), Some(pos));
}

#[test]
fn expand_all_registers_col_names() {
    let child_schema = FactorizedSchema::new();
    let mut node = LogicalNodeEnum::ExpandAll(
        crate::planning::plan::logical::logical_nodes::traversal::LogicalExpandAllNode {
            id: next_node_id(),
            space_id: 1,
            edge_types: vec!["knows".to_string()],
            direction: "out".to_string(),
            any_edge_type: false,
            step_limit: Some(1),
            step_limits: None,
            join_input: false,
            sample: false,
            edge_props: vec![],
            vertex_props: vec![],
            filter: None,
            src_vids: vec![],
            include_empty_paths: false,
            input_var: Some("a".to_string()),
            path_semantic: None,
            dst_tag: None,
            deps: vec![],
            output_var: None,
            col_names: vec!["e".to_string()],
            column_types: vec![],
        },
    );
    let out = node.compute_factorized_schema(&[child_schema]).unwrap();
    out.validate_at_most_one_unflat().unwrap();
    let pos = out.unflat_group_pos().expect("unflat");
    assert_eq!(out.get_group_pos_by_name_opt("e"), Some(pos));
}

#[test]
fn traverse_output_resolves_downstream_without_flatten_all() {
    use crate::optimizer::factorization::GroupDependencyAnalyzer;
    let child_schema = FactorizedSchema::new();
    let mut node = LogicalNodeEnum::Traverse(
        crate::planning::plan::logical::logical_nodes::traversal::LogicalTraverseNode {
            id: next_node_id(),
            space_id: 1,
            start_vids: "1".to_string(),
            end_vids: None,
            edge_types: vec!["knows".to_string()],
            direction: linkrs_core::types::graph_schema::EdgeDirection::Out,
            min_steps: 1,
            max_steps: 1,
            edge_alias: None,
            vertex_alias: Some("b".to_string()),
            e_filter: None,
            v_filter: None,
            first_step_filter: None,
            path_semantic: None,
            dst_tag: None,
            input: None,
            output_var: None,
            col_names: vec![],
            column_types: vec![],
        },
    );
    let schema = node.compute_factorized_schema(&[child_schema]).unwrap();
    schema.validate_at_most_one_unflat().unwrap();
    let pos = schema.unflat_group_pos().expect("unflat");
    let empty_store: HashMap<ExpressionId, linkrs_core::Expression> = HashMap::new();
    let mut analyzer = GroupDependencyAnalyzer::with_expr_store(&schema, false, &empty_store);
    analyzer.visit_expression(&linkrs_core::Expression::Variable("b".to_string()));
    assert!(!analyzer.has_unresolved());
    assert!(analyzer.dependent_groups().contains(&pos));
    let to_flatten = crate::optimizer::factorization::flatten_resolver::FlattenAllButOne::get_groups_pos_to_flatten_for_groups(
        analyzer.dependent_groups(),
        &schema,
    );
    assert!(to_flatten.is_empty());
}

#[test]
fn bi_expand_merges_both_children() {
    let mut left_schema = FactorizedSchema::new();
    let g0 = left_schema.create_flat_group(false);
    left_schema
        .insert_to_group_and_scope(test_id(1), g0)
        .unwrap();
    let mut right_schema = FactorizedSchema::new();
    let g1 = right_schema.create_flat_group(false);
    right_schema
        .insert_to_group_and_scope(test_id(2), g1)
        .unwrap();
    let mut node = LogicalNodeEnum::BiExpand(
        crate::planning::plan::logical::logical_nodes::traversal::LogicalBiExpandNode {
            id: next_node_id(),
            space_id: 1,
            left_direction: linkrs_core::types::graph_schema::EdgeDirection::Out,
            right_direction: linkrs_core::types::graph_schema::EdgeDirection::In,
            edge_types: vec!["knows".to_string()],
            max_hops: 2,
            meeting_point_var: None,
            dst_tag: None,
            left: Box::new(scan()),
            right: Box::new(scan()),
            output_var: None,
            col_names: vec![],
            column_types: vec![],
        },
    );
    let out = node
        .compute_factorized_schema(&[left_schema, right_schema])
        .unwrap();
    out.validate_at_most_one_unflat().unwrap();
    assert!(out.is_expression_in_scope(&test_id(1)));
    assert!(out.is_expression_in_scope(&test_id(2)));
    assert!(out.has_unflat_group());
}

fn ctx_var(var: &str) -> (ContextualExpression, ExpressionId) {
    let ctx = Arc::new(ExpressionAnalysisContext::new());
    let id = ctx.register_expression(ExpressionMeta::new(linkrs_core::Expression::Variable(
        var.to_string(),
    )));
    (ContextualExpression::new(id.clone(), ctx), id)
}

fn assign_node(rhs: ContextualExpression) -> LogicalNodeEnum {
    LogicalNodeEnum::Assign(LogicalAssignNode {
        id: next_node_id(),
        input: Some(Box::new(scan())),
        assignments: vec![("c".into(), rhs)],
        output_var: None,
        col_names: vec![],
        column_types: vec![],
    })
}

#[test]
fn assign_on_flat_group_keeps_unflat() {
    let (rhs, rhs_id) = ctx_var("a");
    let mut child = FactorizedSchema::new();
    let g0 = child.create_flat_group(false);
    let g1 = child.create_group();
    child.insert_to_group_and_scope(test_id(1), g0).unwrap();
    child.insert_to_group_and_scope(test_id(2), g1).unwrap();
    child
        .insert_to_group_and_scope_with_name(test_id(3), Some("a".to_string()), g0)
        .unwrap();
    let mut node = assign_node(rhs);
    let out = node.compute_factorized_schema(&[child]).unwrap();
    out.validate_at_most_one_unflat().unwrap();
    assert_eq!(
        out.unflat_group_pos(),
        Some(g1),
        "assignment reading only a flat group must not flatten the unflat group"
    );
    assert!(
        out.is_expression_in_scope(&rhs_id),
        "assigned expression must be registered in scope"
    );
}

#[test]
fn assign_on_unflat_group_targets_dependent_group() {
    let (rhs, rhs_id) = ctx_var("a");
    let mut child = FactorizedSchema::new();
    let g0 = child.create_flat_group(false);
    let g1 = child.create_group();
    child.insert_to_group_and_scope(test_id(1), g0).unwrap();
    child.insert_to_group_and_scope(test_id(2), g1).unwrap();
    child
        .insert_to_group_and_scope_with_name(test_id(3), Some("a".to_string()), g1)
        .unwrap();
    let mut node = assign_node(rhs);
    let out = node.compute_factorized_schema(&[child]).unwrap();
    out.validate_at_most_one_unflat().unwrap();
    assert_eq!(
        out.unflat_group_pos(),
        Some(g1),
        "assignment on an unflat group must target that group, not flatten-all"
    );
    assert!(
        out.is_expression_in_scope(&rhs_id),
        "assigned expression must be registered in scope"
    );
}

#[test]
fn semi_join_compute_flattens_probe_keys() {
    let ctx = Arc::new(ExpressionAnalysisContext::new());
    let key_id = ctx.register_expression(ExpressionMeta::new(linkrs_core::Expression::Variable(
        "a".to_string(),
    )));
    let key = ContextualExpression::new(key_id.clone(), ctx);
    let mut left_schema = FactorizedSchema::new();
    let lg0 = left_schema.create_flat_group(false);
    let lg1 = left_schema.create_group();
    left_schema
        .insert_to_group_and_scope(test_id(1), lg0)
        .unwrap();
    left_schema
        .insert_to_group_and_scope(key_id.clone(), lg1)
        .unwrap();
    let mut right_schema = FactorizedSchema::new();
    let rg0 = right_schema.create_flat_group(false);
    let rg1 = right_schema.create_group();
    right_schema
        .insert_to_group_and_scope(test_id(2), rg0)
        .unwrap();
    right_schema
        .insert_to_group_and_scope(test_id(3), rg1)
        .unwrap();
    let mut node = LogicalNodeEnum::SemiJoin(LogicalSemiJoinNode {
        id: next_node_id(),
        left: Box::new(scan()),
        right: Box::new(scan()),
        hash_keys: vec![key],
        probe_keys: vec![],
        join_condition: None,
        anti: false,
        output_var: None,
        col_names: vec![],
        column_types: vec![],
    });
    let out = node
        .compute_factorized_schema(&[left_schema, right_schema])
        .unwrap();
    out.validate_at_most_one_unflat().unwrap();
    // Key-aware rule (same as the rewriter): the unflat probe key group
    // is flattened before the membership test, so the surviving unflat
    // group is the right non-key group and the key itself ends up flat.
    assert!(out.get_group(lg1).is_some_and(|g| g.is_flat()));
    assert!(out
        .get_group_pos(&key_id)
        .is_some_and(|pos| { out.get_group(pos).is_some_and(|g| g.is_flat()) }));
}

#[test]
fn right_join_compute_flattens_build_keys() {
    use crate::planning::plan::logical::logical_nodes::join::LogicalRightJoinNode;
    let ctx = Arc::new(ExpressionAnalysisContext::new());
    let key_id = ctx.register_expression(ExpressionMeta::new(linkrs_core::Expression::Variable(
        "k".to_string(),
    )));
    let key = ContextualExpression::new(key_id.clone(), ctx);
    let mut left_schema = FactorizedSchema::new();
    let lg0 = left_schema.create_flat_group(false);
    let lg1 = left_schema.create_group();
    left_schema
        .insert_to_group_and_scope(test_id(1), lg0)
        .unwrap();
    left_schema
        .insert_to_group_and_scope(test_id(11), lg1)
        .unwrap();
    let mut right_schema = FactorizedSchema::new();
    let rg0 = right_schema.create_flat_group(false);
    let rg1 = right_schema.create_group();
    right_schema
        .insert_to_group_and_scope(test_id(2), rg0)
        .unwrap();
    right_schema
        .insert_to_group_and_scope(key_id.clone(), rg1)
        .unwrap();
    let mut node = LogicalNodeEnum::RightJoin(LogicalRightJoinNode {
        id: next_node_id(),
        left: Box::new(scan()),
        right: Box::new(scan()),
        hash_keys: vec![],
        probe_keys: vec![key],
        output_var: None,
        col_names: vec![],
        column_types: vec![],
    });
    let out = node
        .compute_factorized_schema(&[left_schema, right_schema])
        .unwrap();
    out.validate_at_most_one_unflat().unwrap();
    // Right policy mirrors the rewriter: the build (right) key group is
    // flattened fully while the left side keeps its single unflat group.
    assert_eq!(out.unflat_group_pos(), Some(lg1));
    assert!(out.get_group(rg1 + 2).is_none_or(|g| g.is_flat()));
}

#[test]
fn multi_unflat_join_keeps_first_in_position_order() {
    let mut left_schema = FactorizedSchema::new();
    let lg0 = left_schema.create_group();
    left_schema
        .insert_to_group_and_scope(test_id(1), lg0)
        .unwrap();
    let mut right_schema = FactorizedSchema::new();
    let rg0 = right_schema.create_group();
    right_schema
        .insert_to_group_and_scope(test_id(2), rg0)
        .unwrap();
    let mut node = LogicalNodeEnum::InnerJoin(
        crate::planning::plan::logical::logical_nodes::join::LogicalInnerJoinNode {
            id: next_node_id(),
            left: Box::new(scan()),
            right: Box::new(scan()),
            hash_keys: vec![],
            probe_keys: vec![],
            recommended_algorithm: None,
            output_var: None,
            col_names: vec![],
            column_types: vec![],
        },
    );
    let out = node
        .compute_factorized_schema(&[left_schema, right_schema])
        .unwrap();
    out.validate_at_most_one_unflat().unwrap();
    assert_eq!(out.unflat_group_pos(), Some(0));
}

#[test]
fn full_outer_join_compute_flattens_both_key_sides() {
    use crate::planning::plan::logical::logical_nodes::join::LogicalFullOuterJoinNode;
    let ctx = Arc::new(ExpressionAnalysisContext::new());
    let left_key_id = ctx.register_expression(ExpressionMeta::new(
        linkrs_core::Expression::Variable("k".to_string()),
    ));
    let right_key_id = ctx.register_expression(ExpressionMeta::new(
        linkrs_core::Expression::Variable("k".to_string()),
    ));
    let left_key = ContextualExpression::new(left_key_id.clone(), ctx.clone());
    let right_key = ContextualExpression::new(right_key_id.clone(), ctx);
    let mut left_schema = FactorizedSchema::new();
    let lg0 = left_schema.create_flat_group(false);
    let lg1 = left_schema.create_group();
    left_schema
        .insert_to_group_and_scope(test_id(1), lg0)
        .unwrap();
    left_schema
        .insert_to_group_and_scope(left_key_id.clone(), lg1)
        .unwrap();
    let mut right_schema = FactorizedSchema::new();
    let rg0 = right_schema.create_flat_group(false);
    let rg1 = right_schema.create_group();
    right_schema
        .insert_to_group_and_scope(test_id(2), rg0)
        .unwrap();
    right_schema
        .insert_to_group_and_scope(right_key_id.clone(), rg1)
        .unwrap();
    let mut node = LogicalNodeEnum::FullOuterJoin(LogicalFullOuterJoinNode {
        id: next_node_id(),
        left: Box::new(scan()),
        right: Box::new(scan()),
        hash_keys: vec![left_key],
        probe_keys: vec![right_key],
        output_var: None,
        col_names: vec![],
        column_types: vec![],
    });
    let out = node
        .compute_factorized_schema(&[left_schema, right_schema])
        .unwrap();
    out.validate_at_most_one_unflat().unwrap();
    // Full-outer policy flattens both key sides: either side can emit
    // unmatched rows, so no unflat key group survives. With no other
    // unflat groups the output is fully flat, unlike Inner which
    // would keep the right key group alive.
    assert!(out.is_flat_schema());
    assert!(out
        .get_group_pos(&left_key_id)
        .is_some_and(|pos| { out.get_group(pos).is_some_and(|g| g.is_flat()) }));
    assert!(out
        .get_group_pos(&right_key_id)
        .is_some_and(|pos| { out.get_group(pos).is_some_and(|g| g.is_flat()) }));
}

#[test]
fn multi_unflat_non_key_fallback_keeps_first_with_keys_flat() {
    let ctx = Arc::new(ExpressionAnalysisContext::new());
    let left_key_id = ctx.register_expression(ExpressionMeta::new(
        linkrs_core::Expression::Variable("lk".to_string()),
    ));
    let right_key_id = ctx.register_expression(ExpressionMeta::new(
        linkrs_core::Expression::Variable("rk".to_string()),
    ));
    let left_key = ContextualExpression::new(left_key_id.clone(), ctx.clone());
    let right_key = ContextualExpression::new(right_key_id.clone(), ctx);
    // Key groups are already flat; each side carries one unflat
    // non-key group into the merge.
    let mut left_schema = FactorizedSchema::new();
    let lg0 = left_schema.create_flat_group(false);
    let lg1 = left_schema.create_group();
    left_schema
        .insert_to_group_and_scope(left_key_id.clone(), lg0)
        .unwrap();
    left_schema
        .insert_to_group_and_scope(test_id(11), lg1)
        .unwrap();
    let mut right_schema = FactorizedSchema::new();
    let rg0 = right_schema.create_flat_group(false);
    let rg1 = right_schema.create_group();
    right_schema
        .insert_to_group_and_scope(right_key_id.clone(), rg0)
        .unwrap();
    right_schema
        .insert_to_group_and_scope(test_id(22), rg1)
        .unwrap();
    let mut node = LogicalNodeEnum::InnerJoin(
        crate::planning::plan::logical::logical_nodes::join::LogicalInnerJoinNode {
            id: next_node_id(),
            left: Box::new(scan()),
            right: Box::new(scan()),
            hash_keys: vec![left_key],
            probe_keys: vec![right_key],
            recommended_algorithm: None,
            output_var: None,
            col_names: vec![],
            column_types: vec![],
        },
    );
    let out = node
        .compute_factorized_schema(&[left_schema, right_schema])
        .unwrap();
    out.validate_at_most_one_unflat().unwrap();
    // Key groups stayed flat, so the positional fallback keeps the
    // first unflat non-key group (left position 1).
    assert_eq!(out.unflat_group_pos(), Some(lg1));
    assert!(out
        .get_group_pos(&left_key_id)
        .is_some_and(|pos| { out.get_group(pos).is_some_and(|g| g.is_flat()) }));
}

fn select_node(cond: ContextualExpression) -> LogicalNodeEnum {
    LogicalNodeEnum::Select(LogicalSelectNode {
        id: next_node_id(),
        condition: cond,
        if_branch: Some(Box::new(scan())),
        else_branch: Some(Box::new(scan())),
        output_var: None,
        col_names: vec![],
        column_types: vec![],
    })
}

#[test]
fn select_condition_on_flat_keeps_factorization() {
    let (cond, _) = ctx_var("a");
    let mut branch = FactorizedSchema::new();
    let g0 = branch.create_flat_group(false);
    let g1 = branch.create_group();
    branch.insert_to_group_and_scope(test_id(1), g0).unwrap();
    branch.insert_to_group_and_scope(test_id(2), g1).unwrap();
    branch
        .insert_to_group_and_scope_with_name(test_id(3), Some("a".to_string()), g0)
        .unwrap();
    let mut node = select_node(cond);
    let out = node
        .compute_factorized_schema(&[branch.clone(), branch])
        .unwrap();
    out.validate_at_most_one_unflat().unwrap();
    assert_eq!(
        out.unflat_group_pos(),
        Some(g1),
        "condition on a flat group must not flatten the unflat branch group"
    );
}

#[test]
fn select_unresolvable_condition_falls_back_to_flat() {
    let ctx = Arc::new(ExpressionAnalysisContext::new());
    let id = ctx.register_expression(ExpressionMeta::new(linkrs_core::Expression::Variable(
        "ghost".to_string(),
    )));
    let cond = ContextualExpression::new(id, ctx);
    let mut branch = FactorizedSchema::new();
    let g0 = branch.create_flat_group(false);
    let g1 = branch.create_group();
    branch.insert_to_group_and_scope(test_id(1), g0).unwrap();
    branch.insert_to_group_and_scope(test_id(2), g1).unwrap();
    let mut node = select_node(cond);
    let out = node.compute_factorized_schema(&[branch]).unwrap();
    out.validate_at_most_one_unflat().unwrap();
    assert!(
        out.is_flat_schema(),
        "unresolvable condition must conservatively flatten"
    );
}

#[test]
fn loop_condition_on_flat_keeps_factorization() {
    let (cond, _) = ctx_var("a");
    let mut body = FactorizedSchema::new();
    let g0 = body.create_flat_group(false);
    let g1 = body.create_group();
    body.insert_to_group_and_scope(test_id(1), g0).unwrap();
    body.insert_to_group_and_scope(test_id(2), g1).unwrap();
    body.insert_to_group_and_scope_with_name(test_id(3), Some("a".to_string()), g0)
        .unwrap();
    let mut node = LogicalNodeEnum::Loop(LogicalLoopNode::new_with_body(cond, scan()));
    let out = node.compute_factorized_schema(&[body]).unwrap();
    out.validate_at_most_one_unflat().unwrap();
    assert_eq!(out.unflat_group_pos(), Some(g1));
}

#[test]
fn dispatch_module_stays_in_sync_with_compute() {
    // R5 drift guard: every operator maps to the submodule that owns its
    // compute arm. The match in `dispatch_module` is exhaustive, so a new
    // variant fails compilation until both tables are updated.
    assert_eq!(super::dispatch_module(&scan()), "access");
    assert_eq!(
        super::dispatch_module(&LogicalNodeEnum::Flatten(LogicalFlattenNode::new(
            0,
            scan()
        ))),
        "operation"
    );
}