Skip to content
6 changes: 0 additions & 6 deletions crates/storage-cassandra/src/data/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -49,12 +49,6 @@ enum PaginationBinds {
BasePkAndSk { pk_text: String, sk: SortKeyValue },
}

/// Compute upper bound for `begins_with` on strings.
/// Appends the maximum Unicode codepoint to create an exclusive upper bound.
pub(super) fn string_upper_bound(prefix: &str) -> String {
format!("{prefix}\u{10FFFF}")
}

/// Compute upper bound for `begins_with` on binary data.
/// Returns `None` when the prefix is empty or all-0xFF (no finite upper bound).
pub(super) fn binary_upper_bound(prefix: &[u8]) -> Option<Vec<u8>> {
Expand Down
74 changes: 63 additions & 11 deletions crates/storage-cassandra/src/data/sk_range.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,9 @@ use std::cmp::Ordering;
use extenddb_core::expression::{CompareOp, ExpressionMaps, SortKeyCondition};
use extenddb_core::types::ScalarAttributeType;
use extenddb_storage::error::StorageError;
use extenddb_storage::util::{SortKeyValue, parse_sk};
use extenddb_storage::util::{SortKeyValue, parse_sk, string_prefix_upper_bound};

use super::query::{binary_upper_bound, string_upper_bound};
use super::query::binary_upper_bound;
use super::resolve_expr_to_av;

/// One side of a sort key range: the CQL operator and the bound value.
Expand Down Expand Up @@ -66,10 +66,10 @@ impl SkRange {
SortKeyCondition::BeginsWith { prefix, .. } => {
let prefix_av = resolve_expr_to_av(prefix, expression_maps)?;
let prefix_sk = parse_sk(&prefix_av, sk_type)?;
// Exclusive upper bound by type; None when the prefix is empty
// or all 0xFF (no finite upper bound).
// Exclusive upper bound by type; None when the prefix is empty,
// all U+10FFFF or all 0xFF (no finite upper bound).
let upper_sk_opt = match &prefix_sk {
SortKeyValue::S(s) => Some(SortKeyValue::S(string_upper_bound(s))),
SortKeyValue::S(s) => string_prefix_upper_bound(s).map(SortKeyValue::S),
SortKeyValue::B(b) => binary_upper_bound(b).map(SortKeyValue::B),
SortKeyValue::N(_) => {
return Err(StorageError::Validation(
Expand Down Expand Up @@ -261,11 +261,57 @@ mod tests {
);
}

#[test]
fn begins_with_range_admits_every_key_with_the_prefix_and_no_other() {
// Sort keys measured on the service on 2026-10-25: begins_with(sk, p)
// returns exactly the keys that start with p, in UTF-8 byte order.
let keys = [
"a",
"ab",
"a\u{D7FF}",
"a\u{E000}",
"a\u{FFFF}",
"a\u{10000}",
"a\u{10FFFF}",
"a\u{10FFFF}b",
"a\u{10FFFF}\u{10FFFF}",
"b",
"\u{10FFFF}",
"\u{10FFFF}x",
"\u{10FFFF}\u{10FFFF}",
"z",
];
for prefix in [
"a",
"a\u{10FFFF}",
"a\u{10FFFF}\u{10FFFF}",
"a\u{D7FF}",
"\u{10FFFF}",
] {
let range = SkRange {
lower: Some((">=", s(prefix))),
upper: string_prefix_upper_bound(prefix).map(|u| ("<", SortKeyValue::S(u))),
};
for key in keys {
assert_eq!(
range.contains(&s(key)),
key.starts_with(prefix),
"prefix {} key {}",
prefix.escape_unicode(),
key.escape_unicode()
);
}
}
}

#[test]
fn begins_with_in_reverse_replaces_the_exclusive_upper_bound() {
let mut range = SkRange {
lower: Some((">=", s("ab"))),
upper: Some(("<", SortKeyValue::S(string_upper_bound("ab")))),
upper: Some((
"<",
SortKeyValue::S(string_prefix_upper_bound("ab").unwrap()),
)),
};
range.advance_past(s("abz"), false);
assert_eq!(
Expand Down Expand Up @@ -370,17 +416,23 @@ mod tests {
),
(" AND sk >= ? AND sk <= ?".to_owned(), strs(&["aa", "m"]))
);
// begins_with on a string: inclusive prefix, exclusive prefix plus U+10FFFF.
// begins_with on a string: inclusive prefix, exclusive next prefix.
assert_eq!(
range(
"pk = :p AND begins_with(sk, :v)",
&[(":v", s_av("ab"))],
ScalarAttributeType::S
),
(
" AND sk >= ? AND sk < ?".to_owned(),
strs(&["ab", "ab\u{10FFFF}"])
)
(" AND sk >= ? AND sk < ?".to_owned(), strs(&["ab", "ac"]))
);
// A prefix of U+10FFFF alone has no finite upper bound.
assert_eq!(
range(
"pk = :p AND begins_with(sk, :v)",
&[(":v", s_av("\u{10FFFF}"))],
ScalarAttributeType::S
),
(" AND sk >= ?".to_owned(), strs(&["\u{10FFFF}"]))
);
// begins_with on binary: the last byte below 0xFF is incremented and the
// rest dropped; an all-0xFF prefix has no finite upper bound.
Expand Down
109 changes: 95 additions & 14 deletions crates/storage-postgres/src/data/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ use extenddb_core::types::{
};
use extenddb_storage::error::StorageError;
use extenddb_storage::util::SortKeyValue;
use extenddb_storage::util::{parse_sk, pk_to_text, sk_info};
use extenddb_storage::util::{parse_sk, pk_to_text, sk_info, string_prefix_upper_bound};

/// Extra pagination bind values for index queries.
///
Expand Down Expand Up @@ -82,8 +82,9 @@ pub(crate) struct SkSqlInfo {
pub(crate) fn build_sk_sql(
sk_cond: &SortKeyCondition,
sk_col: &str,
maps: &ExpressionMaps,
param_idx: &mut u32,
) -> SkSqlInfo {
) -> Result<SkSqlInfo, StorageError> {
// Apply COLLATE "C" for string sort key columns to get byte-order comparison.
let collate = if sk_col == "sk_s" || sk_col.ends_with("_s") {
" COLLATE \"C\""
Expand All @@ -102,7 +103,7 @@ pub(crate) fn build_sk_sql(
};
let frag = format!(" AND {sk_col}{collate} {sql_op} ${param_idx}");
*param_idx += 1;
SkSqlInfo { fragment: frag }
Ok(SkSqlInfo { fragment: frag })
}
SortKeyCondition::Between { .. } => {
let frag = format!(
Expand All @@ -111,9 +112,9 @@ pub(crate) fn build_sk_sql(
hi = *param_idx + 1
);
*param_idx += 2;
SkSqlInfo { fragment: frag }
Ok(SkSqlInfo { fragment: frag })
}
SortKeyCondition::BeginsWith { .. } => {
SortKeyCondition::BeginsWith { prefix, .. } => {
let is_binary = sk_col == "sk_b" || sk_col.ends_with("_b");
if is_binary {
// For binary sort keys, bind both prefix and incremented-prefix
Expand All @@ -125,22 +126,50 @@ pub(crate) fn build_sk_sql(
hi = *param_idx + 1
);
*param_idx += 2;
SkSqlInfo { fragment: frag }
} else {
// For string sort keys, append the max Unicode code point.
// With COLLATE "C" (byte order), prefix || chr(1114111) is
// strictly greater than any string starting with prefix.
Ok(SkSqlInfo { fragment: frag })
} else if string_prefix_bounded(prefix, sk_col, maps)? {
// For string sort keys the exclusive upper bound is the next
// prefix in code point order, computed in Rust (see
// bind_sk_condition). With COLLATE "C" the comparison is in
// byte order, which is code point order.
let frag = format!(
" AND {sk_col}{collate} >= ${p} AND {sk_col}{collate} < (${p} || chr(1114111))",
p = *param_idx
" AND {sk_col}{collate} >= ${lo} AND {sk_col}{collate} < ${hi}",
lo = *param_idx,
hi = *param_idx + 1
);
*param_idx += 2;
Ok(SkSqlInfo { fragment: frag })
} else {
// A prefix of only U+10FFFF has no upper bound.
let frag = format!(" AND {sk_col}{collate} >= ${p}", p = *param_idx);
*param_idx += 1;
SkSqlInfo { fragment: frag }
Ok(SkSqlInfo { fragment: frag })
}
}
}
}

/// Whether a string `begins_with` prefix has a finite upper bound, so the SQL
/// carries a placeholder for it. [`bind_sk_condition`] binds that bound.
///
/// A numeric prefix on a numeric column is refused here: it passes `parse_sk`
/// in [`bind_sk_condition`] and would bind one value for two placeholders.
/// Other non-string prefixes are refused by `parse_sk` before anything is
/// bound.
fn string_prefix_bounded(
prefix: &Expr,
sk_col: &str,
maps: &ExpressionMaps,
) -> Result<bool, StorageError> {
match resolve_expr_to_av(prefix, maps)? {
AttributeValue::S(s) => Ok(string_prefix_upper_bound(&s).is_some()),
AttributeValue::N(_) if sk_col.ends_with("_n") => Err(StorageError::Internal(
"begins_with is not supported on numeric sort keys".to_owned(),
)),
_ => Ok(true),
}
}

/// Execute a query SQL statement with dynamic parameter binding.
///
/// # Parameter binding order:
Expand Down Expand Up @@ -249,7 +278,14 @@ fn bind_sk_condition<'q>(
let q = bind_sk_value(query, &sk);
Ok(q.bind(upper))
} else {
Ok(bind_sk_value(query, &sk))
let q = bind_sk_value(query, &sk);
match &sk {
SortKeyValue::S(p) => match string_prefix_upper_bound(p) {
Some(upper) => Ok(q.bind(upper)),
None => Ok(q),
},
_ => Ok(q),
}
}
}
SortKeyCondition::Between { low, high, .. } => {
Expand Down Expand Up @@ -359,3 +395,48 @@ pub(crate) async fn execute_scan_sql(
pub(crate) fn build_key(item: &Item, key_schema: &[KeySchemaElement]) -> Item {
extract_key(item, key_schema)
}

#[cfg(test)]
mod tests {
use std::collections::HashMap;

use extenddb_core::expression::{Expr, ExpressionMaps, PathElement, SortKeyCondition};
use extenddb_core::types::AttributeValue;

use super::build_sk_sql;

fn begins_with(value: AttributeValue) -> (SortKeyCondition, ExpressionMaps) {
let cond = SortKeyCondition::BeginsWith {
path: vec![PathElement::Attribute("sk".to_owned())],
prefix: Expr::Placeholder(":v".to_owned()),
};
let maps = ExpressionMaps::new(HashMap::new(), HashMap::from([(":v".to_owned(), value)]));
(cond, maps)
}

#[test]
fn a_string_prefix_gets_an_upper_bound_placeholder_unless_it_has_none() {
let (cond, maps) = begins_with(AttributeValue::S("ab".to_owned()));
let mut idx = 2;
let info = build_sk_sql(&cond, "sk_s", &maps, &mut idx).unwrap();
assert_eq!(
info.fragment,
" AND sk_s COLLATE \"C\" >= $2 AND sk_s COLLATE \"C\" < $3"
);
assert_eq!(idx, 4);

let (cond, maps) = begins_with(AttributeValue::S("\u{10FFFF}".to_owned()));
let mut idx = 2;
let info = build_sk_sql(&cond, "sk_s", &maps, &mut idx).unwrap();
assert_eq!(info.fragment, " AND sk_s COLLATE \"C\" >= $2");
assert_eq!(idx, 3);
}

#[test]
fn a_numeric_prefix_on_a_numeric_column_is_refused_before_any_sql() {
let (cond, maps) = begins_with(AttributeValue::N("1".to_owned()));
let mut idx = 2;
assert!(build_sk_sql(&cond, "sk_n", &maps, &mut idx).is_err());
assert_eq!(idx, 2);
}
}
2 changes: 1 addition & 1 deletion crates/storage-postgres/src/data/query_scan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -191,7 +191,7 @@ impl PostgresEngine {
&& let Some((pos, sk_type)) = range_column(sk_cond.path())
.or_else(|| sk_info_val.map(|(_, sk_type)| (0, sk_type)))
{
let info = build_sk_sql(sk_cond, &sk_column_n(pos, sk_type), &mut param_idx);
let info = build_sk_sql(sk_cond, &sk_column_n(pos, sk_type), maps, &mut param_idx)?;
sql.push_str(&info.fragment);
Some(sk_type)
} else {
Expand Down
24 changes: 14 additions & 10 deletions crates/storage-sqlite/src/data/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,16 +6,17 @@
//!
//! Sort-key comparisons bind through [`super::sk_bound`], so `N` keys are
//! compared as their order-preserving TEXT encoding (D2): `>`, `<`, `BETWEEN`
//! all remain numerically correct. `begins_with` applies only to `S` (string
//! prefix via the maximal code point `char(1114111)`) and `B` (byte-range via
//! an incremented upper bound); DynamoDB rejects `begins_with` on `N`.
//! all remain numerically correct. `begins_with` applies only to `S` and `B`,
//! each read as a range up to the next prefix (see
//! [`string_prefix_upper_bound`] and [`binary_prefix_upper_bound`]); DynamoDB
//! rejects `begins_with` on `N`.

use extenddb_core::expression::{self, CompareOp, Expr, ExpressionMaps, SortKeyCondition};
use extenddb_core::types::{
AttributeValue, Item, KeySchemaElement, ScalarAttributeType, extract_key,
};
use extenddb_storage::error::StorageError;
use extenddb_storage::util::{SortKeyValue, parse_sk};
use extenddb_storage::util::{SortKeyValue, parse_sk, string_prefix_upper_bound};

use super::BoundValue;

Expand Down Expand Up @@ -103,12 +104,15 @@ pub(super) fn build_sk_sql_and_binds(
// rows.
None => Ok((format!(" AND {sk_col} >= ?"), vec![SortKeyValue::B(b)])),
},
// For strings the SQL upper bound is `? || char(1114111)`,
// so the same prefix is bound twice.
SortKeyValue::S(s) => Ok((
format!(" AND {sk_col} >= ? AND {sk_col} < (? || char(1114111))"),
vec![SortKeyValue::S(s.clone()), SortKeyValue::S(s)],
)),
// The exclusive upper bound is the next prefix in code point
// order; a prefix of only U+10FFFF has none.
SortKeyValue::S(s) => match string_prefix_upper_bound(&s) {
Some(upper) => Ok((
format!(" AND {sk_col} >= ? AND {sk_col} < ?"),
vec![SortKeyValue::S(s), SortKeyValue::S(upper)],
)),
None => Ok((format!(" AND {sk_col} >= ?"), vec![SortKeyValue::S(s)])),
},
SortKeyValue::N(_) => Err(StorageError::Validation(
"begins_with is not supported on numeric sort keys".to_owned(),
)),
Expand Down
2 changes: 2 additions & 0 deletions crates/storage/src/util/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ mod arn;
mod backup_id;
mod catalog_version;
mod key;
mod prefix;
mod shard_stream;

pub use arn::{
Expand All @@ -29,6 +30,7 @@ pub use key::{
merge_attribute_definitions, parse_sk, pk_to_text, recover_sort_key_definitions,
refuse_conditions_beyond_first_range, sk_column, sk_column_n, sk_info,
};
pub use prefix::string_prefix_upper_bound;
pub use shard_stream::shard_in_stream;

/// Whether an UpdateTable's GSI actions must be refused because the GSIs named
Expand Down
Loading
Loading