use std::collections::HashMap;
use nodedb_physical::PhysicalTaskVisitor;
use nodedb_physical::physical_plan::document::DocumentOp;
use nodedb_physical::physical_plan::query::AggregateSpec;
use nodedb_physical::physical_plan::{PARAMETERIZED_AGGREGATE_FUNCTIONS, SortKeySpec};
use nodedb_query::expr::GroupKeySpec;
use nodedb_sql::types::SqlPlan;
use nodedb_sql::types::filter::Filter;
use nodedb_sql::types::query::{AggregateExpr, Projection, SortKey, WindowSpec};
use nodedb_sql::types_expr::{SqlExpr, SqlValue};
use nodedb_sql::{
AggregateVisitArgs, DocumentIndexLookupVisitArgs, JoinVisitArgs, SubqueryVisitArgs,
};
use nodedb_types::result::QueryResult;
use nodedb_types::value::Value;
use crate::error::LiteError;
use crate::query::engine::LiteQueryEngine;
use crate::query::expr_convert::convert_sql_expr;
use crate::query::filter_convert::{sql_filters_to_metadata, sql_value_to_value};
use crate::query::physical_visitor::LiteDataPlaneVisitor;
use crate::query::query_ops::aggregate::execute_aggregate;
use crate::query::query_ops::joins::inline_hash::execute_inline_hash_join;
use crate::query::query_ops::joins::spec::JoinSpec;
use crate::storage::engine::StorageEngine;
use super::adapter::LiteFut;
use super::having_eval::{apply_having_result, make_agg_alias_map};
use super::projection::join_projections;
use super::scan_post::{
ScanPostArgs, apply_scan_post_processing, apply_window_functions, distinct_rows, filter_rows,
project_rows, sort_rows,
};
pub(super) fn sql_agg_to_spec(agg: &AggregateExpr) -> Result<AggregateSpec, LiteError> {
let function = nodedb_sql::planner::agg_naming::aggregate_function_name(agg);
if PARAMETERIZED_AGGREGATE_FUNCTIONS.contains(&function.as_str()) {
return Err(LiteError::Unsupported {
detail: format!(
"aggregate function {}() is not supported in Lite",
agg.function
),
});
}
let (field, expr) = match agg.args.first() {
None | Some(SqlExpr::Wildcard) => ("*".to_string(), None),
Some(SqlExpr::Column { name, .. }) => (name.clone(), None),
Some(other) => ("*".to_string(), Some(convert_sql_expr(other)?)),
};
Ok(AggregateSpec {
function,
alias: agg.alias.clone(),
user_alias: None,
field,
expr,
param: None,
})
}
fn convert_aggregates(aggs: &[AggregateExpr]) -> Result<Vec<AggregateSpec>, LiteError> {
aggs.iter().map(sql_agg_to_spec).collect()
}
fn sort_key_to_spec(k: &SortKey) -> SortKeySpec {
match &k.expr {
SqlExpr::Column { name, .. } => SortKeySpec::column(name.clone(), k.ascending),
_ => SortKeySpec {
expr: nodedb_query::expr::SqlExpr::Literal(nodedb_types::value::Value::Null),
ascending: k.ascending,
nulls_first: !k.ascending,
},
}
}
fn encode_filters(filters: &[Filter]) -> Result<Vec<u8>, LiteError> {
if filters.is_empty() {
return Ok(Vec::new());
}
match sql_filters_to_metadata(filters, &[])?.meta {
None => Ok(Vec::new()),
Some(mf) => zerompk::to_msgpack_vec(&mf).map_err(|e| LiteError::Serialization {
detail: format!("encode filters: {e}"),
}),
}
}
fn result_to_maps(result: QueryResult) -> Vec<HashMap<String, Value>> {
let cols = result.columns;
result
.rows
.into_iter()
.map(|row| cols.iter().cloned().zip(row).collect())
.collect()
}
fn encode_result_msgpack(result: &QueryResult) -> Result<Vec<u8>, LiteError> {
let maps: Vec<HashMap<String, Value>> = result
.rows
.iter()
.map(|row| {
result
.columns
.iter()
.cloned()
.zip(row.iter().cloned())
.collect()
})
.collect();
zerompk::to_msgpack_vec(&maps).map_err(|e| LiteError::Serialization {
detail: format!("encode join side msgpack: {e}"),
})
}
pub(super) fn lower_aggregate<'a, S: StorageEngine + 'a>(
engine: &'a LiteQueryEngine<S>,
permit: Option<&'a crate::engine::fts::coordinator::TextMutationPermit>,
args: AggregateVisitArgs<'_>,
) -> Result<LiteFut<'a>, LiteError> {
let AggregateVisitArgs {
input,
group_by,
aggregates,
having,
limit,
grouping_sets,
sort_keys,
} = args;
let input = input.clone();
let group_cols: Vec<GroupKeySpec> = group_by
.iter()
.enumerate()
.map(|(i, e)| match e {
SqlExpr::Column { name, .. } => Ok(GroupKeySpec::column(name.clone())),
other => {
let expr = convert_sql_expr(other)?;
Ok(GroupKeySpec {
output_name: format!("group_key_{i}"),
field: None,
expr: Some(expr),
})
}
})
.collect::<Result<Vec<_>, LiteError>>()?;
let agg_specs = convert_aggregates(aggregates)?;
let agg_alias_map = make_agg_alias_map(aggregates);
let having_post = having.to_vec();
let sort_pairs: Vec<SortKeySpec> = sort_keys.iter().map(sort_key_to_spec).collect();
let gs: Vec<Vec<u32>> = grouping_sets
.unwrap_or(&[])
.iter()
.map(|s| s.iter().map(|&i| i as u32).collect())
.collect();
Ok(Box::pin(async move {
let source_result = engine.execute_plan_admitted(&input, permit).await?;
let rows = result_to_maps(source_result);
let result = execute_aggregate(rows, &group_cols, &agg_specs, &[], &[], &sort_pairs, &gs)?;
let mut result = apply_having_result(result, &having_post, &agg_alias_map);
result.rows.truncate(limit);
Ok(result)
}))
}
pub(super) fn lower_join<'a, S: StorageEngine + 'a>(
engine: &'a LiteQueryEngine<S>,
permit: Option<&'a crate::engine::fts::coordinator::TextMutationPermit>,
args: JoinVisitArgs<'_>,
) -> Result<LiteFut<'a>, LiteError> {
let JoinVisitArgs {
left,
right,
on,
join_type,
limit,
projection,
filters,
..
} = args;
let left = left.clone();
let right = right.clone();
let on = on.to_vec();
let limit = limit.unwrap_or(usize::MAX);
let join_type_str = format!("{join_type:?}").to_lowercase();
let proj = join_projections(projection, "JOIN")?;
let post_filters_bytes = encode_filters(filters)?;
Ok(Box::pin(async move {
let left_result = engine.execute_plan_admitted(&left, permit).await?;
let right_result = engine.execute_plan_admitted(&right, permit).await?;
let left_bytes = encode_result_msgpack(&left_result)?;
let right_bytes = encode_result_msgpack(&right_result)?;
execute_inline_hash_join(
&left_bytes,
&right_bytes,
None,
&JoinSpec {
on: &on,
join_type: &join_type_str,
limit,
projection: &proj,
post_filters: &post_filters_bytes,
},
)
}))
}
pub(super) fn lower_document_index_lookup<'a, S: StorageEngine + 'a>(
engine: &'a LiteQueryEngine<S>,
args: DocumentIndexLookupVisitArgs<'_>,
) -> Result<LiteFut<'a>, LiteError> {
let DocumentIndexLookupVisitArgs {
collection,
field,
value,
filters,
projection,
sort_keys,
limit,
offset,
distinct,
window_functions,
..
} = args;
let collection = collection.to_string();
let field = field.to_string();
let probe = sql_value_to_value(value)?;
let filters = filters.to_vec();
let sort_keys = sort_keys.to_vec();
let window_functions = window_functions.to_vec();
let projection = projection.to_vec();
Ok(Box::pin(async move {
let raw = crate::query::document_ops::index_reads::indexed_fetch(
engine,
&collection,
&field,
&probe,
usize::MAX,
0,
)
.await?;
apply_scan_post_processing(
raw,
ScanPostArgs {
filters: &filters,
sort_keys: &sort_keys,
window_specs: &window_functions,
projection: &projection,
sequences: engine.sequences(),
limit,
offset,
distinct,
},
)
}))
}
pub(super) fn lower_range_scan<'a, S: StorageEngine + 'a>(
engine: &'a LiteQueryEngine<S>,
collection: &str,
field: &str,
lower: Option<&SqlValue>,
upper: Option<&SqlValue>,
limit: usize,
) -> Result<LiteFut<'a>, LiteError> {
let col = nodedb_types::QualifiedCollection::new(nodedb_types::DatabaseId::DEFAULT, collection);
let fld = field.to_string();
let encode_bound = |v: &SqlValue| -> Result<Vec<u8>, LiteError> {
let ndb_val = sql_value_to_value(v)?;
zerompk::to_msgpack_vec(&ndb_val).map_err(|e| LiteError::Serialization {
detail: format!("encode range bound: {e}"),
})
};
let lo_bytes: Option<Vec<u8>> = lower.map(encode_bound).transpose()?;
let hi_bytes: Option<Vec<u8>> = upper.map(encode_bound).transpose()?;
let op = DocumentOp::RangeScan {
collection: col,
field: fld,
lower: lo_bytes,
upper: hi_bytes,
limit,
rls_filters: Vec::new(),
};
let mut phys = LiteDataPlaneVisitor::new(engine);
let fut = phys.document(&op)?;
Ok(Box::pin(fut))
}
pub(super) fn lower_cte<'a, S: StorageEngine + 'a>(
engine: &'a LiteQueryEngine<S>,
permit: Option<&'a crate::engine::fts::coordinator::TextMutationPermit>,
definitions: &[(String, SqlPlan)],
outer: &SqlPlan,
) -> Result<LiteFut<'a>, LiteError> {
let definitions = definitions.to_vec();
let outer = outer.clone();
Ok(Box::pin(async move {
for (_name, def_plan) in &definitions {
let _ = engine.execute_plan_admitted(def_plan, permit).await?;
}
engine.execute_plan_admitted(&outer, permit).await
}))
}
pub(super) fn lower_subquery<'a, S: StorageEngine + 'a>(
engine: &'a LiteQueryEngine<S>,
permit: Option<&'a crate::engine::fts::coordinator::TextMutationPermit>,
args: SubqueryVisitArgs<'_>,
) -> Result<LiteFut<'a>, LiteError> {
let SubqueryVisitArgs {
input,
filters,
projection,
window_functions,
sort_keys,
offset,
distinct,
limit,
} = args;
let input = input.clone();
let filters = filters.to_vec();
let projection = subquery_projection_columns(projection, window_functions)?;
let window_functions = window_functions.to_vec();
let sort_keys = sort_keys.to_vec();
Ok(Box::pin(async move {
let mut result = engine.execute_plan_admitted(&input, permit).await?;
filter_rows(&mut result, &filters)?;
apply_window_functions(&mut result, &window_functions)?;
sort_rows(&mut result, &sort_keys)?;
if distinct {
distinct_rows(&mut result, &projection);
}
project_rows(&mut result, &projection);
if offset > 0 {
result.rows = result.rows.into_iter().skip(offset).collect();
}
if let Some(n) = limit {
result.rows.truncate(n);
}
Ok(result)
}))
}
fn subquery_projection_columns(
projection: &[Projection],
window_functions: &[WindowSpec],
) -> Result<Vec<String>, LiteError> {
let mut names = Vec::with_capacity(projection.len());
for p in projection {
match p {
Projection::Computed { alias, .. }
if window_functions.iter().any(|spec| spec.alias == *alias) =>
{
names.push(alias.clone());
}
Projection::Column(qname) => {
names.push(qname.rsplit('.').next().unwrap_or(qname).to_string());
}
Projection::Star | Projection::QualifiedStar(_) => return Ok(Vec::new()),
Projection::Computed { .. } | Projection::CpComputed { .. } => {
return Err(LiteError::BadRequest {
detail: "a computed projection over an ORDER BY / OFFSET / DISTINCT subquery \
is not supported; select the base columns in the subquery and \
compute them in an outer SELECT"
.into(),
});
}
}
}
Ok(names)
}