diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index da6c251d3..5ac47253a 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -302,8 +302,10 @@ 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 # table_snapshot: the SnapshotEngine consistency check; it needs the # shared page-size hook, compiled only into the crate's own test @@ -316,7 +318,7 @@ jobs: cargo test --release -p extenddb-storage-postgres --test vector_control_plane --test put_create_race --test twi_conflict --test stream_shard_retention --test stream_shard_ownership --test stream_records_page_order --test backup_restore - --test gsi_online_build --test gsi_stale_writer --test table_snapshot + --test gsi_online_build --test gsi_stale_writer --test table_snapshot --test lsi_query_plan # The daemonized server logs to syslog; dump it so server-side failures # are diagnosable from the job log. diff --git a/crates/storage-postgres/src/data/query.rs b/crates/storage-postgres/src/data/query.rs index 6a84dc1b7..1a1cc06aa 100755 --- a/crates/storage-postgres/src/data/query.rs +++ b/crates/storage-postgres/src/data/query.rs @@ -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 }, } @@ -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); diff --git a/crates/storage-postgres/src/data/query_scan.rs b/crates/storage-postgres/src/data/query_scan.rs index c0f0c8b05..f3b6f3dce 100644 --- a/crates/storage-postgres/src/data/query_scan.rs +++ b/crates/storage-postgres/src/data/query_scan.rs @@ -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(); } @@ -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 { let collated = |col: &str, t: ScalarAttributeType| { if t == ScalarAttributeType::S { @@ -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()); @@ -226,7 +224,6 @@ impl PostgresEngine { sk_info_val, &base_sk_info, index_name.is_some(), - is_lsi, forward, ); sql.push_str(&pagination_sql); @@ -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 @@ -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 { @@ -611,11 +613,11 @@ 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" ); } @@ -623,7 +625,7 @@ mod tests { #[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. @@ -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)" @@ -646,7 +647,7 @@ 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)" ); } @@ -654,41 +655,49 @@ mod tests { #[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] @@ -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)" ); } @@ -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::>(); @@ -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()); } } diff --git a/crates/storage-postgres/tests/lsi_query_plan.rs b/crates/storage-postgres/tests/lsi_query_plan.rs new file mode 100644 index 000000000..b6e0e9c02 --- /dev/null +++ b/crates/storage-postgres/tests/lsi_query_plan.rs @@ -0,0 +1,608 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 +//! Storage-level tests for the plan PostgreSQL chooses for a Query on a local +//! secondary index. +//! +//! Every index table whose index has a sort key has an ordering index on +//! `(pk, sk, base_pk, base_sk)`. A Query whose `ORDER BY` is that index minus +//! the constant `pk` is a walk of the index that stops at the page limit. One +//! that skips a column of it gets a sort node over the index scan: an +//! `Incremental Sort` when the key condition leaves several index sort key +//! values, and a top-N `Sort` when `lsk = :l` leaves one. The rows come back in +//! the same order either way, so the wire suites cannot tell the two apart, +//! and these tests read the plan. +//! +//! Each test takes the statement the engine sent from `pg_stat_activity` and +//! explains that statement. It checks both plans PostgreSQL can run a prepared +//! statement with: the custom plan for the bound values, used for the first +//! executions on a connection, and the generic plan it may switch to after +//! them. +//! +//! Each test builds its own throwaway database, applies the shipped migrations +//! to it, and drops it when it passes. A failing test leaves its database behind +//! on purpose, named `eddb_lsip_*`, so the state that failed can be inspected. +//! +//! Requires `EXTENDDB_TEST_PG_CONNECTION_STRING`, a base URL with no database +//! component (for example `postgresql://postgres@127.0.0.1:5432`), pointing at a +//! server whose role may create and drop databases and read other sessions' +//! rows in `pg_stat_activity`. Without it every test here reports a skip and +//! passes, the same convention the wire suites use. PostgreSQL added +//! `EXPLAIN (GENERIC_PLAN)` in version 16; on an older server the test reports +//! a skip as well. + +use std::collections::{BTreeMap, HashMap}; + +use extenddb_core::expression::{ExpressionMaps, KeyCondition, parse_key_condition, tokenize}; +use extenddb_core::types::{ + AttributeDefinition, AttributeValue, BillingMode, CreateTableInput, Item, KeySchemaElement, + KeyType, LsiInput, Projection, ProjectionType, ScalarAttributeType, TableKeyInfo, +}; +use extenddb_storage::{DataEngine, TableEngine}; +use extenddb_storage_postgres::{PostgresConfig, PostgresEngine}; +use sqlx::postgres::PgPoolOptions; +use sqlx::{PgPool, Row}; + +const ACCOUNT: &str = "123456789012"; +const REGION: &str = "us-east-1"; +const TABLE: &str = "t_lsi_plan"; +/// Items in the queried partition, and how many of them share each index +/// sort key value. +const ITEMS: usize = 3_000; +const GROUP: usize = 50; +/// The Query Limit of a page, larger than a group so a page crosses from one +/// index sort key value to the next. +const PAGE: usize = 100; +/// The Query Limit of a page whose key condition fixes one index sort key +/// value (`lsk = :l`): smaller than a group, so the page ends inside it. When +/// the Limit covers the whole group, PostgreSQL may read the group through a +/// bitmap scan and sort it, which it estimates cheaper than an ordered index +/// scan that reads the same rows. +const GROUP_PAGE: usize = 10; + +struct Scratch { + engine: PostgresEngine, + db: PgPool, + admin: PgPool, + db_name: String, +} + +impl Scratch { + async fn cleanup(self) { + let Scratch { + engine, + db, + admin, + db_name, + } = self; + drop(engine); + db.close().await; + sqlx::query(&format!( + "DROP DATABASE IF EXISTS \"{db_name}\" WITH (FORCE)" + )) + .execute(&admin) + .await + .expect("drop the scratch database"); + admin.close().await; + } +} + +fn base_conn() -> Option { + let conn = std::env::var("EXTENDDB_TEST_PG_CONNECTION_STRING").ok()?; + (!conn.trim().is_empty()).then(|| conn.trim_end_matches('/').to_owned()) +} + +/// Every `.sql` file under `migrations/` and `data_migrations/`, each +/// directory in filename order, so a new migration is picked up here without +/// editing this file. +fn shipped_migrations() -> Vec { + let root = std::path::Path::new(env!("CARGO_MANIFEST_DIR")); + let mut out = Vec::new(); + for dir in ["migrations", "data_migrations"] { + let mut files: Vec<_> = std::fs::read_dir(root.join(dir)) + .expect("read the migrations directory") + .map(|e| e.expect("directory entry").path()) + .filter(|p| p.extension().is_some_and(|e| e == "sql")) + .collect(); + files.sort(); + for f in files { + out.push(std::fs::read_to_string(&f).expect("read a migration file")); + } + } + out +} + +fn skip(test: &str) { + eprintln!( + "SKIP {test}: EXTENDDB_TEST_PG_CONNECTION_STRING is not set, so there is no PostgreSQL \ + to build a scratch catalog in." + ); +} + +async fn scratch() -> Scratch { + let base = base_conn().expect("caller checks base_conn() first"); + let db_name = format!("eddb_lsip_{}", uuid::Uuid::new_v4().simple())[..24].to_owned(); + let admin = PgPoolOptions::new() + .max_connections(1) + .connect(&format!("{base}/postgres")) + .await + .expect("connect to the postgres maintenance database"); + sqlx::query(&format!("CREATE DATABASE \"{db_name}\"")) + .execute(&admin) + .await + .expect("create the scratch database"); + let url = format!("{base}/{db_name}"); + let db = PgPoolOptions::new() + .max_connections(2) + .connect(&url) + .await + .expect("connect to the scratch database"); + // One scratch database stands in for both the catalog and the data + // database, so every shipped migration file from both directories is + // applied to it, in filename order. + for sql in shipped_migrations() { + sqlx::raw_sql(&sql) + .execute(&db) + .await + .expect("apply a shipped migration"); + } + sqlx::query("UPDATE settings SET value = '0' WHERE key = 'control_plane_delay_seconds'") + .execute(&db) + .await + .expect("pin the control-plane delay to zero"); + sqlx::query("INSERT INTO accounts (account_id, account_name) VALUES ($1, $2)") + .bind(ACCOUNT) + .bind(format!("acct-{db_name}")) + .execute(&db) + .await + .expect("seed the account row"); + let engine = PostgresEngine::new( + &PostgresConfig { + connection_string: url, + pool_size: 10, + max_item_size_bytes: 400_000, + }, + REGION, + ) + .await + .expect("open a PostgresEngine on the scratch database"); + Scratch { + engine, + db, + admin, + db_name, + } +} + +fn s_attr(name: &str) -> AttributeDefinition { + AttributeDefinition { + attribute_name: name.to_owned(), + attribute_type: ScalarAttributeType::S, + } +} + +fn key(name: &str, key_type: KeyType) -> KeySchemaElement { + KeySchemaElement { + attribute_name: name.to_owned(), + key_type, + } +} + +fn sv(v: &str) -> AttributeValue { + AttributeValue::S(v.to_owned()) +} + +/// A (pk S, sk S) table with an LSI on (pk, lsk S). +fn lsi_input() -> CreateTableInput { + CreateTableInput { + table_name: TABLE.to_owned(), + key_schema: vec![key("pk", KeyType::Hash), key("sk", KeyType::Range)], + attribute_definitions: vec![s_attr("pk"), s_attr("sk"), s_attr("lsk")], + billing_mode: Some(BillingMode::PayPerRequest), + local_secondary_indexes: Some(vec![LsiInput { + index_name: "lsi1".to_owned(), + key_schema: vec![key("pk", KeyType::Hash), key("lsk", KeyType::Range)], + projection: Projection { + projection_type: ProjectionType::All, + non_key_attributes: None, + }, + }]), + ..Default::default() + } +} + +fn lsk(i: usize) -> String { + format!("l{:03}", i % (ITEMS / GROUP)) +} + +fn sk(i: usize) -> String { + format!("s{i:05}") +} + +/// Create the table and fill partition `p` with `ITEMS` items, `GROUP` per +/// index sort key value, plus a smaller partition `q`. Returns the key info a +/// Query on the LSI is given and the SQL name of the LSI table. +async fn lsi_table(s: &Scratch) -> (TableKeyInfo, String) { + s.engine + .create_table(ACCOUNT, lsi_input()) + .await + .expect("create the table"); + let base = s + .engine + .table_key_info(ACCOUNT, TABLE) + .await + .expect("read the key info"); + let maps = ExpressionMaps::new(HashMap::new(), HashMap::new()); + for (pk, count) in [("p", ITEMS), ("q", ITEMS / 10)] { + for i in 0..count { + let item: Item = BTreeMap::from([ + ("pk".to_owned(), sv(pk)), + ("sk".to_owned(), sv(&sk(i))), + ("lsk".to_owned(), sv(&lsk(i))), + ]); + s.engine + .put_item(&base, item, false, None, &maps, None) + .await + .expect("put an item"); + } + } + let index = s + .engine + .index_info(ACCOUNT, TABLE, "lsi1") + .await + .expect("read the index info"); + let sql_table = format!("\"_ddb_{}\"", index.index_id); + sqlx::query(&format!("ANALYZE {sql_table}")) + .execute(&s.db) + .await + .expect("analyze the index table"); + // The key info the engine passes for an index Query: the index key schema, + // with the base table's key schema for the tie-breaker columns. + let key_info = TableKeyInfo { + key_schema: index.key_schema.clone(), + base_key_schema: base.key_schema.clone(), + ..base + }; + (key_info, sql_table) +} + +/// A key condition on the LSI and its value map. `values` are the +/// `:name` placeholders, given without the colon: the engine strips it before +/// building the maps. +fn condition(expr: &str, values: &[(&str, &str)]) -> (KeyCondition, ExpressionMaps) { + let tokens = tokenize(expr).expect("tokenize the key condition"); + let cond = parse_key_condition(&tokens).expect("parse the key condition"); + let values = values + .iter() + .map(|(name, value)| ((*name).to_owned(), sv(value))) + .collect(); + (cond, ExpressionMaps::new(HashMap::new(), values)) +} + +/// The text of the most recent Query statement another session of this +/// database ran on `sql_table`. +async fn last_statement(db: &PgPool, sql_table: &str) -> String { + sqlx::query_scalar( + "SELECT query FROM pg_stat_activity \ + WHERE datname = current_database() AND pid <> pg_backend_pid() \ + AND strpos(query, $1) > 0 \ + ORDER BY query_start DESC LIMIT 1", + ) + .bind(format!("SELECT item_data FROM {sql_table}")) + .fetch_optional(db) + .await + .expect("read pg_stat_activity") + .expect("the engine's Query statement is visible in pg_stat_activity") +} + +/// The plan PostgreSQL would cache for `statement` after its first executions. +/// Sent over the simple query protocol: the extended protocol would make the +/// statement's `$n` placeholders parameters of the `EXPLAIN` itself. +async fn generic_plan(db: &PgPool, statement: &str) -> Vec { + let sql = format!("EXPLAIN (GENERIC_PLAN, COSTS OFF) {statement}"); + sqlx::raw_sql(&sql) + .fetch_all(db) + .await + .unwrap_or_else(|e| panic!("{sql}: {e}")) + .iter() + .map(|row| row.get::(0)) + .collect() +} + +/// The plan for `statement` with `values` bound, as the first executions get. +async fn custom_plan(db: &PgPool, statement: &str, values: &[&str]) -> Vec { + let mut conn = db.acquire().await.expect("a connection"); + let args = values + .iter() + .map(|v| format!("'{}'", v.replace('\'', "''"))) + .collect::>() + .join(", "); + // Simple query protocol, for the same reason as in `generic_plan`. + let sql = format!( + "PREPARE lsi_page AS {statement}; \ + EXPLAIN (COSTS OFF) EXECUTE lsi_page({args}); \ + DEALLOCATE lsi_page" + ); + sqlx::raw_sql(&sql) + .fetch_all(&mut *conn) + .await + .unwrap_or_else(|e| panic!("{sql}: {e}")) + .iter() + .map(|row| row.get::(0)) + .collect() +} + +/// The sort nodes of a plan: `Sort` and `Incremental Sort`. +fn sort_nodes(plan: &[String]) -> Vec<&str> { + plan.iter() + .map(|l| l.trim_start().trim_start_matches("->").trim()) + .filter(|node| { + *node == "Sort" || node.starts_with("Sort (") || node.starts_with("Incremental Sort") + }) + .collect() +} + +fn assert_index_walk(label: &str, statement: &str, plan: &[String]) { + let shown = format!("{statement}\n{}", plan.join("\n")); + eprintln!("{label}:\n{shown}\n"); + assert!( + sort_nodes(plan).is_empty(), + "{label}: the plan sorts the rows the index scan reads:\n{shown}" + ); + assert!( + plan.iter() + .any(|l| l.contains("Index Scan") && l.contains("_pk_sk_s_base_pk_base_sk_s_idx")), + "{label}: the plan does not scan the ordering index:\n{shown}" + ); + // Every predicate, the cursor included, is an index condition. A cursor + // the index cannot seek on comes back as a Filter under the scan, and the + // scan then reads and discards rows before the page starts. + assert!( + !plan.iter().any(|l| l.trim_start().starts_with("Filter:")), + "{label}: the plan filters rows after the index scan instead of seeking on them:\n{shown}" + ); +} + +/// One Query page through the engine; returns the page's (lsk, sk) keys and +/// the statement the engine sent. +async fn page( + s: &Scratch, + key_info: &TableKeyInfo, + sql_table: &str, + (cond, maps): &(KeyCondition, ExpressionMaps), + forward: bool, + start: Option<&Item>, + limit: usize, +) -> (Vec<(String, String)>, String) { + let limit = i64::try_from(limit).expect("a page limit fits in i64"); + let (items, _) = s + .engine + .query( + key_info, + cond, + maps, + forward, + Some(limit), + start, + Some("lsi1"), + ) + .await + .expect("query a page of the LSI"); + let text = |item: &Item, attr: &str| match item.get(attr) { + Some(AttributeValue::S(v)) => v.clone(), + other => panic!("{attr}: {other:?}"), + }; + let keys = items + .iter() + .map(|i| (text(i, "lsk"), text(i, "sk"))) + .collect(); + (keys, last_statement(&s.db, sql_table).await) +} + +/// The page a Query with Limit `limit` returns from `rows`, which are in +/// forward index order: the first `limit` rows past `cursor` in the Query's +/// direction, or from the first row in that direction when there is no +/// cursor. +fn expected_page( + rows: &[(String, String)], + forward: bool, + cursor: Option<&(String, String)>, + limit: usize, +) -> Vec<(String, String)> { + let mut page: Vec<(String, String)> = rows + .iter() + .filter(|row| match cursor { + None => true, + Some(c) if forward => *row > c, + Some(c) => *row < c, + }) + .cloned() + .collect(); + if !forward { + page.reverse(); + } + page.truncate(limit); + page +} + +/// Query one page, compare its rows with `expected`, and check that both +/// plans of the statement the engine sent walk the ordering index. `values` +/// are the statement's binds in placeholder order, for the custom plan. +#[allow(clippy::too_many_arguments)] +async fn check_page( + s: &Scratch, + key_info: &TableKeyInfo, + sql_table: &str, + label: &str, + cond: &(KeyCondition, ExpressionMaps), + forward: bool, + start: Option<&Item>, + limit: usize, + expected: &[(String, String)], + values: &[&str], +) { + let (keys, statement) = page(s, key_info, sql_table, cond, forward, start, limit).await; + assert_eq!(keys, expected, "{label}: rows"); + let generic = generic_plan(&s.db, &statement).await; + assert_index_walk(&format!("{label}, generic plan"), &statement, &generic); + let custom = custom_plan(&s.db, &statement, values).await; + assert_index_walk(&format!("{label}, custom plan"), &statement, &custom); +} + +#[tokio::test] +async fn lsi_query_plans_as_an_index_walk_in_both_directions() { + let test = "lsi_query_plans_as_an_index_walk_in_both_directions"; + if base_conn().is_none() { + return skip(test); + } + let s = scratch().await; + let version: i32 = sqlx::query_scalar("SELECT current_setting('server_version_num')::int") + .fetch_one(&s.db) + .await + .expect("read the server version"); + if version < 160_000 { + eprintln!( + "SKIP {test}: EXPLAIN (GENERIC_PLAN) needs PostgreSQL 16 or later, the server \ + reports server_version_num {version}." + ); + s.cleanup().await; + return; + } + let (key_info, sql_table) = lsi_table(&s).await; + + // Every item of partition p in index order: index sort key, then base + // sort key. + let mut order: Vec<(String, String)> = (0..ITEMS).map(|i| (lsk(i), sk(i))).collect(); + order.sort(); + + // A start key in the middle of a group of items sharing an index sort key. + let middle = order.len() / 2 + GROUP / 2; + let cursor = order[middle].clone(); + let (start_lsk, start_sk) = (cursor.0.as_str(), cursor.1.as_str()); + let start: Item = BTreeMap::from([ + ("pk".to_owned(), sv("p")), + ("lsk".to_owned(), sv(start_lsk)), + ("sk".to_owned(), sv(start_sk)), + ]); + + // Key conditions: the partition alone; one index sort key value, the + // start key's group; a range of index sort key values around it. + let whole = condition("pk = :v", &[("v", "p")]); + let one_value = condition("pk = :v AND lsk = :l", &[("v", "p"), ("l", start_lsk)]); + let group: Vec<(String, String)> = order + .iter() + .filter(|(l, _)| l == start_lsk) + .cloned() + .collect(); + let (low, high) = (lsk(20), lsk(40)); + assert!(low.as_str() < start_lsk && start_lsk < high.as_str()); + let range = condition( + "pk = :v AND lsk BETWEEN :a AND :b", + &[("v", "p"), ("a", low.as_str()), ("b", high.as_str())], + ); + let in_range: Vec<(String, String)> = order + .iter() + .filter(|(l, _)| *l >= low && *l <= high) + .cloned() + .collect(); + + for forward in [true, false] { + check_page( + &s, + &key_info, + &sql_table, + &format!("first page, forward={forward}"), + &whole, + forward, + None, + PAGE, + &expected_page(&order, forward, None, PAGE), + &["p"], + ) + .await; + check_page( + &s, + &key_info, + &sql_table, + &format!("lsk = :l, first page, forward={forward}"), + &one_value, + forward, + None, + GROUP_PAGE, + &expected_page(&group, forward, None, GROUP_PAGE), + &["p", start_lsk], + ) + .await; + } + + for forward in [true, false] { + // The binds the engine sends for a resumed LSI page: the partition + // key, the key condition's sort key values, the start key's index + // sort key, the partition key again as the base partition key, then + // the start key's base sort key. + let label = format!("resumed page, forward={forward}"); + let expected = expected_page(&order, forward, Some(&cursor), PAGE); + check_page( + &s, + &key_info, + &sql_table, + &label, + &whole, + forward, + Some(&start), + PAGE, + &expected, + &["p", start_lsk, "p", start_sk], + ) + .await; + check_page( + &s, + &key_info, + &sql_table, + &format!("lsk = :l, resumed page, forward={forward}"), + &one_value, + forward, + Some(&start), + GROUP_PAGE, + &expected_page(&group, forward, Some(&cursor), GROUP_PAGE), + &["p", start_lsk, start_lsk, "p", start_sk], + ) + .await; + check_page( + &s, + &key_info, + &sql_table, + &format!("lsk BETWEEN, resumed page, forward={forward}"), + &range, + forward, + Some(&start), + PAGE, + &expected_page(&in_range, forward, Some(&cursor), PAGE), + &["p", low.as_str(), high.as_str(), start_lsk, "p", start_sk], + ) + .await; + + // The statement does not read the start key's partition key: the + // cursor binds the queried partition key as base_pk, which keeps the + // page inside partition p, so a start key naming partition q returns + // the same page. DynamoDB refuses such a start key with a + // ValidationException; that check belongs to request validation. + let mut elsewhere = start.clone(); + elsewhere.insert("pk".to_owned(), sv("q")); + let (keys_elsewhere, _) = page( + &s, + &key_info, + &sql_table, + &whole, + forward, + Some(&elsewhere), + PAGE, + ) + .await; + assert_eq!( + keys_elsewhere, expected, + "{label}: a start key naming partition q moved the page boundary" + ); + } + + s.cleanup().await; +} diff --git a/crates/storage-sqlite/src/data/query_scan.rs b/crates/storage-sqlite/src/data/query_scan.rs index 585debd28..bc5e7d989 100644 --- a/crates/storage-sqlite/src/data/query_scan.rs +++ b/crates/storage-sqlite/src/data/query_scan.rs @@ -51,98 +51,20 @@ impl SqliteEngine { (data_table_name(&key_info.table_id), false) }; - // Partition key (composite via netstring when multi-HASH). - let pk_text = if key_condition.extra_pk_conditions.is_empty() { - pk_to_text(&resolve_expr_to_av(&key_condition.pk_value, maps)?)?.into_owned() - } else { - let mut parts = - vec![pk_to_text(&resolve_expr_to_av(&key_condition.pk_value, maps)?)?.into_owned()]; - for (_, value) in &key_condition.extra_pk_conditions { - parts.push(pk_to_text(&resolve_expr_to_av(value, maps)?)?.into_owned()); - } - encode_netstring_composite(&parts) - }; - - let sk_info_val = sk_info(&key_info.key_schema, &key_info.attribute_definitions); - let all_sks = all_sort_key_info(&key_info.key_schema, &key_info.attribute_definitions); - let base_sk_info: Option<(String, ScalarAttributeType)> = if index_name.is_some() { - sk_info(&key_info.base_key_schema, &key_info.attribute_definitions) - .map(|(n, t)| (n.to_owned(), t)) - } else { - None - }; - - let mut sql = format!("SELECT item_data FROM {ddb_table} WHERE pk = ?"); - let mut binds: Vec = vec![BoundValue::Text(pk_text)]; - - // Primary sort-key condition. - if let (Some(sk_cond), Some((_, sk_type))) = (&key_condition.sk_condition, sk_info_val) { - let (sk_sql, sk_binds) = - build_sk_sql_and_binds(sk_cond, sk_column(sk_type), sk_type, maps)?; - sql.push_str(&sk_sql); - for v in &sk_binds { - binds.push(sk_bound(v)); - } - } - - // Extra RANGE-key equality conditions (multi-RANGE schemas). - for (path, value) in &key_condition.extra_sk_conditions { - let Some(attr_name) = resolve_attr_name(path, maps) else { - continue; - }; - if let Some(pos) = all_sks.iter().position(|(n, _)| *n == attr_name) - && pos > 0 - { - let (_, sk_type) = all_sks[pos]; - let _ = write!(sql, " AND {} = ?", sk_column_n(pos, sk_type)); - binds.push(sk_bound(&parse_sk( - &resolve_expr_to_av(value, maps)?, - sk_type, - )?)); - } - } - - // Pagination. - if exclusive_start_key.is_some() && sk_info_val.is_none() && index_name.is_none() { - return Ok((Vec::new(), None)); - } - if let Some(start_key) = exclusive_start_key { - append_query_pagination( - &mut sql, - &mut binds, - start_key, - sk_info_val, - base_sk_info.as_ref(), - key_info, - index_name.is_some(), - is_lsi, - forward, - )?; - } - - // ORDER BY: the same columns the page cursor compares, every one in the - // direction of ScanIndexForward, so a reverse Query is the exact mirror - // of the forward one (as DynamoDB's is) and the cursor predicate is a - // row comparison over the key index. - let order_cols = query_order_columns( - sk_info_val, - base_sk_info.as_ref(), - index_name.is_some(), + let Some((sql, binds)) = query_statement( + &ddb_table, is_lsi, - ); - if !order_cols.is_empty() { - let dir = if forward { "ASC" } else { "DESC" }; - let ordered = order_cols - .iter() - .map(|c| format!("{c} {dir}")) - .collect::>() - .join(", "); - let _ = write!(sql, " ORDER BY {ordered}"); - } - - let fetch_limit = limit.map_or(1_000_001, |l| l + 1); - let _ = write!(sql, " LIMIT {fetch_limit}"); - + key_info, + key_condition, + maps, + forward, + limit, + exclusive_start_key, + index_name.is_some(), + )? + else { + return Ok((Vec::new(), None)); + }; let rows = execute_dynamic_query(&sql, binds, &self.pool).await?; finalize(rows, limit, &key_info.key_schema) } @@ -296,6 +218,114 @@ impl SqliteEngine { } } +/// The statement and binds for one Query page on `ddb_table` (a data table or +/// an index table). `None` when the page is known to be empty without a read: +/// a start key on a hash-only base table, whose partition holds one item. +/// +/// Separate from [`SqliteEngine::query_impl`] so the plan tests can run +/// `EXPLAIN QUERY PLAN` on exactly the statement a Query sends. +#[allow(clippy::too_many_arguments)] +fn query_statement( + ddb_table: &str, + is_lsi: bool, + key_info: &TableKeyInfo, + key_condition: &KeyCondition, + maps: &ExpressionMaps, + forward: bool, + limit: Option, + exclusive_start_key: Option<&Item>, + is_index: bool, +) -> Result)>, StorageError> { + // Partition key (composite via netstring when multi-HASH). + let pk_text = if key_condition.extra_pk_conditions.is_empty() { + pk_to_text(&resolve_expr_to_av(&key_condition.pk_value, maps)?)?.into_owned() + } else { + let mut parts = + vec![pk_to_text(&resolve_expr_to_av(&key_condition.pk_value, maps)?)?.into_owned()]; + for (_, value) in &key_condition.extra_pk_conditions { + parts.push(pk_to_text(&resolve_expr_to_av(value, maps)?)?.into_owned()); + } + encode_netstring_composite(&parts) + }; + + let sk_info_val = sk_info(&key_info.key_schema, &key_info.attribute_definitions); + let all_sks = all_sort_key_info(&key_info.key_schema, &key_info.attribute_definitions); + let base_sk_info: Option<(String, ScalarAttributeType)> = if is_index { + sk_info(&key_info.base_key_schema, &key_info.attribute_definitions) + .map(|(n, t)| (n.to_owned(), t)) + } else { + None + }; + + let mut sql = format!("SELECT item_data FROM {ddb_table} WHERE pk = ?"); + let mut binds: Vec = vec![BoundValue::Text(pk_text.clone())]; + + // Primary sort-key condition. + if let (Some(sk_cond), Some((_, sk_type))) = (&key_condition.sk_condition, sk_info_val) { + let (sk_sql, sk_binds) = + build_sk_sql_and_binds(sk_cond, sk_column(sk_type), sk_type, maps)?; + sql.push_str(&sk_sql); + for v in &sk_binds { + binds.push(sk_bound(v)); + } + } + + // Extra RANGE-key equality conditions (multi-RANGE schemas). + for (path, value) in &key_condition.extra_sk_conditions { + let Some(attr_name) = resolve_attr_name(path, maps) else { + continue; + }; + if let Some(pos) = all_sks.iter().position(|(n, _)| *n == attr_name) + && pos > 0 + { + let (_, sk_type) = all_sks[pos]; + let _ = write!(sql, " AND {} = ?", sk_column_n(pos, sk_type)); + binds.push(sk_bound(&parse_sk( + &resolve_expr_to_av(value, maps)?, + sk_type, + )?)); + } + } + + // Pagination. + if exclusive_start_key.is_some() && sk_info_val.is_none() && !is_index { + return Ok(None); + } + if let Some(start_key) = exclusive_start_key { + append_query_pagination( + &mut sql, + &mut binds, + start_key, + &pk_text, + sk_info_val, + base_sk_info.as_ref(), + key_info, + is_index, + is_lsi, + forward, + )?; + } + + // ORDER BY: the same columns the page cursor compares, every one in the + // direction of ScanIndexForward, so a reverse Query is the exact mirror + // of the forward one (as DynamoDB's is) and the cursor predicate is a + // row comparison over the key index. + let order_cols = query_order_columns(sk_info_val, base_sk_info.as_ref(), is_index); + if !order_cols.is_empty() { + let dir = if forward { "ASC" } else { "DESC" }; + let ordered = order_cols + .iter() + .map(|c| format!("{c} {dir}")) + .collect::>() + .join(", "); + let _ = write!(sql, " ORDER BY {ordered}"); + } + + let fetch_limit = limit.map_or(1_000_001, |l| l + 1); + let _ = write!(sql, " LIMIT {fetch_limit}"); + Ok(Some((sql, binds))) +} + /// Resolve a key-condition path's attribute name, handling `#name` references. fn resolve_attr_name(path: &[PathElement], maps: &ExpressionMaps) -> Option { match path.first() { @@ -411,20 +441,21 @@ fn base_pk_from_start_key( /// hash-only base table, which has no order to resume within a partition. /// /// - 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 `idx_order_*` on `(pk, sk, base_pk, +/// base_sk)`. Without it SQLite reads the index in index sort key order +/// only and sorts the rest in a temporary B-tree. /// - Hash-only index: the full base primary key. fn query_order_columns( sk_info_val: Option<(&str, ScalarAttributeType)>, base_sk_info: Option<&(String, ScalarAttributeType)>, is_index: bool, - is_lsi: bool, ) -> Vec { let base_sk = base_sk_info.map(|(_, t)| format!("base_{}", sk_column(*t))); let mut cols = Vec::with_capacity(3); @@ -433,9 +464,7 @@ fn query_order_columns( if !is_index { return cols; } - if !is_lsi { - cols.push("base_pk".to_owned()); - } + cols.push("base_pk".to_owned()); cols.extend(base_sk); } else if is_index { cols.push("base_pk".to_owned()); @@ -457,11 +486,21 @@ fn query_order_columns( /// a hash-only index partition (measured 2026-09-17). Where the tie-breaker /// was the whole order (a hash-only index) an ascending cursor under a /// descending `ORDER BY` repeated the first page's items and skipped the rest. +/// +/// `partition_key` is the queried partition key text. An LSI binds it as the +/// cursor's `base_pk`, because every row of the partition carries it there. +/// `base_pk` then compares equal on every row, and only the base sort key +/// 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. A GSI binds the start key's base partition key. #[allow(clippy::too_many_arguments)] fn append_query_pagination( sql: &mut String, binds: &mut Vec, start_key: &Item, + partition_key: &str, sk_info_val: Option<(&str, ScalarAttributeType)>, base_sk_info: Option<&(String, ScalarAttributeType)>, key_info: &TableKeyInfo, @@ -469,7 +508,7 @@ fn append_query_pagination( is_lsi: bool, forward: bool, ) -> Result<(), StorageError> { - let cols = query_order_columns(sk_info_val, base_sk_info, is_index, is_lsi); + let cols = query_order_columns(sk_info_val, base_sk_info, is_index); if cols.is_empty() { return Ok(()); } @@ -487,11 +526,11 @@ fn append_query_pagination( // The sort key first; a base table's cursor is the sort key alone. cursor.push(sk_bind(sk_name, sk_type)?); if is_index { - if !is_lsi { - cursor.push(BoundValue::Text(base_pk_from_start_key( - start_key, key_info, - )?)); - } + cursor.push(BoundValue::Text(if is_lsi { + partition_key.to_owned() + } else { + base_pk_from_start_key(start_key, key_info)? + })); if let Some((base_name, base_type)) = base_sk_info { cursor.push(sk_bind(base_name, *base_type)?); } @@ -620,6 +659,19 @@ mod tests { is_lsi: bool, forward: bool, key: &Item, + ) -> (String, Vec) { + paginate_in(sk, base_sk, is_index, is_lsi, forward, key, "p") + } + + /// [`paginate`] with the queried partition key given explicitly. + fn paginate_in( + sk: Option<(&str, ScalarAttributeType)>, + base_sk: Option, + is_index: bool, + is_lsi: bool, + forward: bool, + key: &Item, + partition_key: &str, ) -> (String, Vec) { let ki = key_info(base_sk); let base_sk_info = base_sk.map(|t| ("bsk".to_owned(), t)); @@ -629,6 +681,7 @@ mod tests { &mut sql, &mut binds, key, + partition_key, sk, base_sk_info.as_ref(), &ki, @@ -699,24 +752,46 @@ mod tests { } #[test] - fn lsi_cursor_is_the_index_sort_key_then_the_base_sort_key() { + fn lsi_cursor_compares_the_ordering_index_columns() { + // idx_order_* is (pk, sk, base_pk, base_sk). A cursor and ORDER BY + // that skip base_pk leave SQLite an index that orders the index sort + // key only, and it sorts the rest in a temporary B-tree. let key = start_key(&[("lsk", "l"), ("bpk", "p"), ("bsk", "s")]); + let want_binds = vec!["l".to_owned(), "p".to_owned(), "s".to_owned()]; assert_eq!( paginate(Some(("lsk", S)), Some(S), true, true, true, &key), ( - " AND (sk_s, base_sk_s) > (?, ?)".to_owned(), - vec!["l".to_owned(), "s".to_owned()] + " AND (sk_s, base_pk, base_sk_s) > (?, ?, ?)".to_owned(), + want_binds.clone() ) ); assert_eq!( paginate(Some(("lsk", S)), Some(S), true, true, false, &key), ( - " AND (sk_s, base_sk_s) < (?, ?)".to_owned(), - vec!["l".to_owned(), "s".to_owned()] + " AND (sk_s, base_pk, base_sk_s) < (?, ?, ?)".to_owned(), + want_binds ) ); } + #[test] + fn lsi_cursor_binds_the_queried_partition_key_as_base_pk() { + // Every row of an LSI partition stores the queried partition key as + // base_pk, and the cursor binds that value, not the start key's + // partition key, so the page stays inside the queried partition. A + // GSI binds the start key's base partition key. + let key = start_key(&[("lsk", "l"), ("bpk", "other"), ("bsk", "s")]); + assert_eq!( + paginate_in(Some(("lsk", S)), Some(S), true, true, true, &key, "p").1, + vec!["l".to_owned(), "p".to_owned(), "s".to_owned()] + ); + let key = start_key(&[("gsk", "g"), ("bpk", "other"), ("bsk", "s")]); + assert_eq!( + paginate_in(Some(("gsk", S)), Some(S), true, false, true, &key, "g").1, + vec!["g".to_owned(), "other".to_owned(), "s".to_owned()] + ); + } + #[test] fn hash_only_index_cursor_follows_the_direction() { // The defect: a reverse Query on a hash-only index ordered @@ -789,7 +864,7 @@ mod tests { ]) }; let base_sk_info = base_sk.map(|t| ("bsk".to_owned(), t)); - let cols = query_order_columns(sk, base_sk_info.as_ref(), is_index, is_lsi); + let cols = query_order_columns(sk, base_sk_info.as_ref(), is_index); let marks = vec!["?"; cols.len()].join(", "); let (cols_sql, marks_sql) = if cols.len() == 1 { (cols[0].clone(), "?".to_owned()) @@ -812,9 +887,323 @@ mod tests { assert_eq!(fwd_binds, rev_binds); } // A hash-only base table has nothing to resume within a partition. - assert!(query_order_columns(None, None, false, false).is_empty()); + assert!(query_order_columns(None, None, false).is_empty()); let (sql, binds) = paginate(None, None, false, false, true, &start_key(&[("bpk", "p")])); assert!(sql.is_empty()); assert!(binds.is_empty()); } } + +/// The plan SQLite chooses for the statement an index Query sends. +/// +/// Each index table whose index has a sort key has an ordering index +/// `idx_order_*` on `(pk, sk*, base_pk, base_sk*)`. A Query whose `ORDER BY` +/// is that index minus the constant `pk` is a walk of the index; one that +/// skips a column of it makes SQLite sort the rows the walk reads, shown as +/// `USE TEMP B-TREE FOR ... ORDER BY` in `EXPLAIN QUERY PLAN`. The rows come +/// back in the same order either way, so these tests read the plan. +#[cfg(test)] +mod plan_tests { + use super::*; + use extenddb_core::expression::{parse_key_condition, tokenize}; + use extenddb_core::types::{AttributeValue, CreateTableInput, IndexType}; + use extenddb_storage::DataEngine; + use serde_json::json; + use std::collections::{BTreeMap, HashMap}; + + const ACCOUNT: &str = "000000000000"; + const TABLE: &str = "lsi-plan"; + + async fn engine() -> SqliteEngine { + let engine = SqliteEngine::new(":memory:", 2, "us-east-1", 409_600) + .await + .expect("engine"); + crate::schema::apply(&engine.pool).await.expect("schema"); + sqlx::query( + "INSERT INTO accounts (account_id, account_name) VALUES ('000000000000', 'default')", + ) + .execute(&engine.pool) + .await + .expect("account"); + // Tables become ACTIVE at create time: no transition poller runs here. + sqlx::query( + "INSERT OR REPLACE INTO settings (key, value) VALUES ('control_plane_delay_seconds', '0')", + ) + .execute(&engine.pool) + .await + .expect("settings"); + engine + } + + /// A (pk S, sk S) table with an LSI on (pk, lsk S) and a GSI on (gpk S, + /// gsk S), and 300 items in partition `p` sharing 30 index sort key values. + async fn table(engine: &SqliteEngine) -> TableKeyInfo { + let input: CreateTableInput = serde_json::from_value(json!({ + "TableName": TABLE, + "KeySchema": [ + {"AttributeName": "pk", "KeyType": "HASH"}, + {"AttributeName": "sk", "KeyType": "RANGE"} + ], + "AttributeDefinitions": [ + {"AttributeName": "pk", "AttributeType": "S"}, + {"AttributeName": "sk", "AttributeType": "S"}, + {"AttributeName": "lsk", "AttributeType": "S"}, + {"AttributeName": "gpk", "AttributeType": "S"}, + {"AttributeName": "gsk", "AttributeType": "S"} + ], + "LocalSecondaryIndexes": [{ + "IndexName": "lsi1", + "KeySchema": [ + {"AttributeName": "pk", "KeyType": "HASH"}, + {"AttributeName": "lsk", "KeyType": "RANGE"} + ], + "Projection": {"ProjectionType": "ALL"} + }], + "GlobalSecondaryIndexes": [{ + "IndexName": "gsi1", + "KeySchema": [ + {"AttributeName": "gpk", "KeyType": "HASH"}, + {"AttributeName": "gsk", "KeyType": "RANGE"} + ], + "Projection": {"ProjectionType": "ALL"} + }], + "BillingMode": "PAY_PER_REQUEST" + })) + .expect("input"); + engine + .create_table_impl(ACCOUNT, input, false) + .await + .expect("create table"); + let key_info = engine + .fetch_table_key_info(ACCOUNT, TABLE) + .await + .expect("key info"); + let maps = ExpressionMaps::new(HashMap::new(), HashMap::new()); + for i in 0..300 { + let s = |v: String| AttributeValue::S(v); + let item: Item = BTreeMap::from([ + ("pk".to_owned(), s("p".to_owned())), + ("sk".to_owned(), s(format!("s{i:04}"))), + ("lsk".to_owned(), s(format!("l{:02}", i % 30))), + ("gpk".to_owned(), s("g".to_owned())), + ("gsk".to_owned(), s(format!("l{:02}", i % 30))), + ]); + engine + .put_item(&key_info, item, false, None, &maps, None) + .await + .expect("put an item"); + } + key_info + } + + /// The table's key info as the engine passes it for a Query on `index`, + /// with the index table name and whether the index is an LSI. + async fn index_query_target( + engine: &SqliteEngine, + base: &TableKeyInfo, + index: &str, + ) -> (TableKeyInfo, String, bool) { + let info = engine + .fetch_index_info_by_table_id(&base.table_id, index) + .await + .expect("index info"); + let key_info = TableKeyInfo { + key_schema: info.key_schema.clone(), + base_key_schema: base.key_schema.clone(), + ..base.clone() + }; + ( + key_info, + index_table_name(&info.index_id), + info.index_type == IndexType::Lsi, + ) + } + + fn partition(attr: &str, value: &str) -> (KeyCondition, ExpressionMaps) { + condition(&format!("{attr} = :v"), &[("v", value)]) + } + + /// A key condition and its value map; `values` name the placeholders + /// without the colon, as the engine passes them. + fn condition(expr: &str, values: &[(&str, &str)]) -> (KeyCondition, ExpressionMaps) { + let tokens = tokenize(expr).expect("tokenize"); + let cond = parse_key_condition(&tokens).expect("parse"); + let values = values + .iter() + .map(|(name, value)| ((*name).to_owned(), AttributeValue::S((*value).to_owned()))) + .collect(); + (cond, ExpressionMaps::new(HashMap::new(), values)) + } + + fn s_key(pairs: &[(&str, &str)]) -> Item { + pairs + .iter() + .map(|(k, v)| ((*k).to_owned(), AttributeValue::S((*v).to_owned()))) + .collect() + } + + /// `EXPLAIN QUERY PLAN` detail lines for a statement and its binds. + async fn plan(engine: &SqliteEngine, sql: &str, binds: Vec) -> Vec { + let explain = format!("EXPLAIN QUERY PLAN {sql}"); + let mut query = sqlx::query_as::<_, (i64, i64, i64, String)>(&explain); + for bound in binds { + query = crate::data::bind_bound!(query, bound); + } + query + .fetch_all(&engine.pool) + .await + .expect("explain the statement") + .into_iter() + .map(|(_, _, _, detail)| detail) + .collect() + } + + /// Plan the Query statement for one page and assert it is a walk of the + /// ordering index with no sort step. + #[allow(clippy::too_many_arguments)] + async fn assert_index_walk( + engine: &SqliteEngine, + label: &str, + key_info: &TableKeyInfo, + ddb_table: &str, + is_lsi: bool, + cond: &KeyCondition, + maps: &ExpressionMaps, + forward: bool, + start: Option<&Item>, + ) { + let (sql, binds) = query_statement( + ddb_table, + is_lsi, + key_info, + cond, + maps, + forward, + Some(100), + start, + true, + ) + .expect("build the statement") + .expect("an index Query always reads"); + let plan = plan(engine, &sql, binds).await; + let shown = format!("{sql}\n {}", plan.join("\n ")); + assert!( + !plan.iter().any(|l| l.contains("TEMP B-TREE")), + "{label}: the plan sorts in a temporary B-tree:\n{shown}" + ); + assert!( + plan.iter() + .any(|l| l.contains("USING INDEX") && l.contains("idx_order_")), + "{label}: the plan does not walk the ordering index:\n{shown}" + ); + } + + #[tokio::test] + async fn lsi_query_walks_the_ordering_index_without_a_sort() { + let engine = engine().await; + let base = table(&engine).await; + let (lsi, lsi_table, is_lsi) = index_query_target(&engine, &base, "lsi1").await; + assert!(is_lsi); + // The partition alone, one index sort key value, and a range of them. + let shapes = [ + ("pk only", partition("pk", "p")), + ( + "lsk = :l", + condition("pk = :v AND lsk = :l", &[("v", "p"), ("l", "l07")]), + ), + ( + "lsk BETWEEN", + condition( + "pk = :v AND lsk BETWEEN :a AND :b", + &[("v", "p"), ("a", "l05"), ("b", "l10")], + ), + ), + ]; + // A start key inside a group of rows sharing an index sort key value. + let start = s_key(&[("pk", "p"), ("lsk", "l07"), ("sk", "s0127")]); + for (shape, (cond, maps)) in &shapes { + for forward in [true, false] { + for start in [None, Some(&start)] { + let label = + format!("LSI {shape} forward={forward} resumed={}", start.is_some()); + assert_index_walk( + &engine, &label, &lsi, &lsi_table, true, cond, maps, forward, start, + ) + .await; + } + } + } + } + + #[tokio::test] + async fn gsi_query_walks_the_ordering_index_without_a_sort() { + let engine = engine().await; + let base = table(&engine).await; + let (gsi, gsi_table, is_lsi) = index_query_target(&engine, &base, "gsi1").await; + assert!(!is_lsi); + let (cond, maps) = partition("gpk", "g"); + let start = s_key(&[("gpk", "g"), ("gsk", "l07"), ("pk", "p"), ("sk", "s0127")]); + for forward in [true, false] { + for start in [None, Some(&start)] { + let label = format!("GSI forward={forward} resumed={}", start.is_some()); + assert_index_walk( + &engine, &label, &gsi, &gsi_table, false, &cond, &maps, forward, start, + ) + .await; + } + } + } + + /// The pages of an LSI Query in both directions: ordered by the index sort + /// key, then the base sort key, with each page resuming after the last row + /// of the one before, including inside a group of equal index sort keys. + #[tokio::test] + async fn lsi_query_pages_follow_the_index_order_in_both_directions() { + let engine = engine().await; + let base = table(&engine).await; + let (lsi, _, _) = index_query_target(&engine, &base, "lsi1").await; + let (cond, maps) = partition("pk", "p"); + let mut want: Vec<(String, String)> = (0..300) + .map(|i| (format!("l{:02}", i % 30), format!("s{i:04}"))) + .collect(); + want.sort(); + for forward in [true, false] { + let mut got = Vec::new(); + let mut start: Option = None; + loop { + let (items, last) = engine + .query_impl( + &lsi, + &cond, + &maps, + forward, + Some(7), + start.as_ref(), + Some("lsi1"), + ) + .await + .expect("query a page"); + let text = |item: &Item, attr: &str| match item.get(attr) { + Some(AttributeValue::S(v)) => v.clone(), + other => panic!("{attr}: {other:?}"), + }; + got.extend(items.iter().map(|i| (text(i, "lsk"), text(i, "sk")))); + if last.is_none() { + break; + } + // The engine adds the base sort key to an index page's key. + let tail = items.last().expect("a page with a key has items"); + start = Some(s_key(&[ + ("pk", "p"), + ("lsk", &text(tail, "lsk")), + ("sk", &text(tail, "sk")), + ])); + } + let mut expected = want.clone(); + if !forward { + expected.reverse(); + } + assert_eq!(got, expected, "forward={forward}"); + } + } +} diff --git a/docs/design/04-component-storage.md b/docs/design/04-component-storage.md index dcfda3a3d..0815452e5 100755 --- a/docs/design/04-component-storage.md +++ b/docs/design/04-component-storage.md @@ -752,6 +752,25 @@ WHERE pk = $gsi_pk OR (base_pk = $base_pk AND base_sk_n > $base_sk)))) ``` +**LSI Query:** an LSI cursor compares the same three columns as a GSI cursor, +the index sort key, then `base_pk`, then the base sort key, and the `ORDER BY` +lists them in the same order and direction. Every row of an LSI partition +stores the queried partition key as `base_pk`, and the cursor binds the queried +partition key for it, so `base_pk` compares equal on every row and the order +and page boundaries are those of the index sort key and the base sort key. +Including `base_pk` makes the `ORDER BY` match the ordering index +`(pk, sk_s, base_pk, base_sk_n)` after the fixed `pk`, so the index can return +the page in order without a sort step. With `(sk_s, base_sk_n)` alone the index +orders `sk_s` only, and the database sorts the rows that share an index sort +key value. PostgreSQL adds an Incremental Sort when the key condition leaves +several index sort key values and a Sort when an equality on the index sort key +leaves one; SQLite adds a temporary B-tree. +```sql +WHERE pk = $lsi_pk + AND (sk_s COLLATE "C", base_pk COLLATE "C", base_sk_n) > ($lsi_sk, $lsi_pk, $base_sk) +ORDER BY sk_s COLLATE "C" ASC, base_pk COLLATE "C" ASC, base_sk_n ASC +``` + For GSI queries, the pagination key includes both the GSI key attributes and the base table primary key (needed to uniquely identify the position, since GSI keys are not unique). This is why the GSI PostgreSQL table includes `base_pk` and `base_sk_*` as actual columns.