plexus-engine 0.3.4

use super::*;
use plexus_serde::{Expr, Op, Plan, VectorMetric};

#[test]
fn mock_vector_engine_executes_vector_scan_cosine() {
    let plan = Plan {
        version: version(),
        ops: vec![
            Op::ConstRow,
            Op::VectorScan {
                input: 0,
                collection: "docs".to_string(),
                query_vector: Expr::ListLiteral {
                    items: vec![Expr::FloatLiteral(1.0), Expr::FloatLiteral(0.0)],
                },
                metric: VectorMetric::Cosine,
                top_k: 2,
                approx_hint: false,
                schema: vec![
                    ColDef {
                        name: "doc".to_string(),
                        kind: ColKind::Node,
                        logical_type: LogicalType::NodeRef,
                    },
                    ColDef {
                        name: "score".to_string(),
                        kind: ColKind::Float64,
                        logical_type: LogicalType::Float64,
                    },
                ],
            },
            Op::Return { input: 1 },
        ],
        root_op: 2,
    };

    let mut engine = mock_vector_engine();
    let out = engine.execute_plan(&plan).expect("execute");
    assert_eq!(out.rows.len(), 2);
    assert_eq!(out.rows[0][0], Value::NodeRef(1));
    assert_eq!(out.rows[1][0], Value::NodeRef(3));
    match &out.rows[0][1] {
        Value::Float(v) => assert!((*v - 1.0).abs() < 1e-9),
        other => panic!("unexpected score value: {other:?}"),
    }
}

#[test]
fn mock_vector_engine_executes_rerank_with_vector_similarity() {
    let plan = Plan {
        version: version(),
        ops: vec![
            Op::ConstRow,
            Op::VectorScan {
                input: 0,
                collection: "docs".to_string(),
                query_vector: Expr::ListLiteral {
                    items: vec![Expr::FloatLiteral(1.0), Expr::FloatLiteral(0.0)],
                },
                metric: VectorMetric::Cosine,
                top_k: 3,
                approx_hint: false,
                schema: vec![
                    ColDef {
                        name: "doc".to_string(),
                        kind: ColKind::Node,
                        logical_type: LogicalType::NodeRef,
                    },
                    ColDef {
                        name: "score".to_string(),
                        kind: ColKind::Float64,
                        logical_type: LogicalType::Float64,
                    },
                ],
            },
            Op::Rerank {
                input: 1,
                score_expr: Expr::VectorSimilarity {
                    metric: VectorMetric::DotProduct,
                    lhs: Box::new(Expr::PropAccess {
                        col: 0,
                        prop: "embedding".to_string(),
                    }),
                    rhs: Box::new(Expr::ListLiteral {
                        items: vec![Expr::FloatLiteral(0.0), Expr::FloatLiteral(1.0)],
                    }),
                },
                top_k: 2,
                schema: vec![
                    ColDef {
                        name: "doc".to_string(),
                        kind: ColKind::Node,
                        logical_type: LogicalType::NodeRef,
                    },
                    ColDef {
                        name: "score".to_string(),
                        kind: ColKind::Float64,
                        logical_type: LogicalType::Float64,
                    },
                ],
            },
            Op::Return { input: 2 },
        ],
        root_op: 3,
    };

    let mut engine = mock_vector_engine();
    let out = engine.execute_plan(&plan).expect("execute");
    assert_eq!(
        out.rows,
        vec![
            vec![Value::NodeRef(2), Value::Float(0.0)],
            vec![Value::NodeRef(3), Value::Float(0.9701425001453318)]
        ]
    );
}

#[test]
fn mock_vector_engine_projects_vector_similarity() {
    let plan = Plan {
        version: version(),
        ops: vec![
            Op::ScanNodes {
                labels: vec!["Doc".to_string()],
                schema: vec![ColDef {
                    name: "d".to_string(),
                    kind: ColKind::Node,
                    logical_type: LogicalType::NodeRef,
                }],
                must_labels: vec![],
                forbidden_labels: vec![],
                graph_ref: None,
                est_rows: -1,
                selectivity: 1.0,
            },
            Op::Project {
                input: 0,
                exprs: vec![
                    Expr::PropAccess {
                        col: 0,
                        prop: "title".to_string(),
                    },
                    Expr::VectorSimilarity {
                        metric: VectorMetric::Cosine,
                        lhs: Box::new(Expr::PropAccess {
                            col: 0,
                            prop: "embedding".to_string(),
                        }),
                        rhs: Box::new(Expr::ListLiteral {
                            items: vec![Expr::FloatLiteral(1.0), Expr::FloatLiteral(0.0)],
                        }),
                    },
                ],
                schema: vec![
                    ColDef {
                        name: "title".to_string(),
                        kind: ColKind::String,
                        logical_type: LogicalType::String,
                    },
                    ColDef {
                        name: "score".to_string(),
                        kind: ColKind::Float64,
                        logical_type: LogicalType::Float64,
                    },
                ],
            },
            Op::Sort {
                input: 1,
                keys: vec![1],
                dirs: vec![plexus_serde::SortDir::Desc],
            },
            Op::Return { input: 2 },
        ],
        root_op: 3,
    };

    let mut engine = mock_vector_engine();
    let out = engine.execute_plan(&plan).expect("execute");
    assert_eq!(out.rows[0][0], Value::String("Alpha".to_string()));
    assert_eq!(out.rows[1][0], Value::String("Gamma".to_string()));
}