use std::{mem, sync::Arc};
use arrow_array::{ArrayRef, RecordBatch};
use arrow_buffer::BooleanBuffer;
use arrow_schema::FieldRef;
use reifydb_catalog::catalog::Catalog;
use reifydb_core::{
expression::{Expression, IsVariantExpression},
interface::resolved::ResolvedObject,
value::{batch::filter, column::headers::ColumnHeaders},
};
use reifydb_evaluate::{expression::context::EvalContext, lower::LoweredExpr};
use reifydb_extension::transform::{Transform, context::TransformContext};
use reifydb_transaction::transaction::Transaction;
use reifydb_value::{
reifydb_assertions,
value::column_view::{ColumnView, ViewData},
};
use tracing::instrument;
use super::NoopNode;
use crate::{
Result,
vm::volcano::{
inline::{resolve_is_variants, resolve_sumtype_ctor},
query::{QueryContext, QueryNode, eval_context_from_transform},
udf::{UdfEvalNode, strip_udf_columns},
},
};
pub(crate) struct FilterNode {
input: Box<dyn QueryNode>,
expressions: Vec<Expression>,
source: Option<ResolvedObject>,
operator: &'static str,
udf_names: Vec<String>,
context: Option<(Arc<QueryContext>, Vec<LoweredExpr>)>,
emitted: bool,
empty: Option<RecordBatch>,
}
impl FilterNode {
pub fn new(input: Box<dyn QueryNode>, expressions: Vec<Expression>, operator: &'static str) -> Self {
Self::with_source(input, expressions, None, operator)
}
pub fn with_source(
input: Box<dyn QueryNode>,
expressions: Vec<Expression>,
source: Option<ResolvedObject>,
operator: &'static str,
) -> Self {
Self {
input,
expressions,
source,
operator,
udf_names: Vec::new(),
context: None,
emitted: false,
empty: None,
}
}
#[instrument(level = "trace", skip_all, name = "volcano::filter::eval")]
fn eval_predicate(
session: &EvalContext,
lowered: &LoweredExpr,
columns: &RecordBatch,
row_count: usize,
) -> Result<(FieldRef, ArrayRef)> {
let exec_ctx = session.with_eval(columns.clone(), row_count);
lowered.evaluate(&exec_ctx)
}
#[instrument(level = "trace", skip_all, name = "volcano::filter::mask")]
fn build_mask(result: &ColumnView<'_>, row_count: usize) -> BooleanBuffer {
match &result.data {
ViewData::Bool(container) => {
BooleanBuffer::collect_bool(row_count, |i| result.is_defined(i) && container.value(i))
}
_ => panic!("filter expression must evaluate to a boolean column"),
}
}
#[instrument(level = "trace", skip_all, name = "volcano::filter::compact")]
fn compact(columns: &RecordBatch, mask: &BooleanBuffer) -> Result<RecordBatch> {
filter(columns, mask)
}
}
impl QueryNode for FilterNode {
#[instrument(level = "trace", skip_all, name = "volcano::filter::initialize")]
fn initialize<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &QueryContext) -> Result<()> {
if let Some(source) = self.source.as_ref() {
for expr in &mut self.expressions {
resolve_is_variants(&ctx.services.catalog, rx, source, expr)?;
resolve_variant_equality(&ctx.services.catalog, rx, source, expr)?;
}
}
let (input, expressions, udf_names) = UdfEvalNode::wrap_if_needed(
mem::replace(&mut self.input, Box::new(NoopNode)),
&self.expressions,
&ctx.symbols,
);
self.input = input;
self.expressions = expressions;
self.udf_names = udf_names;
let lowered = self.expressions.iter().map(|e| LoweredExpr::new(e.clone(), self.operator)).collect();
self.context = Some((Arc::new(ctx.clone()), lowered));
self.input.initialize(rx, ctx)?;
Ok(())
}
#[instrument(level = "trace", skip_all, name = "volcano::filter::next")]
fn next<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &mut QueryContext) -> Result<Option<RecordBatch>> {
reifydb_assertions! {
assert!(self.context.is_some(), "FilterNode::next() called before initialize()");
}
let (stored_ctx, _) = self.context.as_ref().unwrap();
let stored_ctx = stored_ctx.clone();
loop {
match self.input.next(rx, ctx)? {
Some(columns) => {
let transform_ctx = TransformContext {
routines: &ctx.services.routines,
runtime_context: &stored_ctx.services.runtime_context,
params: &stored_ctx.params,
};
let columns = self.apply(&transform_ctx, columns)?;
let columns = strip_udf_columns(columns, &self.udf_names)?;
if columns.num_rows() > 0 {
self.emitted = true;
return Ok(Some(columns));
}
if !self.emitted && self.empty.is_none() {
self.empty = Some(columns);
}
}
None => {
if self.emitted {
return Ok(None);
}
self.emitted = true;
return Ok(self.empty.take());
}
}
}
}
fn headers(&self) -> Option<ColumnHeaders> {
self.input.headers()
}
}
impl Transform for FilterNode {
fn apply(&self, ctx: &TransformContext, input: RecordBatch) -> Result<RecordBatch> {
let (stored_ctx, lowered) =
self.context.as_ref().expect("FilterNode::apply() called before initialize()");
let session = eval_context_from_transform(ctx, stored_ctx);
let mut columns = input;
let mut row_count = columns.num_rows();
for lowered_expr in lowered {
if row_count == 0 {
break;
}
let result = Self::eval_predicate(&session, lowered_expr, &columns, row_count)?;
let filter_mask = Self::build_mask(&ColumnView::try_from(&result)?, row_count);
columns = Self::compact(&columns, &filter_mask)?;
row_count = columns.num_rows();
}
Ok(columns)
}
}
pub(crate) fn resolve_variant_equality(
catalog: &Catalog,
rx: &mut Transaction<'_>,
source: &ResolvedObject,
expr: &mut Expression,
) -> Result<()> {
match expr {
Expression::And(e) => {
resolve_variant_equality(catalog, rx, source, &mut e.left)?;
resolve_variant_equality(catalog, rx, source, &mut e.right)
}
Expression::Or(e) => {
resolve_variant_equality(catalog, rx, source, &mut e.left)?;
resolve_variant_equality(catalog, rx, source, &mut e.right)
}
Expression::Prefix(e) => resolve_variant_equality(catalog, rx, source, &mut e.expression),
Expression::Equal(e) => {
let (Expression::Column(column), Expression::SumTypeConstructor(ctor)) =
(e.left.as_ref(), e.right.as_ref())
else {
return Ok(());
};
if !ctor.columns.is_empty() {
return Ok(());
}
let (sumtype, variant) = resolve_sumtype_ctor(catalog, rx, Some(source), &column.0.name, ctor)?;
let is_variant = Expression::IsVariant(IsVariantExpression {
expression: e.left.clone(),
namespace: Some(ctor.namespace.clone()),
sumtype_name: ctor.sumtype_name.clone(),
variant_name: ctor.variant_name.clone(),
tag: Some(sumtype.variants[variant].tag),
fragment: e.fragment.clone(),
});
*expr = is_variant;
Ok(())
}
_ => Ok(()),
}
}