Fix index scan rewrite rule (#2728)

Author: joshua-spacetime

crates/core/src/sql/execute.rs

@@ -495,6 +495,39 @@ pub(crate) mod tests {
         );
     }
 
+    /// Test a query that uses a multi-column index
+    #[test]
+    fn test_multi_column_index() -> anyhow::Result<()> {
+        let db = TestDB::in_memory()?;
+
+        let schema = [
+            ("a", AlgebraicType::U64),
+            ("b", AlgebraicType::U64),
+            ("c", AlgebraicType::U64),
+        ];
+
+        let table_id = db.create_table_for_test_multi_column("t", &schema, [1, 2].into())?;
+
+        insert_rows(
+            &db,
+            table_id,
+            vec![
+                product![0_u64, 1_u64, 2_u64],
+                product![1_u64, 2_u64, 1_u64],
+                product![2_u64, 2_u64, 2_u64],
+            ],
+        )?;
+
+        assert_query_results(
+            &db,
+            "select * from t where c = 1 and b = 2",
+            &AuthCtx::for_testing(),
+            [product![1_u64, 2_u64, 1_u64]],
+        );
+
+        Ok(())
+    }
+
     /// Test querying a table with RLS rules
     #[test]
     fn test_rls_rules() -> anyhow::Result<()> {

crates/physical-plan/src/plan.rs

@@ -1830,10 +1830,10 @@ mod tests {
                 _,
             )) => {
                 assert_eq!(schema.table_id, t_id);
-                assert_eq!(arg, Sarg::Eq(ColId(1), AlgebraicValue::U8(3)));
+                assert_eq!(arg, Sarg::Eq(ColId(3), AlgebraicValue::U8(5)));
                 assert_eq!(
                     prefix,
-                    vec![(ColId(3), AlgebraicValue::U8(5)), (ColId(2), AlgebraicValue::U8(4))]
+                    vec![(ColId(1), AlgebraicValue::U8(3)), (ColId(2), AlgebraicValue::U8(4))]
                 );
             }
             proj => panic!("unexpected plan: {:#?}", proj),
@@ -1941,8 +1941,8 @@ mod tests {
                 _,
             )) => {
                 assert_eq!(schema.table_id, t_id);
-                assert_eq!(arg, Sarg::Eq(ColId(2), AlgebraicValue::U8(1)));
-                assert_eq!(prefix, vec![(ColId(3), AlgebraicValue::U8(2))]);
+                assert_eq!(arg, Sarg::Eq(ColId(3), AlgebraicValue::U8(2)));
+                assert_eq!(prefix, vec![(ColId(2), AlgebraicValue::U8(1))]);
             }
             proj => panic!("unexpected plan: {:#?}", proj),
         };

crates/physical-plan/src/rules.rs

@@ -622,13 +622,20 @@ impl RewriteRule for IxScanEq2Col {
                     .indexes
                     .iter()
                     .filter(|idx| idx.index_algorithm.columns().len() == 2) // TODO: Support prefix scans
-                    .find_map(|idx| {
+                    .map(|idx| (idx.index_id, idx.index_algorithm.columns()))
+                    .find_map(|(index_id, columns)| {
+                        let mut columns = columns.iter();
+                        let x = columns.next()?;
+                        if x.idx() != u.field_pos {
+                            return None;
+                        }
+                        let y = columns.next()?;
+                        if y.idx() != v.field_pos {
+                            return None;
+                        }
                         Some(IxScanInfo {
-                            index_id: idx.index_id,
-                            cols: vec![
-                                (i, idx.index_algorithm.find_col_index(u.field_pos)?),
-                                (j, idx.index_algorithm.find_col_index(v.field_pos)?),
-                            ],
+                            index_id,
+                            cols: vec![(i, x), (j, y)],
                         })
                     })
                 {
@@ -795,14 +802,24 @@ impl RewriteRule for IxScanEq3Col {
                         .indexes
                         .iter()
                         .filter(|idx| idx.index_algorithm.columns().len() == 3)
-                        .find_map(|idx| {
+                        .map(|idx| (idx.index_id, idx.index_algorithm.columns()))
+                        .find_map(|(index_id, columns)| {
+                            let mut columns = columns.iter();
+                            let x = columns.next()?;
+                            if x.idx() != u.field_pos {
+                                return None;
+                            }
+                            let y = columns.next()?;
+                            if y.idx() != v.field_pos {
+                                return None;
+                            }
+                            let z = columns.next()?;
+                            if z.idx() != w.field_pos {
+                                return None;
+                            }
                             Some(IxScanInfo {
-                                index_id: idx.index_id,
-                                cols: vec![
-                                    (i, idx.index_algorithm.find_col_index(u.field_pos)?),
-                                    (j, idx.index_algorithm.find_col_index(v.field_pos)?),
-                                    (k, idx.index_algorithm.find_col_index(w.field_pos)?),
-                                ],
+                                index_id,
+                                cols: vec![(i, x), (j, y), (k, z)],
                             })
                         })
                     {