Skip to content

Commit 30c488e

Browse files
jussisauriojussisaurio
committed
Merge 'Feat: add support for descending indexes' from Jussi Saurio
### Feat: - Adds support for descending indexes ### Testing: - Augments existing compound key index seek fuzz test to test for various combinations of ascending and descending indexed columns. To illustrate, the test runs 10000 queries like this on 8 different tables with all different asc/desc permutations of a three-column primary key index: ```sql query: SELECT * FROM t WHERE x = 2826 LIMIT 5 query: SELECT * FROM t WHERE x = 671 AND y >= 2447 ORDER BY x ASC, y DESC LIMIT 5 query: SELECT * FROM t WHERE x = 2412 AND y = 589 AND z >= 894 ORDER BY x DESC LIMIT 5 query: SELECT * FROM t WHERE x = 1217 AND y = 1437 AND z <= 265 ORDER BY x ASC, y ASC, z DESC LIMIT 5 query: SELECT * FROM t WHERE x < 138 ORDER BY x DESC LIMIT 5 query: SELECT * FROM t WHERE x = 1312 AND y = 2757 AND z > 39 ORDER BY x DESC, y ASC, z ASC LIMIT 5 query: SELECT * FROM t WHERE x = 1829 AND y >= 1629 ORDER BY x ASC, y ASC LIMIT 5 query: SELECT * FROM t WHERE x = 2047 ORDER BY x DESC LIMIT 5 query: SELECT * FROM t WHERE x = 893 AND y > 432 ORDER BY y DESC LIMIT 5 query: SELECT * FROM t WHERE x = 1865 AND y = 784 AND z <= 785 ORDER BY x DESC, y DESC, z DESC LIMIT 5 query: SELECT * FROM t WHERE x = 213 AND y = 1475 AND z <= 2870 ORDER BY x ASC, y ASC, z ASC LIMIT 5 query: SELECT * FROM t WHERE x >= 1780 ORDER BY x ASC LIMIT 5 query: SELECT * FROM t WHERE x = 1983 AND y = 602 AND z = 485 ORDER BY y ASC, z ASC LIMIT 5 query: SELECT * FROM t WHERE x = 2311 AND y >= 31 ORDER BY y DESC LIMIT 5 query: SELECT * FROM t WHERE x = 81 AND y >= 1037 ORDER BY x ASC, y DESC LIMIT 5 query: SELECT * FROM t WHERE x < 2698 ORDER BY x ASC LIMIT 5 query: SELECT * FROM t WHERE x = 1503 AND y = 554 AND z >= 185 ORDER BY x DESC, y DESC, z DESC LIMIT 5 query: SELECT * FROM t WHERE x = 619 AND y > 1414 ORDER BY x DESC, y ASC LIMIT 5 query: SELECT * FROM t WHERE x >= 865 ORDER BY x DESC LIMIT 5 query: SELECT * FROM t WHERE x = 1596 AND y = 622 AND z = 62 ORDER BY x DESC, z ASC LIMIT 5 query: SELECT * FROM t WHERE x = 1555 AND y = 1257 AND z < 1929 ORDER BY x ASC, y ASC, z ASC LIMIT 5 query: SELECT * FROM t WHERE x > 2598 LIMIT 5 query: SELECT * FROM t WHERE x = 302 AND y = 2476 AND z < 2302 ORDER BY z DESC LIMIT 5 query: SELECT * FROM t WHERE x = 2197 AND y = 2195 AND z > 2089 ORDER BY y ASC, z DESC LIMIT 5 query: SELECT * FROM t WHERE x = 1030 AND y = 1717 AND z < 987 LIMIT 5 query: SELECT * FROM t WHERE x = 2899 AND y >= 382 ORDER BY y DESC LIMIT 5 query: SELECT * FROM t WHERE x = 62 AND y = 2980 AND z < 1109 ORDER BY x DESC, y DESC, z DESC LIMIT 5 query: SELECT * FROM t WHERE x = 550 AND y > 221 ORDER BY y DESC LIMIT 5 query: SELECT * FROM t WHERE x = 376 AND y = 1874 AND z < 206 ORDER BY y DESC, z ASC LIMIT 5 query: SELECT * FROM t WHERE x = 859 AND y = 2157 ORDER BY x DESC LIMIT 5 query: SELECT * FROM t WHERE x = 2166 AND y = 2079 AND z < 301 ORDER BY x DESC, y ASC LIMIT 5 ``` the queries are run against both sqlite and limbo. Reviewed-by: Pere Diaz Bou (@pereman2) Reviewed-by: Preston Thorpe (@PThorpe92) Closes #1330
2 parents 7a3fc33 + 95bc644 commit 30c488e

File tree

8 files changed

+741
-316
lines changed

8 files changed

+741
-316
lines changed

‎core/schema.rs

+49-31
Original file line numberDiff line numberDiff line change
@@ -158,16 +158,16 @@ impl PartialEq for Table {
158158
pub struct BTreeTable {
159159
pub root_page: usize,
160160
pub name: String,
161-
pub primary_key_column_names: Vec<String>,
161+
pub primary_key_columns: Vec<(String, SortOrder)>,
162162
pub columns: Vec<Column>,
163163
pub has_rowid: bool,
164164
pub is_strict: bool,
165165
}
166166

167167
impl BTreeTable {
168168
pub fn get_rowid_alias_column(&self) -> Option<(usize, &Column)> {
169-
if self.primary_key_column_names.len() == 1 {
170-
let (idx, col) = self.get_column(&self.primary_key_column_names[0]).unwrap();
169+
if self.primary_key_columns.len() == 1 {
170+
let (idx, col) = self.get_column(&self.primary_key_columns[0].0).unwrap();
171171
if self.column_is_rowid_alias(col) {
172172
return Some((idx, col));
173173
}
@@ -265,7 +265,7 @@ fn create_table(
265265
let table_name = normalize_ident(&tbl_name.name.0);
266266
trace!("Creating table {}", table_name);
267267
let mut has_rowid = true;
268-
let mut primary_key_column_names = vec![];
268+
let mut primary_key_columns = vec![];
269269
let mut cols = vec![];
270270
let is_strict: bool;
271271
match body {
@@ -282,15 +282,17 @@ fn create_table(
282282
} = c.constraint
283283
{
284284
for column in columns {
285-
primary_key_column_names.push(match column.expr {
285+
let col_name = match column.expr {
286286
Expr::Id(id) => normalize_ident(&id.0),
287287
Expr::Literal(Literal::String(value)) => {
288288
value.trim_matches('\'').to_owned()
289289
}
290290
_ => {
291291
todo!("Unsupported primary key expression");
292292
}
293-
});
293+
};
294+
primary_key_columns
295+
.push((col_name, column.order.unwrap_or(SortOrder::Asc)));
294296
}
295297
}
296298
}
@@ -347,10 +349,17 @@ fn create_table(
347349
let mut default = None;
348350
let mut primary_key = false;
349351
let mut notnull = false;
352+
let mut order = SortOrder::Asc;
350353
for c_def in &col_def.constraints {
351354
match &c_def.constraint {
352-
limbo_sqlite3_parser::ast::ColumnConstraint::PrimaryKey { .. } => {
355+
limbo_sqlite3_parser::ast::ColumnConstraint::PrimaryKey {
356+
order: o,
357+
..
358+
} => {
353359
primary_key = true;
360+
if let Some(o) = o {
361+
order = o.clone();
362+
}
354363
}
355364
limbo_sqlite3_parser::ast::ColumnConstraint::NotNull { .. } => {
356365
notnull = true;
@@ -363,8 +372,11 @@ fn create_table(
363372
}
364373

365374
if primary_key {
366-
primary_key_column_names.push(name.clone());
367-
} else if primary_key_column_names.contains(&name) {
375+
primary_key_columns.push((name.clone(), order));
376+
} else if primary_key_columns
377+
.iter()
378+
.any(|(col_name, _)| col_name == &name)
379+
{
368380
primary_key = true;
369381
}
370382

@@ -386,7 +398,7 @@ fn create_table(
386398
};
387399
// flip is_rowid_alias back to false if the table has multiple primary keys
388400
// or if the table has no rowid
389-
if !has_rowid || primary_key_column_names.len() > 1 {
401+
if !has_rowid || primary_key_columns.len() > 1 {
390402
for col in cols.iter_mut() {
391403
col.is_rowid_alias = false;
392404
}
@@ -395,7 +407,7 @@ fn create_table(
395407
root_page,
396408
name: table_name,
397409
has_rowid,
398-
primary_key_column_names,
410+
primary_key_columns,
399411
columns: cols,
400412
is_strict,
401413
})
@@ -621,7 +633,7 @@ pub fn sqlite_schema_table() -> BTreeTable {
621633
name: "sqlite_schema".to_string(),
622634
has_rowid: true,
623635
is_strict: false,
624-
primary_key_column_names: vec![],
636+
primary_key_columns: vec![],
625637
columns: vec![
626638
Column {
627639
name: Some("type".to_string()),
@@ -740,16 +752,16 @@ impl Index {
740752
index_name: &str,
741753
root_page: usize,
742754
) -> Result<Index> {
743-
if table.primary_key_column_names.is_empty() {
755+
if table.primary_key_columns.is_empty() {
744756
return Err(crate::LimboError::InternalError(
745757
"Cannot create automatic index for table without primary key".to_string(),
746758
));
747759
}
748760

749761
let index_columns = table
750-
.primary_key_column_names
762+
.primary_key_columns
751763
.iter()
752-
.map(|col_name| {
764+
.map(|(col_name, order)| {
753765
// Verify that each primary key column exists in the table
754766
let Some((pos_in_table, _)) = table.get_column(col_name) else {
755767
return Err(crate::LimboError::InternalError(format!(
@@ -759,7 +771,7 @@ impl Index {
759771
};
760772
Ok(IndexColumn {
761773
name: normalize_ident(col_name),
762-
order: SortOrder::Asc, // Primary key indexes are always ascending
774+
order: order.clone(),
763775
pos_in_table,
764776
})
765777
})
@@ -905,8 +917,8 @@ mod tests {
905917
let column = table.get_column("c").unwrap().1;
906918
assert!(!column.primary_key, "column 'c' shouldn't be a primary key");
907919
assert_eq!(
908-
vec!["a"],
909-
table.primary_key_column_names,
920+
vec![("a".to_string(), SortOrder::Asc)],
921+
table.primary_key_columns,
910922
"primary key column names should be ['a']"
911923
);
912924
Ok(())
@@ -923,16 +935,19 @@ mod tests {
923935
let column = table.get_column("c").unwrap().1;
924936
assert!(!column.primary_key, "column 'c' shouldn't be a primary key");
925937
assert_eq!(
926-
vec!["a", "b"],
927-
table.primary_key_column_names,
938+
vec![
939+
("a".to_string(), SortOrder::Asc),
940+
("b".to_string(), SortOrder::Asc)
941+
],
942+
table.primary_key_columns,
928943
"primary key column names should be ['a', 'b']"
929944
);
930945
Ok(())
931946
}
932947

933948
#[test]
934949
pub fn test_primary_key_separate_single() -> Result<()> {
935-
let sql = r#"CREATE TABLE t1 (a INTEGER, b TEXT, c REAL, PRIMARY KEY(a));"#;
950+
let sql = r#"CREATE TABLE t1 (a INTEGER, b TEXT, c REAL, PRIMARY KEY(a desc));"#;
936951
let table = BTreeTable::from_sql(sql, 0)?;
937952
let column = table.get_column("a").unwrap().1;
938953
assert!(column.primary_key, "column 'a' should be a primary key");
@@ -941,16 +956,16 @@ mod tests {
941956
let column = table.get_column("c").unwrap().1;
942957
assert!(!column.primary_key, "column 'c' shouldn't be a primary key");
943958
assert_eq!(
944-
vec!["a"],
945-
table.primary_key_column_names,
959+
vec![("a".to_string(), SortOrder::Desc)],
960+
table.primary_key_columns,
946961
"primary key column names should be ['a']"
947962
);
948963
Ok(())
949964
}
950965

951966
#[test]
952967
pub fn test_primary_key_separate_multiple() -> Result<()> {
953-
let sql = r#"CREATE TABLE t1 (a INTEGER, b TEXT, c REAL, PRIMARY KEY(a, b));"#;
968+
let sql = r#"CREATE TABLE t1 (a INTEGER, b TEXT, c REAL, PRIMARY KEY(a, b desc));"#;
954969
let table = BTreeTable::from_sql(sql, 0)?;
955970
let column = table.get_column("a").unwrap().1;
956971
assert!(column.primary_key, "column 'a' should be a primary key");
@@ -959,8 +974,11 @@ mod tests {
959974
let column = table.get_column("c").unwrap().1;
960975
assert!(!column.primary_key, "column 'c' shouldn't be a primary key");
961976
assert_eq!(
962-
vec!["a", "b"],
963-
table.primary_key_column_names,
977+
vec![
978+
("a".to_string(), SortOrder::Asc),
979+
("b".to_string(), SortOrder::Desc)
980+
],
981+
table.primary_key_columns,
964982
"primary key column names should be ['a', 'b']"
965983
);
966984
Ok(())
@@ -977,8 +995,8 @@ mod tests {
977995
let column = table.get_column("c").unwrap().1;
978996
assert!(!column.primary_key, "column 'c' shouldn't be a primary key");
979997
assert_eq!(
980-
vec!["a"],
981-
table.primary_key_column_names,
998+
vec![("a".to_string(), SortOrder::Asc)],
999+
table.primary_key_columns,
9821000
"primary key column names should be ['a']"
9831001
);
9841002
Ok(())
@@ -994,8 +1012,8 @@ mod tests {
9941012
let column = table.get_column("c").unwrap().1;
9951013
assert!(!column.primary_key, "column 'c' shouldn't be a primary key");
9961014
assert_eq!(
997-
vec!["a"],
998-
table.primary_key_column_names,
1015+
vec![("a".to_string(), SortOrder::Asc)],
1016+
table.primary_key_columns,
9991017
"primary key column names should be ['a']"
10001018
);
10011019
Ok(())
@@ -1143,7 +1161,7 @@ mod tests {
11431161
name: "t1".to_string(),
11441162
has_rowid: true,
11451163
is_strict: false,
1146-
primary_key_column_names: vec!["nonexistent".to_string()],
1164+
primary_key_columns: vec![("nonexistent".to_string(), SortOrder::Asc)],
11471165
columns: vec![Column {
11481166
name: Some("a".to_string()),
11491167
ty: Type::Integer,

‎core/storage/btree.rs

+59-14
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
use crate::{
2+
schema::Index,
23
storage::{
34
pager::Pager,
45
sqlite3_ondisk::{
@@ -7,6 +8,7 @@ use crate::{
78
},
89
},
910
translate::plan::IterationDirection,
11+
types::IndexKeySortOrder,
1012
MvCursor,
1113
};
1214

@@ -364,6 +366,7 @@ pub struct BTreeCursor {
364366
/// Reusable immutable record, used to allow better allocation strategy.
365367
reusable_immutable_record: RefCell<Option<ImmutableRecord>>,
366368
empty_record: Cell<bool>,
369+
pub index_key_sort_order: IndexKeySortOrder,
367370
}
368371

369372
impl BTreeCursor {
@@ -388,9 +391,22 @@ impl BTreeCursor {
388391
},
389392
reusable_immutable_record: RefCell::new(None),
390393
empty_record: Cell::new(true),
394+
index_key_sort_order: IndexKeySortOrder::default(),
391395
}
392396
}
393397

398+
pub fn new_index(
399+
mv_cursor: Option<Rc<RefCell<MvCursor>>>,
400+
pager: Rc<Pager>,
401+
root_page: usize,
402+
index: &Index,
403+
) -> Self {
404+
let index_key_sort_order = IndexKeySortOrder::from_index(index);
405+
let mut cursor = Self::new(mv_cursor, pager, root_page);
406+
cursor.index_key_sort_order = index_key_sort_order;
407+
cursor
408+
}
409+
394410
/// Check if the table is empty.
395411
/// This is done by checking if the root page has no cells.
396412
fn is_empty_table(&self) -> Result<CursorResult<bool>> {
@@ -547,8 +563,11 @@ impl BTreeCursor {
547563
let record_values = record.get_values();
548564
let record_slice_same_num_cols =
549565
&record_values[..index_key.get_values().len()];
550-
let order =
551-
compare_immutable(record_slice_same_num_cols, index_key.get_values());
566+
let order = compare_immutable(
567+
record_slice_same_num_cols,
568+
index_key.get_values(),
569+
self.index_key_sort_order,
570+
);
552571
order
553572
};
554573

@@ -602,8 +621,11 @@ impl BTreeCursor {
602621
let record_values = record.get_values();
603622
let record_slice_same_num_cols =
604623
&record_values[..index_key.get_values().len()];
605-
let order =
606-
compare_immutable(record_slice_same_num_cols, index_key.get_values());
624+
let order = compare_immutable(
625+
record_slice_same_num_cols,
626+
index_key.get_values(),
627+
self.index_key_sort_order,
628+
);
607629
order
608630
};
609631
let found = match op {
@@ -849,10 +871,18 @@ impl BTreeCursor {
849871
let SeekKey::IndexKey(index_key) = key else {
850872
unreachable!("index seek key should be a record");
851873
};
852-
let order = compare_immutable(
853-
&self.get_immutable_record().as_ref().unwrap().get_values(),
854-
index_key.get_values(),
855-
);
874+
let order = {
875+
let record = self.get_immutable_record();
876+
let record = record.as_ref().unwrap();
877+
let record_slice_same_num_cols =
878+
&record.get_values()[..index_key.get_values().len()];
879+
let order = compare_immutable(
880+
record_slice_same_num_cols,
881+
index_key.get_values(),
882+
self.index_key_sort_order,
883+
);
884+
order
885+
};
856886
let found = match op {
857887
SeekOp::GT => order.is_gt(),
858888
SeekOp::GE => order.is_ge(),
@@ -901,10 +931,18 @@ impl BTreeCursor {
901931
let SeekKey::IndexKey(index_key) = key else {
902932
unreachable!("index seek key should be a record");
903933
};
904-
let order = compare_immutable(
905-
&self.get_immutable_record().as_ref().unwrap().get_values(),
906-
index_key.get_values(),
907-
);
934+
let order = {
935+
let record = self.get_immutable_record();
936+
let record = record.as_ref().unwrap();
937+
let record_slice_same_num_cols =
938+
&record.get_values()[..index_key.get_values().len()];
939+
let order = compare_immutable(
940+
record_slice_same_num_cols,
941+
index_key.get_values(),
942+
self.index_key_sort_order,
943+
);
944+
order
945+
};
908946
let found = match op {
909947
SeekOp::GT => order.is_lt(),
910948
SeekOp::GE => order.is_le(),
@@ -1031,7 +1069,11 @@ impl BTreeCursor {
10311069
let record = record.as_ref().unwrap();
10321070
let record_slice_equal_number_of_cols =
10331071
&record.get_values().as_slice()[..index_key.get_values().len()];
1034-
let order = record_slice_equal_number_of_cols.cmp(index_key.get_values());
1072+
let order = compare_immutable(
1073+
record_slice_equal_number_of_cols,
1074+
index_key.get_values(),
1075+
self.index_key_sort_order,
1076+
);
10351077
let found = match op {
10361078
SeekOp::GT => order.is_gt(),
10371079
SeekOp::GE => order.is_ge(),
@@ -1278,6 +1320,7 @@ impl BTreeCursor {
12781320
let interior_cell_vs_index_key = compare_immutable(
12791321
record_slice_equal_number_of_cols,
12801322
index_key.get_values(),
1323+
self.index_key_sort_order,
12811324
);
12821325
// in sqlite btrees left child pages have <= keys.
12831326
// in general, in forwards iteration we want to find the first key that matches the seek condition.
@@ -1430,7 +1473,8 @@ impl BTreeCursor {
14301473
self.get_immutable_record()
14311474
.as_ref()
14321475
.unwrap()
1433-
.get_values()
1476+
.get_values(),
1477+
self.index_key_sort_order,
14341478
) == Ordering::Equal {
14351479

14361480
tracing::debug!("insert_into_page: found exact match with cell_idx={cell_idx}, overwriting");
@@ -3017,6 +3061,7 @@ impl BTreeCursor {
30173061
let order = compare_immutable(
30183062
key.to_index_key_values(),
30193063
self.get_immutable_record().as_ref().unwrap().get_values(),
3064+
self.index_key_sort_order,
30203065
);
30213066
match order {
30223067
Ordering::Less | Ordering::Equal => {

0 commit comments

Comments
 (0)