use super::super::access_path::{AccessPath, FallbackReason};
use super::super::select_optimizer::SSTablePredicate;
use crate::types::{RowKey, ScanRow, Value};
#[derive(Debug)]
pub(super) enum PartitionLookupOutcome {
Targeted(Vec<u8>),
MultiTargeted(Vec<Vec<u8>>),
Fallback(FallbackReason),
}
pub(super) fn honest_targeted_path(targeted: AccessPath, engaged: bool) -> AccessPath {
if engaged {
targeted
} else {
AccessPath::FallbackFullScan {
reason: FallbackReason::TombstonesBuildNoPrune,
}
}
}
pub(super) const MAX_IN_TARGETED_LOOKUPS: usize = 64;
pub(super) fn classify_partition_lookup(
predicates: &[SSTablePredicate],
schema: Option<&crate::schema::TableSchema>,
) -> PartitionLookupOutcome {
use super::super::select_optimizer::SSTableFilterOp;
let Some(schema) = schema else {
return PartitionLookupOutcome::Fallback(FallbackReason::NoSchema);
};
if schema.partition_keys.is_empty() {
return PartitionLookupOutcome::Fallback(FallbackReason::NoSchema);
}
let mut per_column_values: Vec<Vec<Value>> = Vec::with_capacity(schema.partition_keys.len());
for pk in &schema.partition_keys {
let predicate = predicates.iter().find(|p| {
!p.is_token()
&& p.column == pk.name
&& matches!(p.operation, SSTableFilterOp::Equal | SSTableFilterOp::In)
});
let Some(predicate) = predicate else {
return PartitionLookupOutcome::Fallback(
FallbackReason::PartitionKeyNotFullyConstrained,
);
};
if predicate.values.is_empty() {
return PartitionLookupOutcome::Fallback(
FallbackReason::PartitionKeyNotFullyConstrained,
);
}
per_column_values.push(predicate.values.clone());
}
let product_size = per_column_values
.iter()
.try_fold(1usize, |acc, vals| acc.checked_mul(vals.len()));
match product_size {
Some(n) if n <= MAX_IN_TARGETED_LOOKUPS => {}
_ => {
return PartitionLookupOutcome::Fallback(
FallbackReason::PartitionKeyNotFullyConstrained,
);
}
}
let combinations = cartesian_product(&per_column_values);
let mut seen: std::collections::HashSet<Vec<u8>> = std::collections::HashSet::new();
let mut keys: Vec<Vec<u8>> = Vec::with_capacity(combinations.len());
for values in &combinations {
match crate::storage::partition_key_codec::encode_partition_key_columns(values, schema) {
Ok(bytes) => {
if seen.insert(bytes.clone()) {
keys.push(bytes);
}
}
Err(_) => {
return PartitionLookupOutcome::Fallback(
FallbackReason::PartitionKeyEncodingFailed,
);
}
}
}
match keys.len() {
0 => PartitionLookupOutcome::Fallback(FallbackReason::PartitionKeyNotFullyConstrained),
1 => PartitionLookupOutcome::Targeted(keys.into_iter().next().unwrap_or_default()),
_ => PartitionLookupOutcome::MultiTargeted(keys),
}
}
#[cfg(not(feature = "tombstones"))]
fn coerce_clustering_value(value: &Value, cql_type: &str) -> Option<Value> {
use crate::schema::CqlType;
let ty = CqlType::parse(cql_type).ok()?;
let as_i128 = |v: &Value| -> Option<i128> {
match v {
Value::TinyInt(i) => Some(*i as i128),
Value::SmallInt(i) => Some(*i as i128),
Value::Integer(i) => Some(*i as i128),
Value::BigInt(i) | Value::Counter(i) | Value::Timestamp(i) => Some(*i as i128),
_ => None,
}
};
match ty {
CqlType::TinyInt => Some(Value::TinyInt(i8::try_from(as_i128(value)?).ok()?)),
CqlType::SmallInt => Some(Value::SmallInt(i16::try_from(as_i128(value)?).ok()?)),
CqlType::Int => Some(Value::Integer(i32::try_from(as_i128(value)?).ok()?)),
CqlType::BigInt | CqlType::Counter => {
Some(Value::BigInt(i64::try_from(as_i128(value)?).ok()?))
}
CqlType::Timestamp => Some(Value::Timestamp(i64::try_from(as_i128(value)?).ok()?)),
CqlType::Text | CqlType::Varchar | CqlType::Ascii => match value {
Value::Text(_) => Some(value.clone()),
_ => None,
},
CqlType::Uuid | CqlType::TimeUuid => match value {
Value::Uuid(_) => Some(value.clone()),
_ => None,
},
CqlType::Boolean => match value {
Value::Boolean(_) => Some(value.clone()),
_ => None,
},
CqlType::Blob => match value {
Value::Blob(_) => Some(value.clone()),
_ => None,
},
_ => None,
}
}
#[cfg(not(feature = "tombstones"))]
pub(super) fn classify_clustering_slice(
predicates: &[SSTablePredicate],
schema: Option<&crate::schema::TableSchema>,
) -> Option<crate::storage::sstable::reader::ClusteringSlice> {
use super::super::select_optimizer::SSTableFilterOp;
use crate::storage::sstable::reader::ClusteringSlice;
let schema = schema?;
let first_ck = schema.clustering_keys.first()?;
let ck_type = first_ck.data_type.as_str();
let coerce = |v: &Value| coerce_clustering_value(v, ck_type);
let non_first_ck_restricted = schema.clustering_keys.iter().skip(1).any(|ck| {
predicates
.iter()
.any(|p| !p.is_token() && p.column == ck.name)
});
if non_first_ck_restricted {
return None;
}
let mut start: Vec<Value> = Vec::new();
let mut start_inclusive = false;
let mut end: Vec<Value> = Vec::new();
let mut end_inclusive = false;
let mut saw_supported = false;
for p in predicates
.iter()
.filter(|p| !p.is_token() && p.column == first_ck.name)
{
match &p.operation {
SSTableFilterOp::Equal => {
let v = coerce(p.values.first()?)?;
start = vec![v.clone()];
start_inclusive = true;
end = vec![v];
end_inclusive = true;
saw_supported = true;
}
SSTableFilterOp::Gt => {
start = vec![coerce(p.values.first()?)?];
start_inclusive = false;
saw_supported = true;
}
SSTableFilterOp::Gte => {
start = vec![coerce(p.values.first()?)?];
start_inclusive = true;
saw_supported = true;
}
SSTableFilterOp::Lt => {
end = vec![coerce(p.values.first()?)?];
end_inclusive = false;
saw_supported = true;
}
SSTableFilterOp::Lte => {
end = vec![coerce(p.values.first()?)?];
end_inclusive = true;
saw_supported = true;
}
SSTableFilterOp::Range => {
if p.values.len() < 2 {
return None;
}
start = vec![coerce(&p.values[0])?];
start_inclusive = true;
end = vec![coerce(&p.values[1])?];
end_inclusive = true;
saw_supported = true;
}
SSTableFilterOp::In | SSTableFilterOp::Prefix | SSTableFilterOp::BloomFilter => {
return None;
}
}
}
if !saw_supported {
return None;
}
Some(ClusteringSlice {
start,
start_inclusive,
end,
end_inclusive,
})
}
pub(super) fn sort_rows_by_token(rows: &mut [(RowKey, ScanRow)]) {
rows.sort_by(|a, b| {
crate::util::cassandra_murmur3::cmp_partition_keys_by_token(&a.0 .0, &b.0 .0)
});
}
#[cfg(not(feature = "tombstones"))]
pub(super) fn requests_clustering_reverse(
order_by: &crate::query::select_ast::OrderByClause,
schema: &crate::schema::TableSchema,
) -> bool {
use crate::query::select_ast::{SelectExpression, SortDirection};
use crate::schema::ClusteringOrder;
if order_by.items.len() != 1 {
return false;
}
let item = &order_by.items[0];
let SelectExpression::Column(col_ref) = &item.expression else {
return false;
};
let Some(first_ck) = schema.clustering_keys.first() else {
return false;
};
if col_ref.column != first_ck.name {
return false;
}
let stored_desc = matches!(first_ck.order, ClusteringOrder::Desc);
let requested_desc = matches!(item.direction, SortDirection::Descending);
requested_desc != stored_desc
}
#[cfg(not(feature = "tombstones"))]
impl super::SelectExecutor {
pub(super) async fn targeted_partition_rows(
&self,
table: &crate::types::TableId,
pk_bytes: &[u8],
predicates: &[SSTablePredicate],
order_by: Option<&crate::query::select_ast::OrderByClause>,
schema: Option<&crate::schema::TableSchema>,
context: &mut super::ExecutionContext,
) -> crate::Result<Vec<(RowKey, ScanRow)>> {
let clustering = classify_clustering_slice(predicates, schema);
if let (Some(order_by), Some(schema)) = (order_by, schema) {
if requests_clustering_reverse(order_by, schema) {
if let Some(rows) = self
.storage
.scan_partition_clustering_reverse(table, pk_bytes, Some(schema))
.await?
{
let path = AccessPath::PartitionLookup;
context.access_path = Some(path.clone());
crate::query::access_path::record(path);
context.reverse_served = true;
return Ok(rows);
}
}
}
let (rows, engaged) = self
.storage
.scan_partition_clustering(table, pk_bytes, clustering.as_ref(), schema)
.await?;
let path = if engaged {
AccessPath::ClusteringSlice
} else {
AccessPath::PartitionLookup
};
context.access_path = Some(path.clone());
crate::query::access_path::record(path);
Ok(rows)
}
}
fn cartesian_product(per_column: &[Vec<Value>]) -> Vec<Vec<Value>> {
let total: usize = per_column.iter().map(Vec::len).product();
if total == 0 {
return Vec::new();
}
let num_cols = per_column.len();
let mut out: Vec<Vec<Value>> = Vec::with_capacity(total);
let mut idx = vec![0usize; num_cols];
for _ in 0..total {
let mut combo = Vec::with_capacity(num_cols);
for (c, col) in per_column.iter().enumerate() {
combo.push(col[idx[c]].clone());
}
out.push(combo);
for c in (0..num_cols).rev() {
idx[c] += 1;
if idx[c] < per_column[c].len() {
break;
}
idx[c] = 0;
}
}
out
}
#[cfg(test)]
mod tests {
use super::super::test_support::{composite_pk_schema, single_pk_schema};
use super::*;
#[test]
fn honest_targeted_path_reports_fallback_when_not_engaged() {
for targeted in [
AccessPath::PartitionLookup,
AccessPath::MultiPartitionLookup,
AccessPath::MetadataPartitionLookup,
AccessPath::StreamingPartitionLookup,
] {
let engaged = honest_targeted_path(targeted.clone(), true);
assert_eq!(engaged, targeted, "engaged must keep the targeted label");
assert!(engaged.is_targeted());
let fallback = honest_targeted_path(targeted, false);
assert_eq!(
fallback,
AccessPath::FallbackFullScan {
reason: FallbackReason::TombstonesBuildNoPrune,
},
"a non-engaged targeted call must report the honest no-prune fallback"
);
assert!(fallback.is_full_scan());
assert!(!fallback.is_targeted());
}
}
#[test]
fn classify_partition_lookup_targets_uuid_literal() {
use super::super::super::select_optimizer::{SSTableFilterOp, SSTablePredicate};
let uuid = [
0x55u8, 0x0e, 0x84, 0x00, 0xe2, 0x9b, 0x41, 0xd4, 0xa7, 0x16, 0x44, 0x66, 0x55, 0x44,
0x00, 0x00,
];
let schema = single_pk_schema("id", "uuid");
let predicate =
SSTablePredicate::column("id", SSTableFilterOp::Equal, vec![Value::Uuid(uuid)]);
match classify_partition_lookup(std::slice::from_ref(&predicate), Some(&schema)) {
PartitionLookupOutcome::Targeted(pk_bytes) => assert_eq!(
pk_bytes,
uuid.to_vec(),
"fast path must encode the UUID literal to the raw 16-byte on-disk key"
),
PartitionLookupOutcome::MultiTargeted(keys) => panic!(
"Issue #956: a single UUID-literal `=` must be a single Targeted lookup, not \
MultiTargeted (got {} keys)",
keys.len()
),
PartitionLookupOutcome::Fallback(reason) => panic!(
"Issue #956: UUID-literal `=` predicate must engage the partition fast path, \
got fallback {reason:?}"
),
}
}
#[test]
fn classify_partition_lookup_falls_back_for_uuid_range_predicate() {
use super::super::super::select_optimizer::{SSTableFilterOp, SSTablePredicate};
let uuid = [1u8; 16];
let schema = single_pk_schema("id", "uuid");
let predicate =
SSTablePredicate::column("id", SSTableFilterOp::Gt, vec![Value::Uuid(uuid)]);
assert!(
matches!(
classify_partition_lookup(std::slice::from_ref(&predicate), Some(&schema)),
PartitionLookupOutcome::Fallback(FallbackReason::PartitionKeyNotFullyConstrained)
),
"a range restriction on the partition key must report the \
PartitionKeyNotFullyConstrained fallback, not a targeted lookup",
);
}
#[test]
fn classify_partition_lookup_in_yields_multi_targeted() {
use super::super::super::select_optimizer::{SSTableFilterOp, SSTablePredicate};
let schema = single_pk_schema("id", "int");
let predicate = SSTablePredicate::column(
"id",
SSTableFilterOp::In,
vec![Value::Integer(1), Value::Integer(2), Value::Integer(3)],
);
match classify_partition_lookup(std::slice::from_ref(&predicate), Some(&schema)) {
PartitionLookupOutcome::MultiTargeted(keys) => {
assert_eq!(keys.len(), 3, "one targeted key per IN element");
assert_eq!(keys[0], 1i32.to_be_bytes().to_vec());
assert_eq!(keys[1], 2i32.to_be_bytes().to_vec());
assert_eq!(keys[2], 3i32.to_be_bytes().to_vec());
}
other => panic!("IN over the complete pk must be MultiTargeted, got {other:?}"),
}
}
#[test]
fn classify_partition_lookup_in_dedupes_and_collapses_singletons() {
use super::super::super::select_optimizer::{SSTableFilterOp, SSTablePredicate};
let schema = single_pk_schema("id", "int");
let one = SSTablePredicate::column("id", SSTableFilterOp::In, vec![Value::Integer(7)]);
assert!(
matches!(
classify_partition_lookup(std::slice::from_ref(&one), Some(&schema)),
PartitionLookupOutcome::Targeted(_)
),
"a single-element IN must collapse to a single Targeted lookup",
);
let dup = SSTablePredicate::column(
"id",
SSTableFilterOp::In,
vec![Value::Integer(5), Value::Integer(5), Value::Integer(6)],
);
match classify_partition_lookup(std::slice::from_ref(&dup), Some(&schema)) {
PartitionLookupOutcome::MultiTargeted(keys) => {
assert_eq!(keys.len(), 2, "duplicate IN elements must be deduplicated");
assert_eq!(keys[0], 5i32.to_be_bytes().to_vec());
assert_eq!(keys[1], 6i32.to_be_bytes().to_vec());
}
other => panic!("IN (5,5,6) must dedupe to 2 MultiTargeted keys, got {other:?}"),
}
}
#[test]
fn classify_partition_lookup_large_in_falls_back() {
use super::super::super::select_optimizer::{SSTableFilterOp, SSTablePredicate};
let schema = single_pk_schema("id", "int");
let values: Vec<Value> = (0..(MAX_IN_TARGETED_LOOKUPS as i32 + 1))
.map(Value::Integer)
.collect();
let predicate = SSTablePredicate::column("id", SSTableFilterOp::In, values);
assert!(
matches!(
classify_partition_lookup(std::slice::from_ref(&predicate), Some(&schema)),
PartitionLookupOutcome::Fallback(FallbackReason::PartitionKeyNotFullyConstrained)
),
"an IN list over the cap must fall back to a full scan",
);
let at_cap: Vec<Value> = (0..(MAX_IN_TARGETED_LOOKUPS as i32))
.map(Value::Integer)
.collect();
let at_cap_pred = SSTablePredicate::column("id", SSTableFilterOp::In, at_cap);
assert!(
matches!(
classify_partition_lookup(std::slice::from_ref(&at_cap_pred), Some(&schema)),
PartitionLookupOutcome::MultiTargeted(_)
),
"an IN list exactly at the cap must still be MultiTargeted",
);
}
#[test]
fn classify_partition_lookup_composite_in_over_cap_falls_back_without_overalloc() {
use super::super::super::select_optimizer::{SSTableFilterOp, SSTablePredicate};
let schema = composite_pk_schema(("a", "int"), ("b", "int"));
let a_vals: Vec<Value> = (0..1000).map(Value::Integer).collect();
let b_vals: Vec<Value> = (0..1000).map(Value::Integer).collect();
let preds = vec![
SSTablePredicate::column("a", SSTableFilterOp::In, a_vals),
SSTablePredicate::column("b", SSTableFilterOp::In, b_vals),
];
assert!(
matches!(
classify_partition_lookup(&preds, Some(&schema)),
PartitionLookupOutcome::Fallback(FallbackReason::PartitionKeyNotFullyConstrained)
),
"a composite IN whose product exceeds the cap must fall back to a full scan",
);
}
#[test]
fn classify_partition_lookup_composite_in_within_cap_is_targeted() {
use super::super::super::select_optimizer::{SSTableFilterOp, SSTablePredicate};
let schema = composite_pk_schema(("a", "int"), ("b", "int"));
let a_vals: Vec<Value> = (0..4).map(Value::Integer).collect();
let b_vals: Vec<Value> = (0..4).map(Value::Integer).collect();
let preds = vec![
SSTablePredicate::column("a", SSTableFilterOp::In, a_vals),
SSTablePredicate::column("b", SSTableFilterOp::In, b_vals),
];
match classify_partition_lookup(&preds, Some(&schema)) {
PartitionLookupOutcome::MultiTargeted(keys) => assert_eq!(
keys.len(),
16,
"4 x 4 composite IN must yield 16 targeted keys (the full product)"
),
other => panic!("composite IN within the cap must be MultiTargeted, got {other:?}"),
}
}
#[test]
fn classify_partition_lookup_token_range_falls_back() {
use super::super::super::select_optimizer::{SSTableFilterOp, SSTablePredicate};
let schema = single_pk_schema("id", "int");
let predicate = SSTablePredicate::token(
vec!["id".to_string()],
SSTableFilterOp::Gte,
vec![Value::BigInt(-100)],
);
assert!(
matches!(
classify_partition_lookup(std::slice::from_ref(&predicate), Some(&schema)),
PartitionLookupOutcome::Fallback(FallbackReason::PartitionKeyNotFullyConstrained)
),
"a token-range restriction must fall back honestly (no fake pruning)",
);
}
#[test]
fn classify_partition_lookup_falls_back_without_schema() {
assert!(
matches!(
classify_partition_lookup(&[], None),
PartitionLookupOutcome::Fallback(FallbackReason::NoSchema)
),
"no schema must report the NoSchema fallback reason",
);
}
#[cfg(not(feature = "dhat-heap"))]
#[test]
fn cartesian_product_builds_each_combo_in_one_allocation() {
use crate::test_alloc_probe::measure;
fn reference(per_column: &[Vec<Value>]) -> Vec<Vec<Value>> {
let mut out: Vec<Vec<Value>> = vec![Vec::new()];
for column_values in per_column {
let mut next = Vec::with_capacity(out.len() * column_values.len());
for prefix in &out {
for value in column_values {
let mut combo = prefix.clone();
combo.push(value.clone());
next.push(combo);
}
}
out = next;
}
out
}
let per_column: Vec<Vec<Value>> = vec![
(0..4).map(Value::Integer).collect(),
(0..4).map(|i| Value::Integer(i + 100)).collect(),
(0..4).map(|i| Value::Integer(i + 1000)).collect(),
];
let (ref_allocs, ref_out) = measure(|| reference(&per_column));
let (new_allocs, new_out) = measure(|| cartesian_product(&per_column));
assert_eq!(
new_out, ref_out,
"cartesian product output must be byte-identical to the prefix-clone build"
);
assert_eq!(new_out.len(), 64, "4×4×4 composite IN expands to 64 keys");
assert!(
new_allocs < ref_allocs,
"index-based cartesian_product must allocate strictly fewer times than \
the prefix-clone reference: {new_allocs} vs {ref_allocs}"
);
}
}