Skip to content

fix(storage): plan an LSI query as an index walk on PostgreSQL and SQLite #432

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
LeeroyHannigan wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(storage): plan an LSI query as an index walk on PostgreSQL and SQLite #432

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 5 additions & 3 deletions .github/workflows/integration.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -302,16 +302,18 @@ jobs:
# write-transaction conflict tests force a deadlock, and the backup restore
# cases need catalog rows the current binary would not write; all need
# direct access. The online GSI build tests reorder queue rows by hand and
# pause writers with a test hook. They build their own throwaway databases;
# the connection string is the server, not a database.
# pause writers with a test hook. The LSI query plan test reads the
# statement the engine sent from pg_stat_activity and explains it, which
# the wire cannot show. They build their own throwaway databases; the
# connection string is the server, not a database.
- name: Run PostgreSQL storage-level tests
env:
EXTENDDB_TEST_PG_CONNECTION_STRING: postgresql://postgres:devpass@127.0.0.1:5432
run: >-
cargo test --release -p extenddb-storage-postgres
--test vector_control_plane --test put_create_race --test twi_conflict
--test stream_shard_retention --test backup_restore
--test gsi_online_build --test gsi_stale_writer
--test gsi_online_build --test gsi_stale_writer --test lsi_query_plan

# The daemonized server logs to syslog; dump it so server-side failures
# are diagnosable from the job log.
Expand Down
14 changes: 5 additions & 9 deletions crates/storage-postgres/src/data/query.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,13 +24,12 @@ pub(crate) enum PaginationBinds {
None,
/// Index query where base table has no SK — only `base_pk` as tie-breaker.
BasePkOnly { pk_text: String },
/// LSI query. `base_sk` alone is a complete tie-breaker, because every
/// row shares the queried partition key.
BaseSkOnly { sk: SortKeyValue },
/// Full base primary key as tie-breaker. Required by a GSI query when the
/// base table has a SK: rows sharing an index SK can come from different
/// Full base primary key as tie-breaker. A GSI query binds the start key's
/// base partition key: rows sharing an index SK can come from different
/// base partitions and can also share a base SK, so `base_sk` alone does
/// not identify a row. Also used by a hash-only index on such a table.
/// not identify a row. An LSI query binds the queried partition key, which
/// is every row's `base_pk`. Also used by a hash-only index on a table with
/// a SK.
BasePkAndSk { pk_text: String, sk: SortKeyValue },
}

Expand Down Expand Up @@ -194,9 +193,6 @@ pub(crate) async fn execute_query_sql(
PaginationBinds::BasePkOnly { pk_text } => {
query = query.bind(pk_text.clone());
}
PaginationBinds::BaseSkOnly { sk } => {
query = bind_sk_value(query, sk);
}
PaginationBinds::BasePkAndSk { pk_text, sk } => {
query = query.bind(pk_text.clone());
query = bind_sk_value(query, sk);
Expand Down
130 changes: 69 additions & 61 deletions crates/storage-postgres/src/data/query_scan.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,10 +44,9 @@ fn build_pagination_where(
sk_info_val: Option<(&str, ScalarAttributeType)>,
base_sk_info: &Option<(String, ScalarAttributeType)>,
is_index: bool,
is_lsi: bool,
forward: bool,
) -> String {
let cols = order_by_columns(sk_info_val, base_sk_info.as_ref(), is_index, is_lsi);
let cols = order_by_columns(sk_info_val, base_sk_info.as_ref(), is_index);
if cols.is_empty() {
return String::new();
}
Expand All @@ -68,20 +67,21 @@ fn build_pagination_where(
/// ordered by the primary key clause the caller writes.
///
/// - Base table with a sort key: the sort key.
/// - LSI: the index sort key, then the base sort key (every row shares the
/// queried partition key, so the base sort key alone identifies a row).
/// - GSI with a sort key: the index sort key, then the FULL base primary key.
/// Rows in a GSI partition are unique on (index SK, base PK, base SK), not
/// on (index SK, base SK): many base partitions can project the same index
/// SK and the same base SK, and comparing base SK alone made a page-two
/// query return nothing whenever the rows sharing an index SK also shared a
/// base SK.
/// - Index with a sort key (GSI or LSI): the index sort key, then the FULL
/// base primary key. Rows in a GSI partition are unique on (index SK, base
/// PK, base SK), not on (index SK, base SK): many base partitions can
/// project the same index SK and the same base SK, and comparing base SK
/// alone made a page-two query return nothing whenever the rows sharing an
/// index SK also shared a base SK. In an LSI partition `base_pk` equals the
/// queried `pk` on every row, so it does not change the order, but it keeps
/// the `ORDER BY` identical to the ordering index `(pk, sk, base_pk,
/// base_sk)`. Without it the index serves only the index SK and PostgreSQL
/// sorts the rows that share one on every page.
/// - Hash-only index: the full base primary key.
fn order_by_columns(
sk_info_val: Option<(&str, ScalarAttributeType)>,
base_sk_info: Option<&(String, ScalarAttributeType)>,
is_index: bool,
is_lsi: bool,
) -> Vec<String> {
let collated = |col: &str, t: ScalarAttributeType| {
if t == ScalarAttributeType::S {
Expand All @@ -97,9 +97,7 @@ fn order_by_columns(
if !is_index {
return cols;
}
if !is_lsi {
cols.push("base_pk COLLATE \"C\"".to_owned());
}
cols.push("base_pk COLLATE \"C\"".to_owned());
cols.extend(base_sk);
} else if is_index {
cols.push("base_pk COLLATE \"C\"".to_owned());
Expand Down Expand Up @@ -226,7 +224,6 @@ impl PostgresEngine {
sk_info_val,
&base_sk_info,
index_name.is_some(),
is_lsi,
forward,
);
sql.push_str(&pagination_sql);
Expand All @@ -237,12 +234,7 @@ impl PostgresEngine {
// of the forward one (as DynamoDB's is) and the cursor predicate above
// is a row comparison the key index can seek on. String columns carry
// COLLATE "C" for DynamoDB's UTF-8 byte order.
let order_cols = order_by_columns(
sk_info_val,
base_sk_info.as_ref(),
index_name.is_some(),
is_lsi,
);
let order_cols = order_by_columns(sk_info_val, base_sk_info.as_ref(), index_name.is_some());
if !order_cols.is_empty() {
let dir = if forward { "ASC" } else { "DESC" };
let ordered = order_cols
Expand All @@ -266,19 +258,29 @@ impl PostgresEngine {
{
// Index that has its own SK, with a base-table tie-breaker. The
// index SK is bound separately (see execute_query_sql); the binds
// here supply the tie-breaker.
// here supply the full base primary key.
//
// An LSI needs the base SK alone, because every row shares the
// queried partition key. A GSI needs the full base primary key,
// because rows sharing an index SK can come from different base
// partitions and can also share a base SK.
// A GSI takes the base partition key from the start key, because
// rows sharing an index SK can come from different base
// partitions and can also share a base SK. An LSI binds the
// queried partition key, which every row of the partition
// carries as its base_pk. base_pk then compares equal on every
// row, and only the base SK separates a row from the cursor.
// The statement does not read the start key's partition key:
// the bind keeps the page inside the queried partition.
// DynamoDB refuses a start key that names another partition
// with a ValidationException; that check belongs to request
// validation, not to this statement.
let base_sk = start_key
.get(base_sk_name.as_str())
.map(|v| parse_sk(v, base_sk_type))
.transpose()?;
if is_lsi {
match base_sk {
Some(sk) => PaginationBinds::BaseSkOnly { sk },
Some(sk) => PaginationBinds::BasePkAndSk {
pk_text: pk_text.clone(),
sk,
},
None => PaginationBinds::None,
}
} else {
Expand Down Expand Up @@ -611,19 +613,19 @@ mod tests {
#[test]
fn base_table_cursor_is_a_single_comparison() {
assert_eq!(
build_pagination_where(3, Some(("sk", S)), &None, false, false, true),
build_pagination_where(3, Some(("sk", S)), &None, false, true),
" AND sk_s COLLATE \"C\" > $3"
);
assert_eq!(
build_pagination_where(3, Some(("sk", N)), &None, false, false, false),
build_pagination_where(3, Some(("sk", N)), &None, false, false),
" AND sk_n < $3"
);
}

#[test]
fn forward_gsi_cursor_is_a_row_comparison_over_the_full_base_key() {
assert_eq!(
build_pagination_where(3, Some(("gsk", S)), &base_s(), true, false, true),
build_pagination_where(3, Some(("gsk", S)), &base_s(), true, true),
" AND (sk_s COLLATE \"C\", base_pk COLLATE \"C\", base_sk_s COLLATE \"C\") > ($3, $4, $5)"
);
// Numeric index sort key with a numeric base sort key: no collation on either.
Expand All @@ -633,7 +635,6 @@ mod tests {
Some(("gsk", N)),
&Some(("bsk".to_owned(), N)),
true,
false,
true
),
" AND (sk_n, base_pk COLLATE \"C\", base_sk_n) > ($3, $4, $5)"
Expand All @@ -646,49 +647,57 @@ mod tests {
// the exact mirror of its forward one, including among items sharing
// an index sort key (measured 2026-09-17).
assert_eq!(
build_pagination_where(3, Some(("gsk", S)), &base_s(), true, false, false),
build_pagination_where(3, Some(("gsk", S)), &base_s(), true, false),
" AND (sk_s COLLATE \"C\", base_pk COLLATE \"C\", base_sk_s COLLATE \"C\") < ($3, $4, $5)"
);
}

#[test]
fn gsi_on_a_hash_only_base_table_uses_base_pk_as_the_tie_breaker() {
assert_eq!(
build_pagination_where(3, Some(("gsk", S)), &None, true, false, true),
build_pagination_where(3, Some(("gsk", S)), &None, true, true),
" AND (sk_s COLLATE \"C\", base_pk COLLATE \"C\") > ($3, $4)"
);
assert_eq!(
build_pagination_where(3, Some(("gsk", S)), &None, true, false, false),
build_pagination_where(3, Some(("gsk", S)), &None, true, false),
" AND (sk_s COLLATE \"C\", base_pk COLLATE \"C\") < ($3, $4)"
);
}

#[test]
fn lsi_cursor_is_a_row_comparison_in_both_directions() {
fn lsi_cursor_compares_the_ordering_index_columns_in_both_directions() {
// The ordering index is (pk, sk, base_pk, base_sk). An LSI cursor and
// ORDER BY that skip base_pk leave PostgreSQL an index that orders the
// index sort key only, so it sorts every group of rows sharing one.
assert_eq!(
build_pagination_where(3, Some(("lsk", S)), &base_s(), true, true, true),
" AND (sk_s COLLATE \"C\", base_sk_s COLLATE \"C\") > ($3, $4)"
build_pagination_where(3, Some(("lsk", S)), &base_s(), true, true),
" AND (sk_s COLLATE \"C\", base_pk COLLATE \"C\", base_sk_s COLLATE \"C\") > ($3, $4, $5)"
);
assert_eq!(
build_pagination_where(3, Some(("lsk", S)), &base_s(), true, false),
" AND (sk_s COLLATE \"C\", base_pk COLLATE \"C\", base_sk_s COLLATE \"C\") < ($3, $4, $5)"
);
assert_eq!(
build_pagination_where(3, Some(("lsk", S)), &base_s(), true, true, false),
" AND (sk_s COLLATE \"C\", base_sk_s COLLATE \"C\") < ($3, $4)"
order_by_columns(Some(("lsk", N)), Some(&("bsk".to_owned(), S)), true),
vec![
"sk_n".to_owned(),
"base_pk COLLATE \"C\"".to_owned(),
"base_sk_s COLLATE \"C\"".to_owned()
]
);
}

#[test]
fn hash_only_index_cursor_is_a_row_comparison_over_the_base_key() {
assert_eq!(
build_pagination_where(3, None, &base_s(), true, false, true),
build_pagination_where(3, None, &base_s(), true, true),
" AND (base_pk COLLATE \"C\", base_sk_s COLLATE \"C\") > ($3, $4)"
);
assert_eq!(
build_pagination_where(3, None, &None, true, false, true),
build_pagination_where(3, None, &None, true, true),
" AND base_pk COLLATE \"C\" > $3"
);
assert_eq!(
build_pagination_where(3, None, &None, false, false, true),
""
);
assert_eq!(build_pagination_where(3, None, &None, false, true), "");
}

#[test]
Expand All @@ -698,15 +707,15 @@ mod tests {
// after (p5, p4) asked for keys above p4, returned p5 again, and
// then ended with p0..p3 never delivered.
assert_eq!(
build_pagination_where(3, None, &None, true, false, false),
build_pagination_where(3, None, &None, true, false),
" AND base_pk COLLATE \"C\" < $3"
);
assert_eq!(
build_pagination_where(3, None, &base_s(), true, false, false),
build_pagination_where(3, None, &base_s(), true, false),
" AND (base_pk COLLATE \"C\", base_sk_s COLLATE \"C\") < ($3, $4)"
);
assert_eq!(
build_pagination_where(3, None, &Some(("bsk".to_owned(), N)), true, false, false),
build_pagination_where(3, None, &Some(("bsk".to_owned(), N)), true, false),
" AND (base_pk COLLATE \"C\", base_sk_n) < ($3, $4)"
);
}
Expand All @@ -719,21 +728,20 @@ mod tests {
Option<(&'a str, ScalarAttributeType)>,
Option<(String, ScalarAttributeType)>,
bool,
bool,
);
let shapes: [Shape<'_>; 7] = [
(Some(("sk", S)), None, false, false),
(Some(("gsk", S)), base_s(), true, false),
(Some(("gsk", N)), Some(("bsk".to_owned(), N)), true, false),
(Some(("gsk", S)), None, true, false),
(Some(("lsk", S)), base_s(), true, true),
(None, base_s(), true, false),
(None, None, true, false),
(Some(("sk", S)), None, false),
(Some(("gsk", S)), base_s(), true),
(Some(("gsk", N)), Some(("bsk".to_owned(), N)), true),
(Some(("gsk", S)), None, true),
(Some(("lsk", S)), base_s(), true),
(None, base_s(), true),
(None, None, true),
];
for (sk, base, is_index, is_lsi) in shapes {
let cols = order_by_columns(sk, base.as_ref(), is_index, is_lsi);
let fwd = build_pagination_where(1, sk, &base, is_index, is_lsi, true);
let rev = build_pagination_where(1, sk, &base, is_index, is_lsi, false);
for (sk, base, is_index) in shapes {
let cols = order_by_columns(sk, base.as_ref(), is_index);
let fwd = build_pagination_where(1, sk, &base, is_index, true);
let rev = build_pagination_where(1, sk, &base, is_index, false);
let params = (1..=cols.len())
.map(|i| format!("${i}"))
.collect::<Vec<_>>();
Expand All @@ -756,6 +764,6 @@ mod tests {
"reverse {sk:?} {base:?}"
);
}
assert!(order_by_columns(None, None, false, false).is_empty());
assert!(order_by_columns(None, None, false).is_empty());
}
}
Loading
Loading