use super::aggregate::{aggregate_col_idx, aggregate_output_name};
use super::pagination::{
compute_page_window, compute_remaining_limit_after_page, compute_split_recommended,
};
use super::{
CursorToken, QueryResult, encode_cursor, extract_pk_key, extract_sort_key, row_after_cursor,
};
use crate::catalog::Catalog;
use crate::catalog::namespace_key;
use crate::catalog::types::{Row, Value};
use crate::query::error::QueryError;
use crate::query::operators::{AggregateOperator, Operator, ScanOperator, compile_expr};
use crate::query::plan::{JoinType, Query, QueryOptions};
use crate::storage::encoded_key::EncodedKey;
use crate::storage::keyspace::KeyspaceSnapshot;
use std::collections::HashMap;
pub(super) struct JoinQueryExecutionRequest<'a> {
pub snapshot: &'a KeyspaceSnapshot,
pub catalog: &'a Catalog,
pub project_id: &'a str,
pub scope_id: &'a str,
pub query: Query,
pub options: QueryOptions,
pub snapshot_seq: u64,
pub max_scan_rows: usize,
pub cursor_state: Option<CursorToken>,
pub cursor_signing_key: Option<&'a [u8; 32]>,
}
pub(super) fn execute_join_query(
request: JoinQueryExecutionRequest<'_>,
) -> Result<QueryResult, QueryError> {
let JoinQueryExecutionRequest {
snapshot,
catalog,
project_id,
scope_id,
query,
options,
snapshot_seq,
max_scan_rows,
cursor_state,
cursor_signing_key,
} = request;
if query.offset.unwrap_or(0) > 0 && query.limit.is_none() {
return Err(QueryError::InvalidQuery {
reason: "OFFSET requires LIMIT".into(),
});
}
if cursor_state.is_some() && query.offset.unwrap_or(0) > 0 {
return Err(QueryError::InvalidQuery {
reason: "OFFSET cannot be combined with cursor pagination".into(),
});
}
let (base_ns_project, base_ns_scope, base_table) =
resolve_table_ref(project_id, scope_id, &query.table);
let base_schema = catalog
.tables
.get(&(
namespace_key(&base_ns_project, &base_ns_scope),
base_table.clone(),
))
.ok_or_else(|| QueryError::TableNotFound {
project_id: base_ns_project.clone(),
table: base_table.clone(),
})?;
let base_alias = query.table_alias.clone().unwrap_or(base_table.clone());
let mut columns: Vec<String> = base_schema
.columns
.iter()
.map(|c| format!("{base_alias}.{}", c.name))
.collect();
let base_count = snapshot
.table(&base_ns_project, &base_ns_scope, &base_table)
.map(|t| t.rows.len())
.unwrap_or(0);
if !options.allow_full_scan && base_count > max_scan_rows {
return Err(QueryError::ScanBoundExceeded {
estimated_rows: base_count as u64,
max_scan_rows: max_scan_rows as u64,
});
}
let mut rows: Vec<Row> = snapshot
.table(&base_ns_project, &base_ns_scope, &base_table)
.map(|t| t.rows.values().cloned().collect())
.unwrap_or_default();
for join in &query.joins {
let (jp, js, jt) = resolve_table_ref(project_id, scope_id, &join.table);
let join_schema = catalog
.tables
.get(&(namespace_key(&jp, &js), jt.clone()))
.ok_or_else(|| QueryError::TableNotFound {
project_id: jp.clone(),
table: jt.clone(),
})?;
let join_alias = join.alias.clone().unwrap_or(jt.clone());
let join_table = snapshot.table(&jp, &js, &jt);
let join_col_offset = columns.len();
let mut next_columns = columns.clone();
next_columns.extend(
join_schema
.columns
.iter()
.map(|c| format!("{join_alias}.{}", c.name)),
);
let (left_idx, right_idx) = match join.join_type {
JoinType::Cross => (None, None),
_ => {
let left = join
.left_column
.as_ref()
.ok_or_else(|| QueryError::InvalidQuery {
reason: "join requires left_column".into(),
})?;
let right = join
.right_column
.as_ref()
.ok_or_else(|| QueryError::InvalidQuery {
reason: "join requires right_column".into(),
})?;
let left_idx = columns.iter().position(|c| c == left).ok_or_else(|| {
QueryError::ColumnNotFound {
table: query.table.clone(),
column: left.clone(),
}
})?;
let right_idx = join_schema
.columns
.iter()
.position(|c| format!("{join_alias}.{}", c.name) == *right || c.name == *right)
.ok_or_else(|| QueryError::ColumnNotFound {
table: join.table.clone(),
column: right.clone(),
})?;
(Some(left_idx), Some(right_idx))
}
};
let can_probe_right_primary_key =
matches!(join.join_type, JoinType::Inner | JoinType::Left)
&& right_idx
.map(|idx| is_single_column_primary_key_join(join_schema, idx))
.unwrap_or(false);
let estimated_join_rows = if matches!(join.join_type, JoinType::Cross) {
rows.len()
.saturating_mul(join_table.map(|table| table.rows.len()).unwrap_or(0).max(1))
} else if can_probe_right_primary_key {
rows.len()
} else {
rows.len()
.saturating_mul(join_table.map(|table| table.rows.len()).unwrap_or(0).max(1))
};
if !options.allow_full_scan && estimated_join_rows > max_scan_rows {
return Err(QueryError::ScanBoundExceeded {
estimated_rows: estimated_join_rows as u64,
max_scan_rows: max_scan_rows as u64,
});
}
let mut joined = Vec::new();
match join.join_type {
JoinType::Cross => {
let join_rows: Vec<&Row> = join_table
.map(|table| table.rows.values().collect())
.unwrap_or_default();
for left in &rows {
for right in &join_rows {
push_joined_row(&mut joined, left, right);
}
}
}
JoinType::Inner | JoinType::Left => {
let right_idx = right_idx.ok_or_else(|| QueryError::InvalidQuery {
reason: "join requires right join key".into(),
})?;
let left_idx = left_idx.ok_or_else(|| QueryError::InvalidQuery {
reason: "join requires left join key".into(),
})?;
if can_probe_right_primary_key {
for left in &rows {
let matched = join_table.and_then(|table| {
let encoded = EncodedKey::from_values(std::slice::from_ref(
&left.values[left_idx],
));
table.rows.get(&encoded)
});
if let Some(right) = matched {
push_joined_row(&mut joined, left, right);
} else if matches!(join.join_type, JoinType::Left) {
push_left_with_nulls(&mut joined, left, join_schema.columns.len());
}
}
} else {
let join_rows: Vec<&Row> = join_table
.map(|table| table.rows.values().collect())
.unwrap_or_default();
let mut right_map: HashMap<Value, Vec<&Row>> = HashMap::new();
for right in &join_rows {
if matches!(right.values[right_idx], Value::Null) {
continue;
}
right_map
.entry(right.values[right_idx].clone())
.or_default()
.push(right);
}
for left in &rows {
if let Some(matches) = right_map.get(&left.values[left_idx]) {
for right in matches {
push_joined_row(&mut joined, left, right);
}
} else if matches!(join.join_type, JoinType::Left) {
push_left_with_nulls(&mut joined, left, join_schema.columns.len());
}
}
}
}
JoinType::Right => {
let left_idx = left_idx.ok_or_else(|| QueryError::InvalidQuery {
reason: "join requires left join key".into(),
})?;
let right_idx = right_idx.ok_or_else(|| QueryError::InvalidQuery {
reason: "join requires right join key".into(),
})?;
let join_rows: Vec<&Row> = join_table
.map(|table| table.rows.values().collect())
.unwrap_or_default();
let mut left_map: HashMap<Value, Vec<&Row>> = HashMap::new();
for left in &rows {
if matches!(left.values[left_idx], Value::Null) {
continue;
}
left_map
.entry(left.values[left_idx].clone())
.or_default()
.push(left);
}
for right in &join_rows {
if let Some(matches) = left_map.get(&right.values[right_idx]) {
for left in matches {
push_joined_row(&mut joined, left, right);
}
} else {
push_nulls_with_right(&mut joined, join_col_offset, right);
}
}
}
}
rows = joined;
if !options.allow_full_scan && rows.len() > max_scan_rows {
return Err(QueryError::ScanBoundExceeded {
estimated_rows: rows.len() as u64,
max_scan_rows: max_scan_rows as u64,
});
}
columns = next_columns;
}
if let Some(predicate) = &query.predicate {
let compiled = compile_expr(predicate, &columns, "join")?;
rows.retain(|r| crate::query::operators::eval_compiled_expr_public(&compiled, r));
}
if !query.aggregates.is_empty() {
let group_by_idx = query
.group_by
.iter()
.map(|name| {
columns
.iter()
.position(|c| c == name)
.ok_or_else(|| QueryError::ColumnNotFound {
table: "join".into(),
column: name.clone(),
})
})
.collect::<Result<Vec<_>, _>>()?;
let agg_col_idx = query
.aggregates
.iter()
.map(|agg| aggregate_col_idx(agg, &columns))
.collect::<Result<Vec<_>, _>>()?;
let mut aggregate = AggregateOperator::new(
Box::new(ScanOperator::new(rows)),
query.aggregates.clone(),
group_by_idx,
agg_col_idx,
);
let mut aggregated_rows = Vec::new();
while let Some(row) = aggregate.next() {
aggregated_rows.push(row);
}
rows = aggregated_rows;
columns = query.group_by.clone();
columns.extend(query.aggregates.iter().map(aggregate_output_name));
}
if let Some(having) = &query.having {
if query.aggregates.is_empty() {
return Err(QueryError::InvalidQuery {
reason: "having requires aggregate or group_by".into(),
});
}
let compiled = compile_expr(having, &columns, "join")?;
rows.retain(|r| crate::query::operators::eval_compiled_expr_public(&compiled, r));
}
order_join_rows(&mut rows, &columns, &query)?;
let rows_examined = rows.len();
let page_window = compute_page_window(&query, &cursor_state, max_scan_rows)?;
let sort_indices: Vec<(usize, crate::query::plan::Order)> = if !query.order_by.is_empty() {
query
.order_by
.iter()
.filter_map(|(name, ord)| columns.iter().position(|c| c == name).map(|i| (i, *ord)))
.collect()
} else {
Vec::new()
};
let pk_indices: Vec<usize> = (0..columns.len()).collect();
let mut sliced = Vec::new();
let mut skipped = 0usize;
for row in rows {
if let Some(cursor) = &cursor_state
&& !row_after_cursor(&row, cursor, &sort_indices, &pk_indices)
{
continue;
}
if skipped < page_window.row_offset_count {
skipped += 1;
continue;
}
sliced.push(row);
if sliced.len() > page_window.effective_page_size {
break;
}
}
let has_more = sliced.len() > page_window.effective_page_size;
if has_more {
sliced.truncate(page_window.effective_page_size);
}
let cursor_last_row = sliced.last().cloned();
sliced = project_join_rows(sliced, &columns, &query)?;
let returned_rows = sliced.len();
let remaining_limit_after_page = compute_remaining_limit_after_page(
query.limit,
cursor_state.as_ref().and_then(|c| c.remaining_limit),
returned_rows,
);
let cursor = if has_more && page_window.row_offset_count == 0 {
let last_row = cursor_last_row.ok_or_else(|| QueryError::InvalidQuery {
reason: "invalid cursor state".into(),
})?;
Some(encode_cursor(
&CursorToken {
snapshot_seq,
last_sort_key: extract_sort_key(&last_row, &sort_indices),
last_pk: extract_pk_key(&last_row, &pk_indices),
page_size: page_window.page_size,
remaining_limit: remaining_limit_after_page,
},
cursor_signing_key,
)?)
} else {
None
};
let split_budget = query.limit.unwrap_or(max_scan_rows);
let split_recommended = compute_split_recommended(rows_examined, split_budget);
Ok(QueryResult {
rows_examined,
rows: sliced,
truncated: has_more,
cursor,
snapshot_seq,
materialized_seq: None,
split_recommended,
})
}
pub(super) fn resolve_table_ref(
project_id: &str,
scope_id: &str,
table_ref: &str,
) -> (String, String, String) {
if let Some(name) = table_ref.strip_prefix("_global.") {
return ("_global".to_string(), "app".to_string(), name.to_string());
}
(
project_id.to_string(),
scope_id.to_string(),
table_ref.to_string(),
)
}
fn is_single_column_primary_key_join(
join_schema: &crate::catalog::schema::TableSchema,
right_idx: usize,
) -> bool {
if join_schema.primary_key.len() != 1 {
return false;
}
join_schema
.columns
.get(right_idx)
.map(|column| column.name == join_schema.primary_key[0])
.unwrap_or(false)
}
fn order_join_rows(rows: &mut [Row], columns: &[String], query: &Query) -> Result<(), QueryError> {
if query.order_by.is_empty() {
return Ok(());
}
let order_pairs: Vec<(usize, crate::query::plan::Order)> = query
.order_by
.iter()
.map(|(col, ord)| {
columns
.iter()
.position(|c| c == col)
.map(|idx| (idx, *ord))
.ok_or_else(|| QueryError::ColumnNotFound {
table: "join".into(),
column: col.clone(),
})
})
.collect::<Result<_, _>>()?;
rows.sort_by(|a, b| {
for (idx, ord) in &order_pairs {
let cmp = a.values[*idx].cmp(&b.values[*idx]);
let ord_cmp = match ord {
crate::query::plan::Order::Asc => cmp,
crate::query::plan::Order::Desc => cmp.reverse(),
};
if !ord_cmp.is_eq() {
return ord_cmp;
}
}
std::cmp::Ordering::Equal
});
Ok(())
}
fn project_join_rows(
rows: Vec<Row>,
columns: &[String],
query: &Query,
) -> Result<Vec<Row>, QueryError> {
if query.select.is_empty() || query.select[0] == "*" {
return Ok(rows);
}
let selected_indices: Vec<usize> = query
.select
.iter()
.map(|col| {
columns
.iter()
.position(|c| c == col)
.ok_or_else(|| QueryError::ColumnNotFound {
table: "join".into(),
column: col.clone(),
})
})
.collect::<Result<_, _>>()?;
Ok(rows
.into_iter()
.map(|row| Row {
values: selected_indices
.iter()
.map(|idx| row.values[*idx].clone())
.collect(),
})
.collect())
}
fn push_joined_row(out: &mut Vec<Row>, left: &Row, right: &Row) {
let mut values = Vec::with_capacity(left.values.len().saturating_add(right.values.len()));
values.extend_from_slice(&left.values);
values.extend_from_slice(&right.values);
out.push(Row { values });
}
fn push_left_with_nulls(out: &mut Vec<Row>, left: &Row, null_count: usize) {
let mut values = Vec::with_capacity(left.values.len().saturating_add(null_count));
values.extend_from_slice(&left.values);
values.extend(std::iter::repeat_n(Value::Null, null_count));
out.push(Row { values });
}
fn push_nulls_with_right(out: &mut Vec<Row>, null_count: usize, right: &Row) {
let mut values = Vec::with_capacity(null_count.saturating_add(right.values.len()));
values.extend(std::iter::repeat_n(Value::Null, null_count));
values.extend_from_slice(&right.values);
out.push(Row { values });
}