use nodedb_physical::physical_plan::SortKeySpec;
use nodedb_sql::types::filter::Filter;
use nodedb_sql::types::query::SortKey;
use nodedb_sql::types_expr::SqlExpr;
use nodedb_sql::{LateralLoopVisitArgs, LateralTopKVisitArgs};
use crate::error::LiteError;
use crate::query::engine::LiteQueryEngine;
use crate::query::filter_convert::sql_filters_to_metadata;
use crate::storage::engine::StorageEngine;
use super::adapter::LiteFut;
use super::projection::join_projections;
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 lateral filters: {e}"),
}),
}
}
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,
},
}
}
pub(super) fn lower_lateral_top_k<'a, S: StorageEngine + 'a>(
engine: &'a LiteQueryEngine<S>,
permit: Option<&'a crate::engine::fts::coordinator::TextMutationPermit>,
args: LateralTopKVisitArgs<'_>,
) -> Result<LiteFut<'a>, LiteError> {
let LateralTopKVisitArgs {
outer,
outer_alias,
inner_collection,
inner_filters,
inner_order_by,
inner_limit,
correlation_keys,
lateral_alias,
projection,
left_join,
} = args;
let outer_sql = outer.clone();
let outer_alias_str = outer_alias.unwrap_or("outer").to_string();
let inner_col = inner_collection.to_string();
let inner_filt_bytes = encode_filters(inner_filters)?;
let inner_ob: Vec<SortKeySpec> = inner_order_by.iter().map(sort_key_to_spec).collect();
let inner_lim = inner_limit;
let corr_keys = correlation_keys.to_vec();
let lat_alias = lateral_alias.to_string();
let proj = join_projections(projection, "LATERAL")?;
let lj = left_join;
Ok(Box::pin(async move {
use crate::query::query_ops::lateral_top_k::{LateralTopKSpec, execute_lateral_top_k_sql};
execute_lateral_top_k_sql(
engine,
permit,
&outer_sql,
LateralTopKSpec {
outer_alias: &outer_alias_str,
inner_collection: &inner_col,
inner_filters: &inner_filt_bytes,
inner_order_by: &inner_ob,
inner_limit: inner_lim,
correlation_keys: &corr_keys,
lateral_alias: &lat_alias,
projection: &proj,
left_join: lj,
},
)
.await
}))
}
pub(super) fn lower_lateral_loop<'a, S: StorageEngine + 'a>(
engine: &'a LiteQueryEngine<S>,
permit: Option<&'a crate::engine::fts::coordinator::TextMutationPermit>,
args: LateralLoopVisitArgs<'_>,
) -> Result<LiteFut<'a>, LiteError> {
let LateralLoopVisitArgs {
outer,
outer_alias,
inner,
correlation_predicates,
lateral_alias,
projection,
outer_row_cap,
left_join,
} = args;
let outer_sql = outer.clone();
let outer_alias_str = outer_alias.unwrap_or("outer").to_string();
let inner_sql = inner.clone();
let corr_preds = correlation_predicates.to_vec();
let lat_alias = lateral_alias.to_string();
let proj = join_projections(projection, "LATERAL")?;
let cap = outer_row_cap;
let lj = left_join;
Ok(Box::pin(async move {
use crate::query::query_ops::lateral_loop::{LateralLoopSqlSpec, execute_lateral_loop_sql};
execute_lateral_loop_sql(
engine,
permit,
&outer_sql,
LateralLoopSqlSpec {
outer_alias: &outer_alias_str,
inner_sql: &inner_sql,
correlation_predicates: &corr_preds,
lateral_alias: &lat_alias,
projection: &proj,
left_join: lj,
outer_row_cap: cap,
},
)
.await
}))
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use nodedb_sql::types::SqlPlan;
use nodedb_sql::types::query::EngineType;
use nodedb_sql::{LateralLoopVisitArgs, LateralTopKVisitArgs};
use crate::PagedbStorageMem;
use crate::query::engine::{LiteQueryEngine, LiteQueryEngineParams};
async fn make_engine() -> LiteQueryEngine<PagedbStorageMem> {
use std::sync::Mutex;
let storage = Arc::new(
PagedbStorageMem::open_in_memory()
.await
.expect("in-memory pagedb"),
);
let crdt = Arc::new(Mutex::new(
crate::engine::crdt::CrdtEngine::new(1).expect("crdt"),
));
let governor = crate::query::engine::test_governor();
let strict = Arc::new(crate::engine::strict::StrictEngine::new(Arc::clone(
&storage,
)));
let columnar = Arc::new(crate::engine::columnar::ColumnarEngine::new(
Arc::clone(&storage),
crate::query::engine::test_scoped_memory(&governor, nodedb_mem::EngineId::Columnar),
));
let htap = Arc::new(crate::engine::htap::HtapBridge::new());
let timeseries = Arc::new(Mutex::new(
crate::engine::timeseries::engine::TimeseriesEngine::new(),
));
let vector_state = Arc::new(crate::engine::vector::VectorState::new(
Arc::clone(&storage),
100,
crate::query::engine::test_scoped_memory(&governor, nodedb_mem::EngineId::Vector),
));
let array_state = Arc::new(tokio::sync::Mutex::new(
crate::engine::array::engine::ArrayEngineState::open(&storage)
.await
.expect("array"),
));
let fts_state = Arc::new(crate::engine::fts::FtsState::new(Arc::clone(&governor)));
let spatial = Arc::new(Mutex::new(
crate::engine::spatial::SpatialIndexManager::new(
crate::query::engine::test_scoped_memory(&governor, nodedb_mem::EngineId::Spatial),
),
));
LiteQueryEngine::new(LiteQueryEngineParams {
crdt,
strict,
columnar,
htap,
storage,
timeseries,
vector_state,
array_state,
fts_state,
sparse_state: Arc::new(crate::engine::sparse_vector::SparseVectorState::new()),
spatial,
csr: Arc::new(Mutex::new(std::collections::HashMap::new())),
governor,
kv_local: crate::query::engine::test_kv_local(),
})
}
fn scan_plan(collection: &str) -> SqlPlan {
SqlPlan::Scan {
collection: collection.to_string(),
alias: None,
engine: EngineType::DocumentSchemaless,
filters: vec![],
projection: vec![],
sort_keys: vec![],
limit: None,
offset: 0,
distinct: false,
window_functions: vec![],
temporal: nodedb_sql::temporal::TemporalScope::default(),
}
}
#[tokio::test]
async fn test_lateral_top_k_lower() {
let engine = make_engine().await;
let outer = scan_plan("users");
let result = super::lower_lateral_top_k(
&engine,
None,
LateralTopKVisitArgs {
outer: &outer,
outer_alias: Some("u"),
inner_collection: "orders",
inner_filters: &[],
inner_order_by: &[],
inner_limit: 3,
correlation_keys: &[("id".to_string(), "user_id".to_string())],
lateral_alias: "o",
projection: &[],
left_join: false,
},
);
assert!(result.is_ok());
}
#[tokio::test]
async fn test_lateral_loop_lower() {
let engine = make_engine().await;
let outer = scan_plan("departments");
let inner = scan_plan("employees");
let result = super::lower_lateral_loop(
&engine,
None,
LateralLoopVisitArgs {
outer: &outer,
outer_alias: Some("d"),
inner: &inner,
correlation_predicates: &[("id".to_string(), "dept_id".to_string())],
lateral_alias: "e",
projection: &[],
outer_row_cap: 1000,
left_join: false,
},
);
assert!(result.is_ok());
}
}