mod arith;
mod cast;
mod compare;
mod conditional;
mod construct;
mod contains;
mod direct;
mod error;
mod field_access;
mod in_list;
mod literal;
mod prefix;
mod routine;
mod schema;
mod variant;
mod xor;
use std::sync::Arc;
use arrow_array::{ArrayRef, RecordBatch, RecordBatchOptions, new_empty_array};
use arrow_schema::{FieldRef, SchemaRef};
use datafusion_expr::{
BinaryExpr, Expr, Operator, ScalarUDF, ScalarUDFImpl, Signature, Volatility, execution_props::ExecutionProps,
expr::ScalarFunction, physical_planning_context::PhysicalPlanningContext,
};
use datafusion_physical_expr::{PhysicalExpr, create_physical_expr};
use reifydb_core::{
expression::{Expression, PrefixExpression, PrefixOperator, name::display_label},
internal_error,
value::column::{arith::ArithOp, factory::default_typed},
};
use reifydb_runtime::sync::mutex::Mutex;
use reifydb_value::{
error::{BinaryOp, Diagnostic, Error, LogicalOp},
fragment::Fragment,
value::{
column_view::ColumnView,
value_type::{
ValueType,
field::{FieldType, from_field, to_field},
},
},
};
use crate::{
Result,
expression::{context::EvalContext, logic::execute_logical_op, prefix::prefix_apply},
lower::{
cast::lower_cast,
direct::{Direct, Spliced},
error::from_datafusion,
schema::positional_schema,
},
};
pub struct LoweredExpr {
expression: Expression,
operator: &'static str,
state: Mutex<Option<Arc<Ready>>>,
}
struct Ready {
schema: SchemaRef,
plan: Plan,
}
enum Plan {
Lowered {
expr: Arc<dyn PhysicalExpr>,
field: FieldRef,
},
Direct(Direct),
Spliced(Spliced),
Failed(Box<Diagnostic>),
}
struct Node {
expr: Expr,
field: FieldRef,
}
enum Stop {
RowChanging,
Error(Error),
}
impl From<Error> for Stop {
fn from(error: Error) -> Self {
Stop::Error(error)
}
}
type Lowered<T> = std::result::Result<T, Stop>;
pub fn compile_check(expression: &Expression) -> Result<()> {
match variant::compile_time_error(expression) {
Some(error) => Err(error),
None => Ok(()),
}
}
impl LoweredExpr {
pub fn new(expression: Expression, operator: &'static str) -> Self {
Self {
expression,
operator,
state: Mutex::new(None),
}
}
pub fn access_column_name(&self) -> Option<&str> {
match &self.expression {
Expression::AccessSource(access) => Some(access.column.name.text()),
_ => None,
}
}
pub fn evaluate(&self, ctx: &EvalContext) -> Result<(FieldRef, ArrayRef)> {
let ready = self.ready(ctx);
match &ready.plan {
Plan::Lowered {
expr,
field,
} => {
let sized;
let batch = if ctx.batch.num_columns() == 0 && ctx.batch.num_rows() != ctx.row_count {
sized = RecordBatch::try_new_with_options(
ctx.batch.schema(),
vec![],
&RecordBatchOptions::new().with_row_count(Some(ctx.row_count)),
)
.map_err(|err| from_datafusion(self.operator, err.into()))?;
&sized
} else {
&ctx.batch
};
let value = expr.evaluate(batch).map_err(|err| from_datafusion(self.operator, err))?;
let array = value
.into_array(ctx.row_count)
.map_err(|err| from_datafusion(self.operator, err))?;
let array = match ctx.take {
Some(take) if take < array.len() => array.slice(0, take),
_ => array,
};
if !field.is_nullable() && array.logical_null_count() > 0 {
let mut field_type = from_field(field)?;
field_type.value_type = field_type
.value_type
.map(|value_type| ValueType::Option(Box::new(value_type)));
return Ok((Arc::new(to_field(field.name(), &field_type)), array));
}
Ok((field.clone(), array))
}
Plan::Direct(direct) => direct.evaluate(ctx),
Plan::Spliced(spliced) => spliced.evaluate(ctx),
Plan::Failed(diagnostic) => Err(Error(diagnostic.clone())),
}
}
pub fn output_field(&self, ctx: &EvalContext) -> Option<FieldRef> {
match &self.ready(ctx).plan {
Plan::Lowered {
field,
..
} => Some(field.clone()),
Plan::Direct(_) | Plan::Spliced(_) | Plan::Failed(_) => None,
}
}
fn ready(&self, ctx: &EvalContext) -> Arc<Ready> {
let schema = ctx.batch.schema();
let mut state = self.state.lock();
if let Some(ready) = state.as_ref()
&& same_schema(&ready.schema, &schema)
{
return ready.clone();
}
let ready = Arc::new(Ready {
plan: self.plan(ctx),
schema,
});
if !matches!(ready.plan, Plan::Failed(_)) {
*state = Some(ready.clone());
}
ready
}
fn plan(&self, ctx: &EvalContext) -> Plan {
match lower_root(ctx, self.operator, &self.expression) {
Ok((expr, field)) => Plan::Lowered {
expr,
field,
},
Err(Stop::Error(error)) => Plan::Failed(error.0),
Err(Stop::RowChanging) => {
let planned = match &self.expression {
Expression::Call(_) => {
Direct::new(ctx, self.operator, &self.expression).map(Plan::Direct)
}
Expression::Alias(alias)
if matches!(alias.expression.as_ref(), Expression::Call(_)) =>
{
Direct::new(ctx, self.operator, &self.expression).map(Plan::Direct)
}
_ => Spliced::new(ctx, self.operator, &self.expression).map(Plan::Spliced),
};
planned.unwrap_or_else(|error| Plan::Failed(error.0))
}
}
}
}
fn same_schema(left: &SchemaRef, right: &SchemaRef) -> bool {
Arc::ptr_eq(left, right) || left == right
}
fn lower_root<'e>(
ctx: &EvalContext,
operator: &'static str,
expression: &'e Expression,
) -> Lowered<(Arc<dyn PhysicalExpr>, FieldRef)> {
if let Some(error) = variant::compile_time_error(expression) {
return Err(Stop::Error(error));
}
let node = lower_node(ctx, operator, expression)?;
let schema = positional_schema(operator, ctx.batch.schema_ref())?;
let expr =
create_physical_expr(&node.expr, &schema, &ExecutionProps::new(), &PhysicalPlanningContext::default())
.map_err(|err| from_datafusion(operator, err))?;
Ok((expr, node.field))
}
fn lower_node<'e>(ctx: &EvalContext, operator: &'static str, expression: &'e Expression) -> Lowered<Node> {
match expression {
Expression::Column(column) => Ok(schema::column(ctx, operator, column)?),
Expression::AccessSource(access) => Ok(schema::access(ctx, access)?),
Expression::Constant(constant) => {
Ok(literal::constant(operator, constant, display_label(expression).text())?)
}
Expression::Prefix(prefix) => match prefix.operator {
PrefixOperator::Not(_) => lower_not(ctx, operator, expression, prefix),
PrefixOperator::Minus(_) | PrefixOperator::Plus(_) => {
prefix::lower_prefix(ctx, operator, expression, prefix)
}
},
Expression::And(and) => lower_logical(
ctx,
operator,
expression,
(&and.left, &and.right),
&and.full_fragment_owned(),
(LogicalOp::And, Operator::And),
),
Expression::Or(or) => lower_logical(
ctx,
operator,
expression,
(&or.left, &or.right),
&or.full_fragment_owned(),
(LogicalOp::Or, Operator::Or),
),
Expression::Equal(e) => compare::lower_compare(
ctx,
operator,
expression,
(&e.left, &e.right),
e.full_fragment_owned(),
(BinaryOp::Equal, Operator::Eq),
),
Expression::NotEqual(e) => compare::lower_compare(
ctx,
operator,
expression,
(&e.left, &e.right),
e.full_fragment_owned(),
(BinaryOp::NotEqual, Operator::NotEq),
),
Expression::LessThan(e) => compare::lower_compare(
ctx,
operator,
expression,
(&e.left, &e.right),
e.full_fragment_owned(),
(BinaryOp::LessThan, Operator::Lt),
),
Expression::LessThanEqual(e) => compare::lower_compare(
ctx,
operator,
expression,
(&e.left, &e.right),
e.full_fragment_owned(),
(BinaryOp::LessThanEqual, Operator::LtEq),
),
Expression::GreaterThan(e) => compare::lower_compare(
ctx,
operator,
expression,
(&e.left, &e.right),
e.full_fragment_owned(),
(BinaryOp::GreaterThan, Operator::Gt),
),
Expression::GreaterThanEqual(e) => compare::lower_compare(
ctx,
operator,
expression,
(&e.left, &e.right),
e.full_fragment_owned(),
(BinaryOp::GreaterThanEqual, Operator::GtEq),
),
Expression::Between(e) => compare::lower_between(ctx, operator, expression, e),
Expression::Add(e) => arith::lower_arith(
ctx,
operator,
expression,
(&e.left, &e.right),
e.full_fragment_owned(),
ArithOp::Add,
),
Expression::Sub(e) => arith::lower_arith(
ctx,
operator,
expression,
(&e.left, &e.right),
e.full_fragment_owned(),
ArithOp::Sub,
),
Expression::Mul(e) => arith::lower_arith(
ctx,
operator,
expression,
(&e.left, &e.right),
e.full_fragment_owned(),
ArithOp::Mul,
),
Expression::Div(e) => arith::lower_arith(
ctx,
operator,
expression,
(&e.left, &e.right),
e.full_fragment_owned(),
ArithOp::Div,
),
Expression::Rem(e) => arith::lower_arith(
ctx,
operator,
expression,
(&e.left, &e.right),
e.full_fragment_owned(),
ArithOp::Rem,
),
Expression::Call(call) => routine::lower_call(ctx, operator, expression, call),
Expression::Variable(variable) => Ok(literal::variable(ctx, operator, variable)?),
Expression::Parameter(parameter) => Ok(literal::parameter(ctx, operator, parameter)?),
Expression::Type(ty) => Ok(literal::type_literal(operator, &ty.ty, &ty.fragment)?),
Expression::FieldAccess(access) => match access.object.as_ref() {
Expression::Variable(variable) => {
Ok(literal::variable_field(ctx, operator, variable, access.field.text())?)
}
_ => field_access::lower_field_access(ctx, operator, expression, access),
},
Expression::Alias(alias) => {
let inner = lower_node(ctx, operator, &alias.expression)?;
Ok(Node {
expr: inner.expr,
field: Arc::new(inner.field.as_ref().clone().with_name(alias.alias.name())),
})
}
Expression::Cast(cast) => lower_cast(ctx, operator, expression, cast),
Expression::In(e) => in_list::lower_in(ctx, operator, expression, e),
Expression::If(e) => conditional::lower_if(ctx, operator, expression, e),
Expression::Xor(e) => xor::lower_xor(ctx, operator, expression, e),
Expression::Tuple(e) => construct::lower_construct(
ctx,
operator,
construct::Shape::Tuple,
e.expressions.iter().collect(),
&e.fragment,
),
Expression::List(e) => construct::lower_construct(
ctx,
operator,
construct::Shape::List,
e.expressions.iter().collect(),
&e.fragment,
),
Expression::Record(e) => construct::lower_construct(
ctx,
operator,
construct::Shape::Record(e.fields.iter().map(|(name, _)| name.text().to_string()).collect()),
e.fields.iter().map(|(_, expression)| expression).collect(),
&e.fragment,
),
Expression::Map(e) => construct::lower_projection(ctx, operator, &e.expressions, &e.fragment),
Expression::Extend(e) => construct::lower_projection(ctx, operator, &e.expressions, &e.fragment),
Expression::Contains(e) => contains::lower_contains(ctx, operator, expression, e),
Expression::IsVariant(e) => Ok(variant::lower_is_variant(ctx, operator, e)?),
Expression::SumTypeConstructor(_) => Err(Stop::Error(
variant::compile_time_error(expression).expect("a sum type constructor is a compile error"),
)),
}
}
fn lower_logical<'e>(
ctx: &EvalContext,
operator: &'static str,
expression: &'e Expression,
(left, right): (&'e Expression, &'e Expression),
fragment: &Fragment,
(logical_op, df_op): (LogicalOp, Operator),
) -> Lowered<Node> {
let left = lower_node(ctx, operator, left)?;
let right = lower_node(ctx, operator, right)?;
execute_logical_op(&logic_probe(&left.field)?, &logic_probe(&right.field)?, fragment, logical_op)?;
let nullable = left.field.is_nullable() || right.field.is_nullable();
Ok(Node {
expr: Expr::BinaryExpr(BinaryExpr::new(
Box::new(bool_operand(left)?),
df_op,
Box::new(bool_operand(right)?),
)),
field: bool_field(display_label(expression).text(), nullable),
})
}
fn lower_not<'e>(
ctx: &EvalContext,
operator: &'static str,
expression: &'e Expression,
prefix: &'e PrefixExpression,
) -> Lowered<Node> {
let inner = lower_node(ctx, operator, &prefix.expression)?;
if !is_untyped_none(&inner.field)? {
let fragment = prefix.full_fragment_owned();
prefix_apply(&empty_of(&inner.field), &prefix.operator, &fragment)?;
if let Some(value_type) = from_field(&inner.field)?.value_type {
prefix_apply(&default_typed("probe", value_type, 1), &prefix.operator, &fragment)?;
}
}
let nullable = inner.field.is_nullable();
Ok(Node {
expr: Expr::Not(Box::new(bool_operand(inner)?)),
field: bool_field(display_label(expression).text(), nullable),
})
}
fn logic_probe(field: &FieldRef) -> Result<(FieldRef, ArrayRef)> {
if is_untyped_none(field)? {
return Ok(empty_of(&bool_field("none", true)));
}
Ok(empty_of(&Arc::new(field.as_ref().clone().with_nullable(false))))
}
fn bool_operand(node: Node) -> Result<Expr> {
if is_untyped_none(&node.field)? {
return Ok(literal::none_bool());
}
Ok(node.expr)
}
fn is_untyped_none(field: &FieldRef) -> Result<bool> {
Ok(ColumnView::try_from(&empty_of(field))?.is_untyped_none())
}
fn inner_type(field: &FieldRef) -> Result<ValueType> {
match from_field(field)?.value_type {
Some(value_type) => Ok(value_type.inner_type().clone()),
None => Err(internal_error!("operand {} has no value type", field.name())),
}
}
fn empty_of(field: &FieldRef) -> (FieldRef, ArrayRef) {
(field.clone(), new_empty_array(field.data_type()))
}
fn bool_field(name: &str, nullable: bool) -> FieldRef {
typed_field(name, ValueType::Boolean, nullable)
}
fn typed_field(name: &str, value_type: ValueType, nullable: bool) -> FieldRef {
let value_type = match nullable {
true => ValueType::Option(Box::new(value_type)),
false => value_type,
};
Arc::new(to_field(name, &FieldType::from(value_type)))
}
fn udf(function: impl ScalarUDFImpl + 'static, args: Vec<Expr>) -> Expr {
Expr::ScalarFunction(ScalarFunction::new_udf(Arc::new(ScalarUDF::new_from_impl(function)), args))
}
fn volatile() -> Signature {
Signature::user_defined(Volatility::Volatile)
}