query-engine/tests/sql_pipeline_tests.rs

use query_engine::catalog::PredicateCatalog;
use query_engine::execution::execute;
use query_engine::planner::sql::plan_select;
use query_engine::sql::parser::parse_select;
use query_engine::{Atom, Instance, Term};

fn parent_instance() -> Instance {
    vec![
        Atom::new(
            "Parent",
            vec![Term::constant("alice"), Term::constant("bob")],
        ),
        Atom::new(
            "Parent",
            vec![Term::constant("bob"), Term::constant("carol")],
        ),
    ]
    .into_iter()
    .collect()
}

#[test]
fn select_star_scans_predicate_as_table() {
    let instance = parent_instance();
    let catalog = PredicateCatalog::from_instance(&instance).unwrap();
    let select = parse_select("SELECT * FROM Parent").unwrap();

    let plan = plan_select(&select, &catalog).unwrap();
    let result = execute(&plan, &instance).unwrap();

    assert_eq!(result.schema().len(), 2);
    assert_eq!(result.rows().len(), 2);
    assert_eq!(result.rows()[0].values().len(), 2);
}

#[test]
fn select_projection_keeps_requested_columns() {
    let instance = parent_instance();
    let catalog = PredicateCatalog::from_instance(&instance).unwrap();
    let select = parse_select("SELECT c0 FROM Parent").unwrap();

    let plan = plan_select(&select, &catalog).unwrap();
    let result = execute(&plan, &instance).unwrap();

    assert_eq!(result.schema().len(), 1);
    assert_eq!(result.rows().len(), 2);
    let mut values = result
        .rows()
        .iter()
        .map(|row| format!("{}", row.values()[0]))
        .collect::<Vec<_>>();
    values.sort();
    assert_eq!(values, vec!["alice".to_string(), "bob".to_string()]);
}

#[test]
fn select_where_filters_rows() {
    let instance = parent_instance();
    let catalog = PredicateCatalog::from_instance(&instance).unwrap();
    let select = parse_select("SELECT c0 FROM Parent WHERE c1 = 'bob'").unwrap();

    let plan = plan_select(&select, &catalog).unwrap();
    let result = execute(&plan, &instance).unwrap();

    assert_eq!(result.rows().len(), 1);
    assert_eq!(format!("{}", result.rows()[0].values()[0]), "alice");
}

#[test]
fn select_alias_and_literal_projection_shape_output() {
    let instance = parent_instance();
    let catalog = PredicateCatalog::from_instance(&instance).unwrap();
    let select =
        parse_select("SELECT c0 AS parent_name, 'seed' AS label, NULL FROM Parent").unwrap();

    let plan = plan_select(&select, &catalog).unwrap();
    let result = execute(&plan, &instance).unwrap();

    assert_eq!(result.schema().fields()[0].name(), "parent_name");
    assert_eq!(result.schema().fields()[1].name(), "label");
    assert_eq!(result.schema().fields()[2].name(), "expr3");
    assert_eq!(result.rows().len(), 2);
    assert_eq!(format!("{}", result.rows()[0].values()[1]), "seed");
    assert_eq!(format!("{}", result.rows()[0].values()[2]), "NULL");
}

#[test]
fn select_uses_explicit_catalog_column_names() {
    let instance = parent_instance();
    let mut catalog = PredicateCatalog::from_instance(&instance).unwrap();
    catalog
        .rename_columns("Parent", ["parent", "child"])
        .unwrap();
    let select = parse_select("SELECT parent FROM Parent WHERE child = 'bob'").unwrap();

    let plan = plan_select(&select, &catalog).unwrap();
    let result = execute(&plan, &instance).unwrap();

    assert_eq!(result.schema().fields()[0].name(), "parent");
    assert_eq!(result.rows().len(), 1);
    assert_eq!(format!("{}", result.rows()[0].values()[0]), "alice");
}

#[test]
fn select_uses_qualified_table_name_in_single_table_query() {
    let instance = parent_instance();
    let mut catalog = PredicateCatalog::from_instance(&instance).unwrap();
    catalog
        .rename_columns("Parent", ["parent", "child"])
        .unwrap();

    let select =
        parse_select("SELECT Parent.parent FROM Parent WHERE Parent.child = 'bob'").unwrap();

    let plan = plan_select(&select, &catalog).unwrap();
    let result = execute(&plan, &instance).unwrap();

    assert_eq!(result.schema().fields()[0].name(), "Parent.parent");
    assert_eq!(result.rows().len(), 1);
    assert_eq!(format!("{}", result.rows()[0].values()[0]), "alice");
}

#[test]
fn select_join_filters_cross_product_by_qualified_columns() {
    let instance: Instance = vec![
        Atom::new(
            "Parent",
            vec![Term::constant("alice"), Term::constant("bob")],
        ),
        Atom::new(
            "Parent",
            vec![Term::constant("bob"), Term::constant("carol")],
        ),
        Atom::new(
            "Ancestor",
            vec![Term::constant("bob"), Term::constant("carol")],
        ),
        Atom::new(
            "Ancestor",
            vec![Term::constant("carol"), Term::constant("dave")],
        ),
    ]
    .into_iter()
    .collect();

    let mut catalog = PredicateCatalog::from_instance(&instance).unwrap();
    catalog
        .rename_columns("Parent", ["parent", "child"])
        .unwrap();
    catalog
        .rename_columns("Ancestor", ["parent", "child"])
        .unwrap();

    let select = parse_select(
        "SELECT Parent.parent, Ancestor.child FROM Parent, Ancestor \
         WHERE Parent.child = Ancestor.parent",
    )
    .unwrap();

    let plan = plan_select(&select, &catalog).unwrap();
    let result = execute(&plan, &instance).unwrap();

    assert_eq!(result.schema().fields()[0].name(), "Parent.parent");
    assert_eq!(result.schema().fields()[1].name(), "Ancestor.child");
    assert_eq!(result.rows().len(), 2);

    let mut rows = result
        .rows()
        .iter()
        .map(|row| format!("{} -> {}", row.values()[0], row.values()[1]))
        .collect::<Vec<_>>();
    rows.sort();
    assert_eq!(
        rows,
        vec!["alice -> carol".to_string(), "bob -> dave".to_string()]
    );
}

#[test]
fn select_self_join_uses_table_aliases() {
    let instance: Instance = vec![
        Atom::new(
            "Parent",
            vec![Term::constant("alice"), Term::constant("bob")],
        ),
        Atom::new(
            "Parent",
            vec![Term::constant("bob"), Term::constant("carol")],
        ),
        Atom::new(
            "Parent",
            vec![Term::constant("carol"), Term::constant("dave")],
        ),
    ]
    .into_iter()
    .collect();

    let mut catalog = PredicateCatalog::from_instance(&instance).unwrap();
    catalog
        .rename_columns("Parent", ["parent", "child"])
        .unwrap();

    let select = parse_select(
        "SELECT p.parent, q.child FROM Parent AS p, Parent AS q \
         WHERE p.child = q.parent",
    )
    .unwrap();

    let plan = plan_select(&select, &catalog).unwrap();
    let result = execute(&plan, &instance).unwrap();

    assert_eq!(result.schema().fields()[0].name(), "p.parent");
    assert_eq!(result.schema().fields()[1].name(), "q.child");

    let mut rows = result
        .rows()
        .iter()
        .map(|row| format!("{} -> {}", row.values()[0], row.values()[1]))
        .collect::<Vec<_>>();
    rows.sort();
    assert_eq!(
        rows,
        vec!["alice -> carol".to_string(), "bob -> dave".to_string()]
    );
}

#[test]
fn select_where_and_applies_multiple_filters() {
    let instance = parent_instance();
    let catalog = PredicateCatalog::from_instance(&instance).unwrap();
    let select = parse_select("SELECT c0 FROM Parent WHERE c1 = 'bob' AND c0 = 'alice'").unwrap();

    let plan = plan_select(&select, &catalog).unwrap();
    let result = execute(&plan, &instance).unwrap();

    assert_eq!(result.rows().len(), 1);
    assert_eq!(format!("{}", result.rows()[0].values()[0]), "alice");
}

#[test]
fn select_order_by_desc_sorts_rows() {
    let instance = parent_instance();
    let catalog = PredicateCatalog::from_instance(&instance).unwrap();
    let select = parse_select("SELECT c0 FROM Parent ORDER BY c0 DESC").unwrap();

    let plan = plan_select(&select, &catalog).unwrap();
    let result = execute(&plan, &instance).unwrap();

    assert_eq!(result.rows().len(), 2);
    assert_eq!(format!("{}", result.rows()[0].values()[0]), "bob");
    assert_eq!(format!("{}", result.rows()[1].values()[0]), "alice");
}

#[test]
fn select_integer_literal_in_projection() {
    let instance = parent_instance();
    let catalog = PredicateCatalog::from_instance(&instance).unwrap();
    let select = parse_select("SELECT c0, 42 AS answer FROM Parent").unwrap();

    let plan = plan_select(&select, &catalog).unwrap();
    let result = execute(&plan, &instance).unwrap();

    assert_eq!(result.schema().len(), 2);
    assert_eq!(result.schema().fields()[1].name(), "answer");
    assert_eq!(
        result.schema().fields()[1].data_type(),
        &query_engine::relational::DataType::Integer
    );
    assert_eq!(result.rows().len(), 2);
    assert_eq!(format!("{}", result.rows()[0].values()[1]), "42");
}

#[test]
fn select_limit_restricts_row_count() {
    let instance = parent_instance();
    let catalog = PredicateCatalog::from_instance(&instance).unwrap();
    let select = parse_select("SELECT c0 FROM Parent ORDER BY c0 ASC LIMIT 1").unwrap();

    let plan = plan_select(&select, &catalog).unwrap();
    let result = execute(&plan, &instance).unwrap();

    assert_eq!(result.rows().len(), 1);
    assert_eq!(format!("{}", result.rows()[0].values()[0]), "alice");
}

#[test]
fn select_where_or_matches_either_condition() {
    let instance: Instance = vec![
        Atom::new(
            "Parent",
            vec![Term::constant("alice"), Term::constant("bob")],
        ),
        Atom::new(
            "Parent",
            vec![Term::constant("bob"), Term::constant("carol")],
        ),
        Atom::new(
            "Parent",
            vec![Term::constant("carol"), Term::constant("dave")],
        ),
    ]
    .into_iter()
    .collect();
    let catalog = PredicateCatalog::from_instance(&instance).unwrap();
    let select = parse_select("SELECT c0 FROM Parent WHERE c1 = 'bob' OR c1 = 'dave'").unwrap();

    let plan = plan_select(&select, &catalog).unwrap();
    let result = execute(&plan, &instance).unwrap();

    assert_eq!(result.rows().len(), 2);
    let mut values = result
        .rows()
        .iter()
        .map(|row| format!("{}", row.values()[0]))
        .collect::<Vec<_>>();
    values.sort();
    assert_eq!(values, vec!["alice".to_string(), "carol".to_string()]);
}

#[test]
fn select_where_not_equal_excludes_matching_rows() {
    let instance = parent_instance();
    let catalog = PredicateCatalog::from_instance(&instance).unwrap();
    let select = parse_select("SELECT c0 FROM Parent WHERE c1 != 'bob'").unwrap();

    let plan = plan_select(&select, &catalog).unwrap();
    let result = execute(&plan, &instance).unwrap();

    assert_eq!(result.rows().len(), 1);
    assert_eq!(format!("{}", result.rows()[0].values()[0]), "bob");
}

#[test]
fn execute_with_table_store_scans_in_memory_rows() {
    use query_engine::execution::TableStore;
    use query_engine::relational::{DataType, Field, Row, Schema, Value};

    let schema = Schema::new(vec![
        Field::new("name", DataType::Text, false),
        Field::new("age", DataType::Integer, false),
    ]);

    let mut store = TableStore::new();
    store.insert(
        "people",
        schema.clone(),
        vec![
            Row::new(vec![Value::text("alice"), Value::Integer(30)]),
            Row::new(vec![Value::text("bob"), Value::Integer(25)]),
        ],
    );

    let mut catalog = PredicateCatalog::new();
    catalog.register_table("people", schema);

    let select = parse_select("SELECT name FROM people WHERE age != 30").unwrap();
    let plan = plan_select(&select, &catalog).unwrap();
    let result = execute(&plan, &store).unwrap();

    assert_eq!(result.rows().len(), 1);
    assert_eq!(format!("{}", result.rows()[0].values()[0]), "bob");
}