mod join;
mod transform;
mod vtable;
use std::sync::Arc;
use reifydb_core::interface::catalog::id::IndexId;
use reifydb_evaluate::stack::Variable;
use reifydb_rql::{
nodes::{
AggregateNode as RqlAggregateNode, AssertNode as RqlAssertNode, GeneratorNode as RqlGeneratorNode,
InlineDataNode as RqlInlineDataNode, PartitionParameterScanNode as RqlPartitionParameterScanNode,
RowListLookupNode as RqlRowListLookupNode, RowParameterLookupNode as RqlRowParameterLookupNode,
RowPointLookupNode as RqlRowPointLookupNode, RowRangeScanNode as RqlRowRangeScanNode,
SeriesParameterScanNode as RqlSeriesParameterScanNode, SeriesScanNode as RqlSeriesScanNode,
SortNode as RqlSortNode, TableScanNode as RqlTableScanNode, TakeLimit, TakeNode as RqlTakeNode,
ViewScanNode as RqlViewScanNode,
},
plan::logical::row_predicate::RowPredicate,
query::QueryPlan as RqlQueryPlan,
};
use reifydb_transaction::transaction::{ScanLayout, Transaction};
use reifydb_value::{
fragment::Fragment,
value::{column_view::ColumnView, system_columns::user_columns},
};
use tracing::instrument;
use super::{
apply_transform::{ApplyTransformNode, UnknownTransformNode},
run_tests::RunTestsQueryNode,
};
#[cfg(feature = "column")]
use crate::vm::volcano::scan::{
column_series::ColumnSeriesScanNode, column_table::ColumnTableScanNode,
column_unsupported::UnsupportedColumnScanNode,
};
use crate::vm::volcano::{
aggregate::AggregateNode,
append::UnsupportedAppendNode,
assert::{AssertNode, AssertWithoutInputNode},
distinct::DistinctNode,
environment::EnvironmentNode,
filter::FilterNode,
generator::GeneratorNode,
inline::InlineDataNode,
lookup::UnsupportedLookupNode,
partition_parameter::{choose_partition, choose_partitions},
query::{QueryContext, QueryNode},
row_lookup::{RowListLookupNode, RowPointLookupNode, RowRangeScanNode, choose_row_lookup},
scan::{
dictionary::DictionaryScanNode, index::IndexScanNode, queue::QueueScan, remote::RemoteFetchNode,
ringbuffer::RingBufferScan, series::SeriesScanNode as VolcanoSeriesScanNode, table::TableScanNode,
view::ViewScanNode,
},
series_parameter::choose_series_range,
sort::SortNode,
take::TakeNode,
top_k::TopKNode,
variable::VariableNode,
window::UnsupportedWindowNode,
};
fn extract_source_name_from_query(plan: &RqlQueryPlan) -> Option<Fragment> {
match plan {
RqlQueryPlan::TableScan(node) => Some(Fragment::internal(node.source.def().name.clone())),
RqlQueryPlan::ViewScan(node) => Some(Fragment::internal(node.source.def().name())),
RqlQueryPlan::RingBufferScan(node) => Some(Fragment::internal(node.source.def().name.clone())),
RqlQueryPlan::DictionaryScan(node) => Some(Fragment::internal(node.source.def().name.clone())),
RqlQueryPlan::SeriesScan(node) => Some(Fragment::internal(node.source.def().name.clone())),
RqlQueryPlan::QueueScan(node) => Some(Fragment::internal(node.source.def().name.clone())),
RqlQueryPlan::RemoteScan(_) => None,
RqlQueryPlan::Assert(node) => node.input.as_ref().and_then(|p| extract_source_name_from_query(p)),
RqlQueryPlan::Filter(node) => extract_source_name_from_query(&node.input),
RqlQueryPlan::Map(node) => node.input.as_ref().and_then(|p| extract_source_name_from_query(p)),
RqlQueryPlan::Take(node) => extract_source_name_from_query(&node.input),
_ => None,
}
}
#[cfg(feature = "column")]
fn row_store_scan(plan: &RqlQueryPlan) -> Option<String> {
match plan {
RqlQueryPlan::ViewScan(node) => Some(format!("view '{}'", node.source.def().name())),
RqlQueryPlan::RingBufferScan(node) => Some(format!("ring buffer '{}'", node.source.def().name)),
RqlQueryPlan::QueueScan(node) => Some(format!("queue '{}'", node.source.def().name)),
RqlQueryPlan::DictionaryScan(node) => Some(format!("dictionary '{}'", node.source.def().name)),
RqlQueryPlan::IndexScan(node) => Some(format!("index scan of table '{}'", node.source.def().name)),
RqlQueryPlan::RowPointLookup(node) => Some(format!("row lookup on '{}'", node.source.name())),
RqlQueryPlan::RowListLookup(node) => Some(format!("row lookup on '{}'", node.source.name())),
RqlQueryPlan::RowRangeScan(node) => Some(format!("row range scan on '{}'", node.source.name())),
_ => None,
}
}
#[instrument(name = "volcano::compile", level = "debug", skip_all)]
pub(crate) fn compile<'a>(
plan: RqlQueryPlan,
rx: &mut Transaction<'a>,
context: Arc<QueryContext>,
) -> Box<dyn QueryNode> {
#[cfg(feature = "column")]
if rx.layout() == ScanLayout::Column
&& let Some(what) = row_store_scan(&plan)
{
return Box::new(UnsupportedColumnScanNode::new(what));
}
match plan {
RqlQueryPlan::Aggregate(RqlAggregateNode {
by,
map,
input,
..
}) => {
let input_node = compile(*input, rx, context.clone());
Box::new(AggregateNode::new(input_node, by, map, context))
}
RqlQueryPlan::Distinct(node) => {
let input = compile(*node.input, rx, context);
Box::new(DistinctNode::new(input, node.columns))
}
RqlQueryPlan::Filter(node) => transform::compile_filter(node, rx, context),
RqlQueryPlan::Gate(node) => {
let input_node = compile(*node.input, rx, context);
Box::new(FilterNode::new(input_node, node.conditions, "gate"))
}
RqlQueryPlan::Map(node) => transform::compile_map(node, rx, context),
RqlQueryPlan::Extend(node) => transform::compile_extend(node, rx, context),
RqlQueryPlan::Patch(node) => transform::compile_patch(node, rx, context),
RqlQueryPlan::Sort(RqlSortNode {
by,
input,
}) => {
let input_node = compile(*input, rx, context);
Box::new(SortNode::new(input_node, by))
}
RqlQueryPlan::Take(RqlTakeNode {
take,
input,
}) => {
let limit = match take {
TakeLimit::Literal(n) => n,
TakeLimit::Variable(ref name) => context
.symbols
.get(name)
.and_then(|var| match var {
Variable::Columns {
batch,
} => user_columns(batch).next().and_then(|(field, array)| {
ColumnView::try_from((array, field.as_ref()))
.unwrap_or_else(|err| {
panic!(
"TAKE variable ${} must be readable: {}",
name, err
)
})
.get_value(0)
.to_usize()
}),
_ => None,
})
.unwrap_or_else(|| panic!("TAKE variable ${} must be a numeric value", name)),
};
if let RqlQueryPlan::Sort(sort_node) = *input {
let input_node = compile(*sort_node.input, rx, context);
return Box::new(TopKNode::new(input_node, sort_node.by, limit));
}
let input_node = compile(*input, rx, context);
Box::new(TakeNode::new(input_node, limit))
}
RqlQueryPlan::JoinInner(node) => join::compile_inner_join(node, rx, context),
RqlQueryPlan::JoinLeft(node) => join::compile_left_join(node, rx, context),
RqlQueryPlan::JoinNatural(node) => join::compile_natural_join(node, rx, context),
RqlQueryPlan::Lookup(node) => {
Box::new(UnsupportedLookupNode::new(node.alias.unwrap_or(Fragment::None)))
}
RqlQueryPlan::Assert(RqlAssertNode {
conditions,
input,
message,
}) => {
if let Some(input) = input {
let input_node = compile(*input, rx, context);
Box::new(AssertNode::new(input_node, conditions, message))
} else {
Box::new(AssertWithoutInputNode::new(conditions, message))
}
}
RqlQueryPlan::TableScan(node) => match rx.layout() {
ScanLayout::Row => {
Box::new(TableScanNode::new(node.source.clone(), node.partition, context, rx).unwrap())
}
#[cfg(feature = "column")]
ScanLayout::Column if node.partition.is_some() => Box::new(UnsupportedColumnScanNode::new(format!(
"a partition scan of table '{}'",
node.source.fully_qualified_name()
))),
#[cfg(feature = "column")]
ScanLayout::Column => Box::new(ColumnTableScanNode::new(node.source.clone(), context)),
#[cfg(not(feature = "column"))]
ScanLayout::Column => unreachable!("column scan layout requires the engine column feature"),
},
RqlQueryPlan::ViewScan(node) => {
Box::new(ViewScanNode::new(node.source.clone(), node.partition, context, rx).unwrap())
}
RqlQueryPlan::RingBufferScan(node) => {
Box::new(RingBufferScan::new(node.source.clone(), context, rx).unwrap())
}
RqlQueryPlan::QueueScan(node) => Box::new(QueueScan::new(node.source.clone(), context, rx).unwrap()),
RqlQueryPlan::DictionaryScan(node) => {
Box::new(DictionaryScanNode::new(node.source.clone(), context).unwrap())
}
RqlQueryPlan::SeriesScan(node) => match rx.layout() {
ScanLayout::Row => Box::new(
VolcanoSeriesScanNode::new(
node.source.clone(),
node.key_range_start,
node.key_range_end,
node.variant_tag,
node.partition,
context,
)
.unwrap(),
),
#[cfg(feature = "column")]
ScanLayout::Column => Box::new(ColumnSeriesScanNode::new(
node.source.clone(),
node.key_range_start,
node.key_range_end,
node.variant_tag,
node.partition,
context,
)),
#[cfg(not(feature = "column"))]
ScanLayout::Column => unreachable!("column scan layout requires the engine column feature"),
},
RqlQueryPlan::IndexScan(node) => {
let table = node.source.def().clone();
let Some(pk) = table.primary_key.clone() else {
unimplemented!()
};
Box::new(IndexScanNode::new(table, IndexId::primary(pk.id), context).unwrap())
}
RqlQueryPlan::RemoteScan(node) => {
Box::new(RemoteFetchNode::new(node.address, node.token, node.remote_rql, node.variables))
}
RqlQueryPlan::TableVirtualScan(node) => vtable::compile_virtual_scan(node, context),
RqlQueryPlan::RowPointLookup(RqlRowPointLookupNode {
source,
row_number,
}) => {
let resolved_source = source;
Box::new(
RowPointLookupNode::new(resolved_source, row_number, context)
.expect("Failed to create RowPointLookupNode"),
)
}
RqlQueryPlan::RowListLookup(RqlRowListLookupNode {
source,
row_numbers,
}) => {
let resolved_source = source;
Box::new(
RowListLookupNode::new(resolved_source, row_numbers, context)
.expect("Failed to create RowListLookupNode"),
)
}
RqlQueryPlan::RowRangeScan(RqlRowRangeScanNode {
source,
start,
end,
}) => {
let resolved_source = source;
Box::new(
RowRangeScanNode::new(resolved_source, start, end, context)
.expect("Failed to create RowRangeScanNode"),
)
}
RqlQueryPlan::PartitionParameterScan(RqlPartitionParameterScanNode {
keys,
scan,
}) => {
if let Some(partitions) = choose_partitions(&keys, rx.layout(), &context.symbols)
&& let RqlQueryPlan::TableScan(node) = scan.as_ref()
{
return Box::new(
TableScanNode::with_partitions(node.source.clone(), partitions, context, rx)
.unwrap(),
);
}
let scan = match (*scan, choose_partition(&keys, rx.layout(), &context.symbols)) {
(RqlQueryPlan::TableScan(node), Some(partition)) => {
RqlQueryPlan::TableScan(RqlTableScanNode {
partition: Some(partition),
..node
})
}
(RqlQueryPlan::ViewScan(node), Some(partition)) => {
RqlQueryPlan::ViewScan(RqlViewScanNode {
partition: Some(partition),
..node
})
}
(scan, _) => scan,
};
compile(scan, rx, context)
}
RqlQueryPlan::SeriesParameterScan(RqlSeriesParameterScanNode {
starts,
ends,
partition_keys,
scan,
}) => {
let scan = match *scan {
RqlQueryPlan::SeriesScan(node) => match choose_series_range(
&starts,
&ends,
partition_keys.as_deref(),
&node.source.def().key,
rx.layout(),
&context.symbols,
) {
Some(range) => RqlQueryPlan::SeriesScan(RqlSeriesScanNode {
key_range_start: node.key_range_start.max(range.start),
key_range_end: match (node.key_range_end, range.end) {
(Some(planned), Some(end)) => Some(planned.min(end)),
(planned, end) => planned.or(end),
},
partition: range.partition.or(node.partition),
..node
}),
None => RqlQueryPlan::SeriesScan(node),
},
other => other,
};
compile(scan, rx, context)
}
RqlQueryPlan::RowParameterLookup(RqlRowParameterLookupNode {
source,
predicate,
fallback,
}) => match choose_row_lookup(&predicate, rx.layout(), &context.symbols) {
Some(RowPredicate::Point(row_number)) => Box::new(
RowPointLookupNode::new(source, row_number, context)
.expect("Failed to create RowPointLookupNode"),
),
Some(RowPredicate::List(row_numbers)) => Box::new(
RowListLookupNode::new(source, row_numbers, context)
.expect("Failed to create RowListLookupNode"),
),
Some(RowPredicate::Range {
start,
end,
}) => Box::new(
RowRangeScanNode::new(source, start, end, context)
.expect("Failed to create RowRangeScanNode"),
),
None => compile(*fallback, rx, context),
},
RqlQueryPlan::InlineData(RqlInlineDataNode {
rows,
}) => Box::new(InlineDataNode::new(rows, context)),
RqlQueryPlan::Generator(RqlGeneratorNode {
name,
expressions,
}) => Box::new(GeneratorNode::new(name, expressions)),
RqlQueryPlan::Variable(node) => Box::new(VariableNode::new(node.variable_expr)),
RqlQueryPlan::Environment(_) => Box::new(EnvironmentNode::new()),
RqlQueryPlan::Apply(node) => {
let Some(transform) = context.services.transforms.get_transform(node.operator.text()) else {
return Box::new(UnknownTransformNode::new(node.operator));
};
let input = node.input.expect("Apply requires input");
let input_node = compile(*input, rx, context);
Box::new(ApplyTransformNode::new(input_node, transform))
}
RqlQueryPlan::RunTests(node) => Box::new(RunTestsQueryNode::new(node, context.clone())),
RqlQueryPlan::CallFunction(node) => Box::new(GeneratorNode::new(node.name, node.arguments)),
RqlQueryPlan::Window(node) => Box::new(UnsupportedWindowNode::new(node.fragment)),
RqlQueryPlan::Append(node) => Box::new(UnsupportedAppendNode::new(node.fragment)),
}
}