Add tests for using expr indexes in query plans for ORDER BY and WHERE predicates

tests/integration/index_method/mod.rs

@@ -64,6 +64,7 @@ fn test_vector_sparse_ivf_create_destroy() {
                 pos_in_table: 1,
                 collation: None,
                 default: None,
+                expr: None,
             }],
             parameters: HashMap::new(),
         })
@@ -106,6 +107,7 @@ fn test_vector_sparse_ivf_insert_query() {
                 pos_in_table: 1,
                 collation: None,
                 default: None,
+                expr: None,
             }],
             parameters: HashMap::new(),
         })
@@ -196,6 +198,7 @@ fn test_vector_sparse_ivf_update() {
                 pos_in_table: 1,
                 collation: None,
                 default: None,
+                expr: None,
             }],
             parameters: HashMap::new(),
         })

tests/integration/query_processing/mod.rs

@@ -6,3 +6,4 @@ mod test_write_path;
 mod encryption;
 mod test_multi_thread;
 mod test_transactions;
+mod test_expr_index;

tests/integration/query_processing/test_expr_index.rs

@@ -0,0 +1,76 @@
+use std::sync::Arc;
+
+use crate::common::{limbo_exec_rows, TempDatabase};
+use turso_core::StepResult;
+
+fn explain_plans(conn: &Arc<turso_core::Connection>, sql: &str) -> anyhow::Result<Vec<String>> {
+    let mut stmt = conn.prepare(format!("EXPLAIN QUERY PLAN {sql}"))?;
+    let mut plans = Vec::new();
+    loop {
+        match stmt.step()? {
+            StepResult::Row => {
+                let row = stmt.row().unwrap();
+                plans.push(row.get::<String>(3)?);
+            }
+            StepResult::IO => {
+                stmt.run_once()?;
+            }
+            StepResult::Done => break,
+            _ => unreachable!(),
+        }
+    }
+    Ok(plans)
+}
+
+#[test]
+fn expression_index_used_for_where() -> anyhow::Result<()> {
+    let _ = env_logger::try_init();
+    let tmp_db = TempDatabase::new_with_rusqlite("CREATE TABLE t (a INT, b INT, c INT);");
+    let conn = tmp_db.connect_limbo();
+
+    conn.execute("INSERT INTO t VALUES (1, 2, 0)")?;
+    conn.execute("INSERT INTO t VALUES (3, 4, 0)")?;
+    conn.execute("INSERT INTO t VALUES (5, 6, 0)")?;
+
+    conn.execute("CREATE INDEX idx_expr ON t(a + b)")?;
+
+    let plans = explain_plans(&conn, "SELECT * FROM t WHERE a + b = 7")?;
+    assert!(
+        plans.iter().any(|p| p.contains("idx_expr")),
+        "expected query plan to mention idx_expr, got {plans:?}"
+    );
+
+    let rows = limbo_exec_rows(&tmp_db, &conn, "SELECT a, b FROM t WHERE a + b = 7");
+    assert_eq!(rows, vec![vec![3.into(), 4.into()]]);
+    Ok(())
+}
+
+#[test]
+fn expression_index_used_for_order_by() -> anyhow::Result<()> {
+    let _ = env_logger::try_init();
+    let tmp_db = TempDatabase::new_with_rusqlite("CREATE TABLE t (a INT, b INT);");
+    let conn = tmp_db.connect_limbo();
+
+    conn.execute("INSERT INTO t VALUES (2, 2)")?;
+    conn.execute("INSERT INTO t VALUES (1, 3)")?;
+    conn.execute("INSERT INTO t VALUES (0, 5)")?;
+
+    conn.execute("CREATE INDEX idx_expr_order ON t(a + b)")?;
+
+    let plans = explain_plans(
+        &conn,
+        "SELECT a, b FROM t WHERE a + b > 0 ORDER BY a + b DESC LIMIT 1 OFFSET 0",
+    )?;
+    assert!(
+        plans.iter().any(|p| p.contains("idx_expr_order")),
+        "expected query plan to mention idx_expr_order, got {plans:?}"
+    );
+
+    let rows = limbo_exec_rows(
+        &tmp_db,
+        &conn,
+        "SELECT a, b FROM t WHERE a + b > 0 ORDER BY a + b DESC LIMIT 1",
+    );
+    assert_eq!(rows, vec![vec![0.into(), 5.into()]]);
+    Ok(())
+}