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_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()]
    );
}