From b0ac042092b3c83c99aca03a59c03c00a9ddf2e0 Mon Sep 17 00:00:00 2001
From: PThorpe92 <preston@turso.tech>
Date: Tue, 18 Nov 2025 20:17:52 -0500
Subject: [PATCH] Add tests for using expr indexes in query plans for ORDER BY
 and WHERE predicates

---
 tests/integration/index_method/mod.rs         |  3 +
 tests/integration/query_processing/mod.rs     |  1 +
 .../query_processing/test_expr_index.rs       | 76 +++++++++++++++++++
 3 files changed, 80 insertions(+)
 create mode 100644 tests/integration/query_processing/test_expr_index.rs

diff --git a/tests/integration/index_method/mod.rs b/tests/integration/index_method/mod.rs
index 8285a258f..a2f1f3bee 100644
--- a/tests/integration/index_method/mod.rs
+++ b/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(),
         })
diff --git a/tests/integration/query_processing/mod.rs b/tests/integration/query_processing/mod.rs
index e24945f20..95b322581 100644
--- a/tests/integration/query_processing/mod.rs
+++ b/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;
diff --git a/tests/integration/query_processing/test_expr_index.rs b/tests/integration/query_processing/test_expr_index.rs
new file mode 100644
index 000000000..db01c584f
--- /dev/null
+++ b/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(())
+}