use alloc::borrow::Cow;
use alloc::boxed::Box;
use alloc::vec::Vec;
use core::ops::Bound;
use spg_sql::ast::{BinOp, Expr, Literal, SelectStatement};
use spg_storage::{Catalog, ColumnSchema, IndexKey, Row, Table, Value};
use crate::eval::{self, EvalContext};
use crate::{
CancelToken, Engine, EngineError, QueryResult, apply_offset_and_limit, build_projection,
memoize,
};
pub(crate) fn try_nsw_knn(
stmt: &SelectStatement,
table: &Table,
schema_cols: &[ColumnSchema],
table_alias: &str,
snapshot: &spg_storage::snapshot::Snapshot,
) -> Option<Vec<usize>> {
if stmt.distinct {
return None;
}
let limit = usize::try_from(stmt.limit_literal()?).ok()?;
if limit == 0 {
return None;
}
if stmt.order_by.len() != 1 {
return None;
}
let order = &stmt.order_by[0];
if order.desc {
return None;
}
let Expr::Binary { lhs, op, rhs } = &order.expr else {
return None;
};
let metric = match op {
BinOp::L2Distance => spg_storage::NswMetric::L2,
BinOp::InnerProduct => spg_storage::NswMetric::InnerProduct,
BinOp::CosineDistance => spg_storage::NswMetric::Cosine,
_ => return None,
};
let ((Expr::Column(col), literal) | (literal, Expr::Column(col))) =
(lhs.as_ref(), rhs.as_ref())
else {
return None;
};
if let Some(q) = &col.qualifier
&& q != table_alias
{
return None;
}
let col_pos = schema_cols.iter().position(|s| s.name == col.name)?;
let query = literal_to_vector(literal)?;
let idx = spg_storage::nsw_index_on(table, col_pos)?;
if let Some(where_expr) = &stmt.where_ {
let over_fetch = limit.saturating_mul(10).max(NSW_OVER_FETCH_FLOOR);
let candidates = spg_storage::nsw_query(table, &idx.name, &query, over_fetch, metric);
let ctx = EvalContext::new(schema_cols, Some(table_alias));
let mut kept: Vec<usize> = Vec::with_capacity(limit);
for i in candidates {
if !table.is_row_visible(i, snapshot) {
continue;
}
let row = &table.rows()[i];
let cond = eval::eval_expr(where_expr, row, &ctx).ok()?;
if crate::eval::predicate_is_true(&cond, "WHERE", ctx.mysql_dialect).ok()? {
kept.push(i);
if kept.len() >= limit {
break;
}
}
}
Some(kept)
} else {
Some(
spg_storage::nsw_query(table, &idx.name, &query, limit, metric)
.into_iter()
.filter(|&i| table.is_row_visible(i, snapshot))
.collect(),
)
}
}
const NSW_OVER_FETCH_FLOOR: usize = 32;
pub(crate) fn try_pk_walk_top_n<'a>(
stmt: &SelectStatement,
catalog: &'a spg_storage::Catalog,
table: &'a Table,
schema_cols: &[ColumnSchema],
table_alias: &str,
engine: &Engine,
cancel: CancelToken<'_>,
) -> Option<Vec<Cow<'a, Row<'static>>>> {
if stmt.distinct || stmt.limit_with_ties {
return None;
}
if stmt.group_by.is_some() || stmt.having.is_some() {
return None;
}
let limit = usize::try_from(stmt.limit_literal()?).ok()?;
if limit == 0 {
return None;
}
let offset = stmt
.offset_literal()
.and_then(|n| usize::try_from(n).ok())
.unwrap_or(0);
const WALKER_OFFSET_CAP: usize = 65_536;
if offset > WALKER_OFFSET_CAP {
return None;
}
let want = offset.checked_add(limit)?;
if stmt.order_by.len() != 1 {
return None;
}
let order = &stmt.order_by[0];
let Expr::Column(col) = &order.expr else {
return None;
};
if let Some(q) = &col.qualifier
&& q != table_alias
{
return None;
}
let col_pos = schema_cols
.iter()
.position(|s| s.name.eq_ignore_ascii_case(&col.name))?;
let index = table.index_on(col_pos)?;
if !matches!(index.kind, spg_storage::IndexKind::BTree(_)) {
return None;
}
let where_expr = stmt.where_.as_ref();
let ctx = EvalContext::new(schema_cols, Some(table_alias));
let table_name = table.schema().name.as_str();
let mut kept: Vec<Cow<'a, Row>> = Vec::with_capacity(want);
let mut memo = memoize::MemoizeCache::new();
let compiled_where: Option<eval::CompiledExpr> = where_expr
.filter(|w| eval::fully_compilable(w))
.map(|w| eval::compile_expr(w, &ctx));
let mut eval_stack: Vec<spg_storage::Value<'static>> = Vec::new();
let walker: Box<dyn Iterator<Item = (&spg_storage::IndexKey, &Vec<spg_storage::RowLocator>)>> =
if order.desc {
Box::new(index.iter_desc())
} else {
Box::new(index.iter_asc())
};
let scan_snapshot = engine.current_snapshot();
for (key, locators) in walker {
for loc in locators {
let row_cow: Cow<'a, Row> = match *loc {
spg_storage::RowLocator::Hot(row_idx) => {
if !table.is_row_visible(row_idx, &scan_snapshot) {
continue;
}
match table.rows().get(row_idx) {
Some(r) => Cow::Borrowed(r),
None => continue,
}
}
spg_storage::RowLocator::Cold { segment_id, .. } => {
match catalog.resolve_cold_locator(table_name, segment_id, key) {
Some(r) => Cow::Owned(r),
None => continue,
}
}
};
if let Some(cw) = &compiled_where {
let cond = eval::eval_compiled(cw, row_cow.as_ref(), &ctx, &mut eval_stack).ok()?;
if !crate::eval::predicate_is_true(&cond, "WHERE", ctx.mysql_dialect).ok()? {
continue;
}
} else if let Some(w) = where_expr {
let cond = engine
.eval_expr_with_correlated(w, row_cow.as_ref(), &ctx, cancel, Some(&mut memo))
.ok()?;
if !crate::eval::predicate_is_true(&cond, "WHERE", ctx.mysql_dialect).ok()? {
continue;
}
}
kept.push(row_cow);
if kept.len() >= want {
return Some(kept);
}
}
}
Some(kept)
}
pub(crate) fn literal_to_vector(e: &Expr) -> Option<Vec<f32>> {
match e {
Expr::Literal(Literal::Vector(v)) => Some(v.clone()),
Expr::Cast { expr, .. } => literal_to_vector(expr),
_ => None,
}
}
pub(crate) fn materialise_in_order(
stmt: &SelectStatement,
schema_cols: &[ColumnSchema],
table_alias: &str,
ordered_rows: &[Cow<'_, Row<'static>>],
mysql: bool,
) -> Result<QueryResult, EngineError> {
let ctx = EvalContext::new(schema_cols, Some(table_alias));
let projection = build_projection(&stmt.items, schema_cols, table_alias, mysql)?;
let mut output_rows: Vec<Row<'static>> = Vec::with_capacity(ordered_rows.len());
for row_cow in ordered_rows {
let row = row_cow.as_ref();
let mut values = Vec::with_capacity(projection.len());
for p in &projection {
values.push(eval::eval_expr(&p.expr, row, &ctx)?);
}
output_rows.push(Row::new(values));
}
apply_offset_and_limit(
&mut output_rows,
stmt.offset_literal(),
stmt.limit_literal(),
);
let columns: Vec<ColumnSchema> = projection
.into_iter()
.map(|p| ColumnSchema::new(p.output_name, p.ty, p.nullable))
.collect();
Ok(QueryResult::Rows {
columns,
rows: output_rows,
})
}
pub(crate) fn try_index_seek_positions(
where_expr: &Expr,
schema_cols: &[ColumnSchema],
table: &Table,
table_alias: &str,
snapshot: &spg_storage::snapshot::Snapshot,
) -> Option<Vec<usize>> {
let seek_cap = table.rows().len() / 4;
if let Some((col_pos, lo, hi)) = parse_range_bounds(where_expr, schema_cols, table_alias)
&& let Some(idx) = table.index_on(col_pos)
&& let Some(locators) =
idx.lookup_range_capped_by(bound_as_ref(&lo), bound_as_ref(&hi), seek_cap, |l| {
match l {
spg_storage::RowLocator::Hot(i) => table.is_row_visible(i, snapshot),
spg_storage::RowLocator::Cold { .. } => true,
}
})
{
let mut out = Vec::with_capacity(locators.len());
let mut all_hot = true;
for loc in &locators {
match *loc {
spg_storage::RowLocator::Hot(i) => out.push(i),
spg_storage::RowLocator::Cold { .. } => {
all_hot = false;
break;
}
}
}
if all_hot {
table.note_index_scan(out.len() as u64);
return Some(out);
}
}
if let Expr::Binary {
lhs,
op: BinOp::And,
rhs,
} = where_expr
{
if let Some(p) = try_index_seek_positions(lhs, schema_cols, table, table_alias, snapshot) {
return Some(p);
}
return try_index_seek_positions(rhs, schema_cols, table, table_alias, snapshot);
}
let Expr::Binary {
lhs,
op: BinOp::Eq,
rhs,
} = where_expr
else {
return None;
};
let (col_pos, value) = resolve_col_literal_pair(lhs, rhs, schema_cols, table_alias)
.or_else(|| resolve_col_literal_pair(rhs, lhs, schema_cols, table_alias))?;
let idx = table.index_on(col_pos)?;
let key = IndexKey::from_value(&value)?;
let locators = idx.lookup_eq(&key);
let mut out = Vec::with_capacity(locators.len());
for loc in locators {
match *loc {
spg_storage::RowLocator::Hot(i) => {
if table.is_row_visible(i, snapshot) {
out.push(i);
}
}
spg_storage::RowLocator::Cold { .. } => return None,
}
}
table.note_index_scan(out.len() as u64);
Some(out)
}
fn flip_comparison(op: BinOp) -> Option<BinOp> {
match op {
BinOp::Gt => Some(BinOp::Lt),
BinOp::GtEq => Some(BinOp::LtEq),
BinOp::Lt => Some(BinOp::Gt),
BinOp::LtEq => Some(BinOp::GtEq),
_ => None,
}
}
fn parse_one_sided_range(
e: &Expr,
schema_cols: &[ColumnSchema],
table_alias: &str,
) -> Option<(usize, Bound<IndexKey>, Bound<IndexKey>)> {
let Expr::Binary { lhs, op, rhs } = e else {
return None;
};
let (col_pos, value, op) =
if let Some((p, v)) = resolve_col_literal_pair(lhs, rhs, schema_cols, table_alias) {
(p, v, *op)
} else if let Some((p, v)) = resolve_col_literal_pair(rhs, lhs, schema_cols, table_alias) {
(p, v, flip_comparison(*op)?)
} else {
return None;
};
let key = IndexKey::from_value(&value)?;
let bounds = match op {
BinOp::Gt => (Bound::Excluded(key), Bound::Unbounded),
BinOp::GtEq => (Bound::Included(key), Bound::Unbounded),
BinOp::Lt => (Bound::Unbounded, Bound::Excluded(key)),
BinOp::LtEq => (Bound::Unbounded, Bound::Included(key)),
_ => return None,
};
Some((col_pos, bounds.0, bounds.1))
}
fn parse_range_bounds(
where_expr: &Expr,
schema_cols: &[ColumnSchema],
table_alias: &str,
) -> Option<(usize, Bound<IndexKey>, Bound<IndexKey>)> {
let bounds = parse_range_bounds_inner(where_expr, schema_cols, table_alias)?;
if schema_cols
.get(bounds.0)
.is_some_and(|c| c.user_enum_type.is_some())
{
return None;
}
Some(bounds)
}
fn parse_range_bounds_inner(
where_expr: &Expr,
schema_cols: &[ColumnSchema],
table_alias: &str,
) -> Option<(usize, Bound<IndexKey>, Bound<IndexKey>)> {
let Expr::Binary {
lhs,
op: BinOp::And,
rhs,
} = where_expr
else {
return None;
};
let (c1, lo1, hi1) = parse_one_sided_range(lhs, schema_cols, table_alias)?;
let (c2, lo2, hi2) = parse_one_sided_range(rhs, schema_cols, table_alias)?;
if c1 != c2 {
return None;
}
let lo = match (lo1, lo2) {
(Bound::Unbounded, b) | (b, Bound::Unbounded) => b,
_ => return None, };
let hi = match (hi1, hi2) {
(Bound::Unbounded, b) | (b, Bound::Unbounded) => b,
_ => return None,
};
if matches!(lo, Bound::Unbounded) || matches!(hi, Bound::Unbounded) {
return None;
}
Some((c1, lo, hi))
}
fn bound_as_ref(b: &Bound<IndexKey>) -> Bound<&IndexKey> {
match b {
Bound::Included(k) => Bound::Included(k),
Bound::Excluded(k) => Bound::Excluded(k),
Bound::Unbounded => Bound::Unbounded,
}
}
fn value_from_key(
key: &spg_storage::IndexKey,
declared: spg_storage::DataType,
) -> Option<spg_storage::Value<'static>> {
use spg_storage::{IndexKey as K, Value};
Some(match (key, declared) {
(K::Int(n), spg_storage::DataType::SmallInt) => Value::SmallInt(i16::try_from(*n).ok()?),
(K::Int(n), spg_storage::DataType::Int) => Value::Int(i32::try_from(*n).ok()?),
(K::Int(n), spg_storage::DataType::BigInt) => Value::BigInt(*n),
(K::Text(t), spg_storage::DataType::Text) => Value::text(t.clone()),
(K::Bool(b), spg_storage::DataType::Bool) => Value::Bool(*b),
(K::Uuid(u), spg_storage::DataType::Uuid) => Value::Uuid(*u),
_ => return None,
})
}
pub(crate) fn try_index_only_range(
where_expr: &Expr,
schema_cols: &[ColumnSchema],
table: &Table,
table_alias: &str,
snapshot: &spg_storage::snapshot::Snapshot,
projected: usize,
) -> Option<Vec<spg_storage::Value<'static>>> {
let mut out: Vec<spg_storage::Value<'static>> = Vec::new();
match index_only_range_each(
where_expr,
schema_cols,
table,
table_alias,
snapshot,
projected,
&mut |v| {
out.push(v);
Ok(())
},
) {
Some(Ok(_)) => Some(out),
Some(Err(_)) | None => None,
}
}
pub(crate) fn index_only_range_each(
where_expr: &Expr,
schema_cols: &[ColumnSchema],
table: &Table,
table_alias: &str,
snapshot: &spg_storage::snapshot::Snapshot,
projected: usize,
sink: &mut dyn FnMut(spg_storage::Value<'static>) -> Result<(), EngineError>,
) -> Option<Result<usize, EngineError>> {
let (ty, lo, hi, idx) =
index_only_precheck(where_expr, schema_cols, table, table_alias, projected)?;
let entries = idx.range_keyed(bound_as_ref(&lo), bound_as_ref(&hi))?;
let mut headers = table.header_runs();
let mut n = 0usize;
for (key, loc) in entries {
let spg_storage::RowLocator::Hot(i) = loc else {
return Some(Err(EngineError::Unsupported(
"index-only scan met a locator outside the hot tier".into(),
)));
};
if !headers.visible(i, snapshot) {
continue;
}
let Some(v) = value_from_key(key, ty) else {
return Some(Err(EngineError::Unsupported(
"index-only scan: index key does not restore the column type".into(),
)));
};
if let Err(e) = sink(v) {
return Some(Err(e));
}
n += 1;
}
Some(Ok(n))
}
pub(crate) fn index_only_precheck<'t>(
where_expr: &Expr,
schema_cols: &[ColumnSchema],
table: &'t Table,
table_alias: &str,
projected: usize,
) -> Option<(
spg_storage::DataType,
Bound<IndexKey>,
Bound<IndexKey>,
&'t spg_storage::Index,
)> {
let (col_pos, lo, hi) = parse_index_only_bounds(where_expr, schema_cols, table_alias)?;
if col_pos != projected {
return None;
}
if table.has_cold_rows_fast() {
return None;
}
let ty = schema_cols[col_pos].ty;
if !key_restores_type(ty) {
return None;
}
let idx = table.index_on(col_pos)?;
Some((ty, lo, hi, idx))
}
fn parse_index_only_bounds(
where_expr: &Expr,
schema_cols: &[ColumnSchema],
table_alias: &str,
) -> Option<(usize, Bound<IndexKey>, Bound<IndexKey>)> {
if let Some(r) = parse_range_bounds(where_expr, schema_cols, table_alias) {
return Some(r);
}
let Expr::Binary {
lhs,
op: BinOp::Eq,
rhs,
} = where_expr
else {
return None;
};
let (col_pos, value) = resolve_col_literal_pair(lhs, rhs, schema_cols, table_alias)
.or_else(|| resolve_col_literal_pair(rhs, lhs, schema_cols, table_alias))?;
if value.is_null() {
return None;
}
let key = IndexKey::from_value(&value)?;
Some((col_pos, Bound::Included(key.clone()), Bound::Included(key)))
}
fn key_restores_type(ty: spg_storage::DataType) -> bool {
use spg_storage::DataType as T;
matches!(
ty,
T::SmallInt | T::Int | T::BigInt | T::Text | T::Bool | T::Uuid
)
}
fn try_range_seek<'a>(
where_expr: &Expr,
schema_cols: &[ColumnSchema],
table: &'a Table,
table_alias: &str,
snapshot: &spg_storage::snapshot::Snapshot,
) -> Option<Vec<Cow<'a, Row<'static>>>> {
let (col_pos, lo, hi) = parse_range_bounds(where_expr, schema_cols, table_alias)?;
let idx = table.index_on(col_pos)?;
let cap = table.rows().len() / 4;
let locators =
idx.lookup_range_capped_by(bound_as_ref(&lo), bound_as_ref(&hi), cap, |l| match l {
spg_storage::RowLocator::Hot(i) => table.is_row_visible(i, snapshot),
spg_storage::RowLocator::Cold { .. } => true,
})?;
let mut out: Vec<Cow<'a, Row>> = Vec::with_capacity(locators.len());
for loc in &locators {
match *loc {
spg_storage::RowLocator::Hot(i) => {
if let Some(row) = table.rows().get(i) {
out.push(Cow::Borrowed(row));
}
}
spg_storage::RowLocator::Cold { .. } => return None,
}
}
table.note_index_scan(out.len() as u64);
Some(out)
}
pub(crate) fn try_range_count(
where_expr: &Expr,
schema_cols: &[ColumnSchema],
table: &Table,
table_alias: &str,
snapshot: &spg_storage::snapshot::Snapshot,
) -> Option<i64> {
let (col_pos, lo, hi) = parse_range_bounds(where_expr, schema_cols, table_alias)?;
let idx = table.index_on(col_pos)?;
let locators = idx.lookup_range_capped(bound_as_ref(&lo), bound_as_ref(&hi), usize::MAX)?;
let mut count: i64 = 0;
for loc in &locators {
match *loc {
spg_storage::RowLocator::Hot(i) => {
if table.is_row_visible(i, snapshot) {
count += 1;
}
}
spg_storage::RowLocator::Cold { .. } => return None,
}
}
Some(count)
}
pub(crate) fn try_index_seek<'a>(
where_expr: &Expr,
schema_cols: &[ColumnSchema],
catalog: &'a Catalog,
table: &'a Table,
table_alias: &str,
snapshot: &spg_storage::snapshot::Snapshot,
) -> Option<Vec<Cow<'a, Row<'static>>>> {
if let Some(rows) = try_range_seek(where_expr, schema_cols, table, table_alias, snapshot) {
return Some(rows);
}
if let Expr::Binary {
lhs,
op: BinOp::And,
rhs,
} = where_expr
{
if let Some(rows) = try_index_seek(lhs, schema_cols, catalog, table, table_alias, snapshot)
{
return Some(rows);
}
return try_index_seek(rhs, schema_cols, catalog, table, table_alias, snapshot);
}
if let Some(rows) = try_inlist_seek(
where_expr,
schema_cols,
catalog,
table,
table_alias,
snapshot,
) {
return Some(rows);
}
let Expr::Binary {
lhs,
op: BinOp::Eq,
rhs,
} = where_expr
else {
return None;
};
let (col_pos, value) = resolve_col_literal_pair(lhs, rhs, schema_cols, table_alias)
.or_else(|| resolve_col_literal_pair(rhs, lhs, schema_cols, table_alias))?;
let idx = table.index_on(col_pos)?;
let key = IndexKey::from_value(&value)?;
let locators = idx.lookup_eq(&key);
let table_name = table.schema().name.as_str();
let mut out: Vec<Cow<'a, Row>> = Vec::with_capacity(locators.len());
for loc in locators {
match *loc {
spg_storage::RowLocator::Hot(i) => {
if !table.is_row_visible(i, snapshot) {
continue;
}
if let Some(row) = table.rows().get(i) {
out.push(Cow::Borrowed(row));
}
}
spg_storage::RowLocator::Cold { segment_id, .. } => {
if let Some(row) = catalog.resolve_cold_locator(table_name, segment_id, &key) {
out.push(Cow::Owned(row));
}
}
}
}
table.note_index_scan(out.len() as u64);
Some(out)
}
fn try_inlist_seek<'a>(
where_expr: &Expr,
schema_cols: &[ColumnSchema],
catalog: &'a Catalog,
table: &'a Table,
table_alias: &str,
snapshot: &spg_storage::snapshot::Snapshot,
) -> Option<Vec<Cow<'a, Row<'static>>>> {
let Expr::InList {
expr,
list,
negated: false,
} = where_expr
else {
return None;
};
let Expr::Column(c) = expr.as_ref() else {
return None;
};
if !c
.qualifier
.as_deref()
.is_none_or(|q| q.eq_ignore_ascii_case(table_alias))
{
return None;
}
let col_pos = schema_cols.iter().position(|s| s.name == c.name)?;
let idx = table.index_on(col_pos)?;
let mut keys: Vec<IndexKey> = Vec::with_capacity(list.len());
for e in list {
let Expr::Literal(l) = e else {
return None;
};
keys.push(IndexKey::from_value(&eval::literal_to_value(l))?);
}
let table_name = table.schema().name.as_str();
let mut out: Vec<Cow<'a, Row>> = Vec::new();
for key in &keys {
for loc in idx.lookup_eq(key) {
match *loc {
spg_storage::RowLocator::Hot(i) => {
if !table.is_row_visible(i, snapshot) {
continue;
}
if let Some(row) = table.rows().get(i) {
out.push(Cow::Borrowed(row));
}
}
spg_storage::RowLocator::Cold { segment_id, .. } => {
if let Some(row) = catalog.resolve_cold_locator(table_name, segment_id, key) {
out.push(Cow::Owned(row));
}
}
}
}
}
table.note_index_scan(out.len() as u64);
Some(out)
}
pub(crate) fn try_gin_seek<'a>(
where_expr: &Expr,
schema_cols: &[ColumnSchema],
catalog: &'a Catalog,
table: &'a Table,
table_alias: &str,
ctx: &eval::EvalContext<'_>,
snapshot: &spg_storage::snapshot::Snapshot,
) -> Option<Vec<Cow<'a, Row<'static>>>> {
if let Expr::Binary {
lhs,
op: BinOp::And,
rhs,
} = where_expr
{
if let Some(rows) =
try_gin_seek(lhs, schema_cols, catalog, table, table_alias, ctx, snapshot)
{
return Some(rows);
}
return try_gin_seek(rhs, schema_cols, catalog, table, table_alias, ctx, snapshot);
}
if let Expr::Binary {
lhs,
op: BinOp::Or,
rhs,
} = where_expr
{
let left = try_gin_seek(lhs, schema_cols, catalog, table, table_alias, ctx, snapshot)?;
let right = try_gin_seek(rhs, schema_cols, catalog, table, table_alias, ctx, snapshot)?;
let mut out: Vec<Cow<'a, Row>> = Vec::with_capacity(left.len() + right.len());
out.extend(left);
out.extend(right);
return Some(out);
}
let Expr::Binary {
lhs,
op: BinOp::TsMatch,
rhs,
} = where_expr
else {
return None;
};
let (col_pos, query) = resolve_gin_col_query(lhs, rhs, schema_cols, table_alias, ctx)
.or_else(|| resolve_gin_col_query(rhs, lhs, schema_cols, table_alias, ctx))?;
let idx = table
.indices()
.iter()
.find(|i| i.column_position == col_pos && (i.is_gin() || i.is_gin_fulltext()))?;
let candidates = gin_query_candidates(idx, &query)?;
let _ = catalog; let mut out: Vec<Cow<'a, Row>> = Vec::with_capacity(candidates.len());
for loc in candidates {
match loc {
spg_storage::RowLocator::Hot(i) => {
if !table.is_row_visible(i, snapshot) {
continue;
}
if let Some(row) = table.rows().get(i) {
out.push(Cow::Borrowed(row));
}
}
spg_storage::RowLocator::Cold { .. } => {}
}
}
Some(out)
}
pub(crate) fn try_gin_jsonb_seek<'a>(
where_expr: &Expr,
schema_cols: &[ColumnSchema],
table: &'a Table,
table_alias: &str,
snapshot: &spg_storage::snapshot::Snapshot,
) -> Option<Vec<Cow<'a, Row<'static>>>> {
if let Expr::Binary {
lhs,
op: BinOp::And,
rhs,
} = where_expr
{
if let Some(rows) = try_gin_jsonb_seek(lhs, schema_cols, table, table_alias, snapshot) {
return Some(rows);
}
return try_gin_jsonb_seek(rhs, schema_cols, table, table_alias, snapshot);
}
let Expr::Binary {
lhs,
op: BinOp::JsonContains,
rhs,
} = where_expr
else {
return None;
};
let col_pos = resolve_jsonb_column(lhs, schema_cols, table_alias)?;
let literal = resolve_jsonb_literal(rhs)?;
let idx = table
.indices()
.iter()
.find(|i| i.column_position == col_pos && i.is_gin_jsonb())?;
let tokens = spg_storage::jsonb_gin::extract_tokens(&literal);
if tokens.is_empty() {
return None;
}
let mut candidates: Vec<spg_storage::RowLocator> = idx.gin_jsonb_lookup(&tokens[0]).to_vec();
candidates.sort_by_key(locator_sort_key);
candidates.dedup_by_key(|l| locator_sort_key(l));
for tok in &tokens[1..] {
let mut next: Vec<spg_storage::RowLocator> = idx.gin_jsonb_lookup(tok).to_vec();
next.sort_by_key(locator_sort_key);
next.dedup_by_key(|l| locator_sort_key(l));
let mut out: Vec<spg_storage::RowLocator> = Vec::new();
let (mut i, mut j) = (0usize, 0usize);
while i < candidates.len() && j < next.len() {
let lk = locator_sort_key(&candidates[i]);
let rk = locator_sort_key(&next[j]);
match lk.cmp(&rk) {
core::cmp::Ordering::Less => i += 1,
core::cmp::Ordering::Greater => j += 1,
core::cmp::Ordering::Equal => {
out.push(candidates[i]);
i += 1;
j += 1;
}
}
}
candidates = out;
if candidates.is_empty() {
break;
}
}
let mut out: Vec<Cow<'a, Row>> = Vec::with_capacity(candidates.len());
for loc in candidates {
if let spg_storage::RowLocator::Hot(i) = loc
&& table.is_row_visible(i, snapshot)
&& let Some(row) = table.rows().get(i)
{
out.push(Cow::Borrowed(row));
}
}
Some(out)
}
fn resolve_jsonb_column(
e: &Expr,
schema_cols: &[ColumnSchema],
table_alias: &str,
) -> Option<usize> {
if let Expr::Column(c) = e {
if let Some(q) = &c.qualifier
&& !q.eq_ignore_ascii_case(table_alias)
{
return None;
}
let pos = schema_cols
.iter()
.position(|s| s.name.eq_ignore_ascii_case(&c.name))?;
if matches!(
schema_cols[pos].ty,
spg_storage::DataType::Json | spg_storage::DataType::Jsonb
) {
return Some(pos);
}
}
None
}
fn resolve_jsonb_literal(e: &Expr) -> Option<alloc::string::String> {
use spg_sql::ast::Literal;
match e {
Expr::Cast { expr, .. } => match expr.as_ref() {
Expr::Literal(Literal::String(s)) => Some(s.clone()),
_ => None,
},
Expr::Literal(Literal::String(s)) => Some(s.clone()),
_ => None,
}
}
pub(crate) fn try_trgm_seek<'a>(
where_expr: &Expr,
schema_cols: &[ColumnSchema],
table: &'a Table,
table_alias: &str,
snapshot: &spg_storage::snapshot::Snapshot,
) -> Option<Vec<Cow<'a, Row<'static>>>> {
if let Expr::Binary {
lhs,
op: BinOp::And,
rhs,
} = where_expr
{
if let Some(rows) = try_trgm_seek(lhs, schema_cols, table, table_alias, snapshot) {
return Some(rows);
}
return try_trgm_seek(rhs, schema_cols, table, table_alias, snapshot);
}
let Expr::Like { expr, pattern, .. } = where_expr else {
return None;
};
let Expr::Column(c) = expr.as_ref() else {
return None;
};
if let Some(q) = &c.qualifier
&& q != table_alias
{
return None;
}
let col_pos = schema_cols
.iter()
.position(|s| s.name.eq_ignore_ascii_case(&c.name))?;
let idx = table
.indices()
.iter()
.find(|i| i.column_position == col_pos && i.is_gin_trgm())?;
let Expr::Literal(spg_sql::ast::Literal::String(pat)) = pattern.as_ref() else {
return None;
};
let trigrams = spg_storage::trgm::trigrams_from_like_pattern(pat)?;
let mut iter = trigrams.iter();
let first = iter.next()?;
let mut acc: Vec<spg_storage::RowLocator> = {
let mut v = idx
.gin_trgm_lookup(spg_storage::trgm::trigram_str(first))
.to_vec();
v.sort_by_key(locator_sort_key);
v.dedup_by_key(|l| locator_sort_key(l));
v
};
for tri in iter {
let mut next: Vec<spg_storage::RowLocator> = idx
.gin_trgm_lookup(spg_storage::trgm::trigram_str(tri))
.to_vec();
next.sort_by_key(locator_sort_key);
next.dedup_by_key(|l| locator_sort_key(l));
let mut merged: Vec<spg_storage::RowLocator> =
Vec::with_capacity(acc.len().min(next.len()));
let (mut i, mut j) = (0usize, 0usize);
while i < acc.len() && j < next.len() {
let lk = locator_sort_key(&acc[i]);
let rk = locator_sort_key(&next[j]);
match lk.cmp(&rk) {
core::cmp::Ordering::Less => i += 1,
core::cmp::Ordering::Greater => j += 1,
core::cmp::Ordering::Equal => {
merged.push(acc[i]);
i += 1;
j += 1;
}
}
}
acc = merged;
if acc.is_empty() {
break;
}
}
let mut out: Vec<Cow<'a, Row>> = Vec::with_capacity(acc.len());
for loc in acc {
if let spg_storage::RowLocator::Hot(i) = loc
&& table.is_row_visible(i, snapshot)
&& let Some(row) = table.rows().get(i)
{
out.push(Cow::Borrowed(row));
}
}
Some(out)
}
pub(crate) fn resolve_gin_col_query(
col_side: &Expr,
query_side: &Expr,
schema_cols: &[ColumnSchema],
table_alias: &str,
ctx: &eval::EvalContext<'_>,
) -> Option<(usize, spg_storage::TsQueryAst)> {
let column = match col_side {
Expr::Column(c) => c,
Expr::FunctionCall { name, args }
if name.eq_ignore_ascii_case("to_tsvector") && !args.is_empty() =>
{
if let Expr::Column(c) = args.last().unwrap() {
c
} else {
return None;
}
}
_ => return None,
};
let c = column;
if let Some(q) = &c.qualifier
&& q != table_alias
{
return None;
}
let pos = schema_cols.iter().position(|s| s.name == c.name)?;
let empty_row = Row::new(Vec::new());
let v = eval::eval_expr(query_side, &empty_row, ctx).ok()?;
let Value::TsQuery(q) = v else { return None };
Some((pos, q))
}
pub(crate) fn gin_query_candidates(
idx: &spg_storage::Index,
query: &spg_storage::TsQueryAst,
) -> Option<Vec<spg_storage::RowLocator>> {
use spg_storage::TsQueryAst;
match query {
TsQueryAst::Term { word, .. } => {
let mut v: Vec<spg_storage::RowLocator> = idx.gin_lookup_word(word).to_vec();
v.sort_by_key(locator_sort_key);
v.dedup_by_key(|l| locator_sort_key(l));
Some(v)
}
TsQueryAst::And(l, r) => {
let mut left = gin_query_candidates(idx, l)?;
let mut right = gin_query_candidates(idx, r)?;
left.sort_by_key(locator_sort_key);
right.sort_by_key(locator_sort_key);
let mut out: Vec<spg_storage::RowLocator> = Vec::new();
let (mut i, mut j) = (0usize, 0usize);
while i < left.len() && j < right.len() {
let lk = locator_sort_key(&left[i]);
let rk = locator_sort_key(&right[j]);
match lk.cmp(&rk) {
core::cmp::Ordering::Less => i += 1,
core::cmp::Ordering::Greater => j += 1,
core::cmp::Ordering::Equal => {
out.push(left[i]);
i += 1;
j += 1;
}
}
}
Some(out)
}
TsQueryAst::Or(l, r) => {
let mut out = gin_query_candidates(idx, l)?;
out.extend(gin_query_candidates(idx, r)?);
out.sort_by_key(locator_sort_key);
out.dedup_by_key(|l| locator_sort_key(l));
Some(out)
}
TsQueryAst::Not(_) | TsQueryAst::Phrase { .. } => None,
}
}
pub(crate) fn locator_sort_key(l: &spg_storage::RowLocator) -> (u8, u64, u64) {
match *l {
spg_storage::RowLocator::Hot(i) => (0, i as u64, 0),
spg_storage::RowLocator::Cold {
segment_id,
page_offset,
} => (1, u64::from(segment_id), u64::from(page_offset)),
}
}
pub(crate) fn try_pk_predicate(
where_expr: &Expr,
schema_cols: &[ColumnSchema],
table_alias: &str,
) -> Option<(usize, IndexKey)> {
let Expr::Binary {
lhs,
op: BinOp::Eq,
rhs,
} = where_expr
else {
return None;
};
let (col_pos, value) = resolve_col_literal_pair(lhs, rhs, schema_cols, table_alias)
.or_else(|| resolve_col_literal_pair(rhs, lhs, schema_cols, table_alias))?;
let key = IndexKey::from_value(&value)?;
Some((col_pos, key))
}
pub(crate) fn resolve_col_literal_pair(
col_side: &Expr,
lit_side: &Expr,
schema_cols: &[ColumnSchema],
table_alias: &str,
) -> Option<(usize, Value<'static>)> {
let Expr::Column(c) = col_side else {
return None;
};
if let Some(q) = &c.qualifier
&& q != table_alias
{
return None;
}
let pos = schema_cols.iter().position(|s| s.name == c.name)?;
let Expr::Literal(l) = lit_side else {
return None;
};
let v = match l {
Literal::Integer(n) => {
if let Ok(small) = i32::try_from(*n) {
Value::Int(small)
} else {
Value::BigInt(*n)
}
}
Literal::Float(x) => Value::Float(*x),
Literal::Numeric { unscaled, scale } => Value::Numeric {
scaled: *unscaled,
scale: *scale,
kind: spg_storage::NumericKind::Finite,
},
Literal::NumericBig(s) => crate::conversions::big_literal_to_value(s),
Literal::String(s) => Value::text(s.clone()),
Literal::Bool(b) => Value::Bool(*b),
Literal::Null => Value::Null,
Literal::Vector(_)
| Literal::Interval { .. }
| Literal::TextArray(_)
| Literal::IntArray(_)
| Literal::BigIntArray(_) => return None,
};
let ty = schema_cols[pos].ty;
if matches!(l, Literal::String(_))
&& !matches!(
ty,
spg_storage::DataType::Text
| spg_storage::DataType::Varchar(_)
| spg_storage::DataType::Char(_)
)
{
return Some((
pos,
crate::conversions::coerce_value(v, ty, &c.name, pos).ok()?,
));
}
Some((pos, v))
}