use std::sync::Arc;
use arrow::datatypes::{DataType, Field};
use datafusion_common::datatype::DataTypeExt;
use datafusion_common::{
Result, ScalarValue, SplitPoint, TableReference, exec_datafusion_err, internal_err,
plan_datafusion_err,
};
use datafusion_execution::TaskContext;
use datafusion_execution::registry::FunctionRegistry;
use datafusion_expr::dml::{
InsertOp, MergeIntoAction, MergeIntoClause, MergeIntoClauseKind, MergeIntoOp,
};
use datafusion_expr::expr::{
Alias, Lambda, LambdaVariable, NullTreatment, Placeholder, Sort,
};
use datafusion_expr::expr::{Unnest, WildcardOptions};
use datafusion_expr::logical_plan::Subquery;
use datafusion_expr::{
Between, BinaryExpr, Case, Cast, Expr, GroupingSet,
GroupingSet::GroupingSets,
Like, Operator, TryCast, WindowFrame,
expr::{self, InList, WindowFunction},
};
use datafusion_expr::{ExprFunctionExt, WriteOp};
use datafusion_proto_common::{FromProtoError as Error, from_proto::FromOptionalField};
use crate::protobuf::{self, CubeNode, GroupingSetNode, PlaceholderNode, RollupNode};
use super::{AsLogicalPlan, LogicalExtensionCodec};
pub fn parse_write_op(
node: &protobuf::DmlNode,
ctx: &TaskContext,
codec: &dyn LogicalExtensionCodec,
) -> Result<WriteOp, Error> {
let typ = node.dml_type();
Ok(match typ {
protobuf::dml_node::Type::Update => WriteOp::Update,
protobuf::dml_node::Type::Delete => WriteOp::Delete,
protobuf::dml_node::Type::InsertAppend => WriteOp::Insert(InsertOp::Append),
protobuf::dml_node::Type::InsertOverwrite => WriteOp::Insert(InsertOp::Overwrite),
protobuf::dml_node::Type::InsertReplace => WriteOp::Insert(InsertOp::Replace),
protobuf::dml_node::Type::Ctas => WriteOp::Ctas,
protobuf::dml_node::Type::Truncate => WriteOp::Truncate,
protobuf::dml_node::Type::MergeInto => {
let merge_into = node.merge_into.as_deref().ok_or_else(|| {
Error::General(
"DmlNode with MERGE_INTO type is missing the merge_into payload"
.to_string(),
)
})?;
WriteOp::MergeInto(Box::new(parse_merge_into_op(merge_into, ctx, codec)?))
}
})
}
fn parse_merge_into_op(
op: &protobuf::MergeIntoOpNode,
ctx: &TaskContext,
codec: &dyn LogicalExtensionCodec,
) -> Result<MergeIntoOp, Error> {
let on = op.on.as_ref().ok_or_else(|| {
Error::General("MergeIntoOpNode is missing required `on` expression".to_string())
})?;
let on = parse_expr(on, ctx, codec)?;
let clauses = op
.clauses
.iter()
.map(|c| parse_merge_into_clause(c, ctx, codec))
.collect::<Result<Vec<_>, Error>>()?;
Ok(MergeIntoOp { on, clauses })
}
fn parse_merge_into_clause(
clause: &protobuf::MergeIntoClauseNode,
ctx: &TaskContext,
codec: &dyn LogicalExtensionCodec,
) -> Result<MergeIntoClause, Error> {
let kind = protobuf::merge_into_clause_node::Kind::try_from(clause.kind)
.map_err(|_| {
Error::General(format!(
"MergeIntoClauseNode has unknown kind tag {}",
clause.kind
))
})
.map(MergeIntoClauseKind::from)?;
let predicate = clause
.predicate
.as_ref()
.map(|e| parse_expr(e, ctx, codec))
.transpose()?;
let action = clause.action.as_ref().ok_or_else(|| {
Error::General("MergeIntoClauseNode is missing required `action`".to_string())
})?;
let action = parse_merge_into_action(action, ctx, codec)?;
Ok(MergeIntoClause {
kind,
predicate,
action,
})
}
fn parse_merge_into_action(
action: &protobuf::MergeIntoActionNode,
ctx: &TaskContext,
codec: &dyn LogicalExtensionCodec,
) -> Result<MergeIntoAction, Error> {
use protobuf::merge_into_action_node::Action;
let action = action.action.as_ref().ok_or_else(|| {
Error::General("MergeIntoActionNode is missing the `action` oneof".to_string())
})?;
Ok(match action {
Action::Update(update) => {
let assignments = update
.assignments
.iter()
.map(|a| {
let value = a.value.as_ref().ok_or_else(|| {
Error::General(format!(
"MergeAssignment for column `{}` is missing its value",
a.column
))
})?;
Ok((a.column.clone(), parse_expr(value, ctx, codec)?))
})
.collect::<Result<Vec<_>, Error>>()?;
MergeIntoAction::Update(assignments)
}
Action::Insert(insert) => MergeIntoAction::Insert {
columns: insert.columns.clone(),
values: parse_exprs(&insert.values, ctx, codec)?,
},
Action::Delete(_) => MergeIntoAction::Delete,
})
}
pub fn parse_expr(
proto: &protobuf::LogicalExprNode,
ctx: &TaskContext,
codec: &dyn LogicalExtensionCodec,
) -> Result<Expr, Error> {
use protobuf::{logical_expr_node::ExprType, window_expr_node};
let expr_type = proto
.expr_type
.as_ref()
.ok_or_else(|| Error::required("expr_type"))?;
match expr_type {
ExprType::BinaryExpr(binary_expr) => {
let op = from_proto_binary_op(&binary_expr.op)?;
let operands = parse_exprs(&binary_expr.operands, ctx, codec)?;
if operands.len() < 2 {
return Err(proto_error(
"A binary expression must always have at least 2 operands",
));
}
Ok(operands
.into_iter()
.reduce(|left, right| {
Expr::BinaryExpr(BinaryExpr::new(Box::new(left), op, Box::new(right)))
})
.expect("Binary expression could not be reduced to a single expression."))
}
ExprType::Column(column) => Ok(Expr::Column(column.into())),
ExprType::Literal(literal) => {
let scalar_value: ScalarValue = literal.try_into()?;
Ok(Expr::Literal(scalar_value, None))
}
ExprType::WindowExpr(expr) => {
let window_function = expr
.window_function
.as_ref()
.ok_or_else(|| Error::required("window_function"))?;
let partition_by = parse_exprs(&expr.partition_by, ctx, codec)?;
let mut order_by = parse_sorts(&expr.order_by, ctx, codec)?;
let window_frame = expr
.window_frame
.as_ref()
.map::<Result<WindowFrame, _>, _>(|window_frame| {
let window_frame = WindowFrame::try_from(window_frame.clone())?;
window_frame
.regularize_order_bys(&mut order_by)
.map(|_| window_frame)
})
.transpose()?
.ok_or_else(|| {
exec_datafusion_err!("missing window frame during deserialization")
})?;
let null_treatment = match expr.null_treatment {
Some(null_treatment) => {
let null_treatment = protobuf::NullTreatment::try_from(null_treatment)
.map_err(|_| {
proto_error(format!(
"Received a WindowExprNode message with unknown NullTreatment {null_treatment}",
))
})?;
Some(NullTreatment::from(null_treatment))
}
None => None,
};
let agg_fn = match window_function {
window_expr_node::WindowFunction::Udaf(udaf_name) => {
let udaf_function = match &expr.fun_definition {
Some(buf) => codec.try_decode_udaf(udaf_name, buf)?,
None => ctx
.udaf(udaf_name)
.or_else(|_| codec.try_decode_udaf(udaf_name, &[]))?,
};
expr::WindowFunctionDefinition::AggregateUDF(udaf_function)
}
window_expr_node::WindowFunction::Udwf(udwf_name) => {
let udwf_function = match &expr.fun_definition {
Some(buf) => codec.try_decode_udwf(udwf_name, buf)?,
None => ctx
.udwf(udwf_name)
.or_else(|_| codec.try_decode_udwf(udwf_name, &[]))?,
};
expr::WindowFunctionDefinition::WindowUDF(udwf_function)
}
};
let args = parse_exprs(&expr.exprs, ctx, codec)?;
let mut builder = Expr::from(WindowFunction::new(agg_fn, args))
.partition_by(partition_by)
.order_by(order_by)
.window_frame(window_frame)
.null_treatment(null_treatment);
if expr.distinct {
builder = builder.distinct();
};
if let Some(filter) = parse_optional_expr(expr.filter.as_deref(), ctx, codec)?
{
builder = builder.filter(filter);
}
builder.build().map_err(Error::DataFusionError)
}
ExprType::Alias(alias) => Ok(Expr::Alias(Alias::new(
parse_required_expr(alias.expr.as_deref(), ctx, "expr", codec)?,
alias
.relation
.first()
.map(|r| TableReference::try_from(r.clone()))
.transpose()?,
alias.alias.clone(),
))),
ExprType::IsNullExpr(is_null) => Ok(Expr::IsNull(Box::new(parse_required_expr(
is_null.expr.as_deref(),
ctx,
"expr",
codec,
)?))),
ExprType::IsNotNullExpr(is_not_null) => Ok(Expr::IsNotNull(Box::new(
parse_required_expr(is_not_null.expr.as_deref(), ctx, "expr", codec)?,
))),
ExprType::NotExpr(not) => Ok(Expr::Not(Box::new(parse_required_expr(
not.expr.as_deref(),
ctx,
"expr",
codec,
)?))),
ExprType::IsTrue(msg) => Ok(Expr::IsTrue(Box::new(parse_required_expr(
msg.expr.as_deref(),
ctx,
"expr",
codec,
)?))),
ExprType::IsFalse(msg) => Ok(Expr::IsFalse(Box::new(parse_required_expr(
msg.expr.as_deref(),
ctx,
"expr",
codec,
)?))),
ExprType::IsUnknown(msg) => Ok(Expr::IsUnknown(Box::new(parse_required_expr(
msg.expr.as_deref(),
ctx,
"expr",
codec,
)?))),
ExprType::IsNotTrue(msg) => Ok(Expr::IsNotTrue(Box::new(parse_required_expr(
msg.expr.as_deref(),
ctx,
"expr",
codec,
)?))),
ExprType::IsNotFalse(msg) => Ok(Expr::IsNotFalse(Box::new(parse_required_expr(
msg.expr.as_deref(),
ctx,
"expr",
codec,
)?))),
ExprType::IsNotUnknown(msg) => Ok(Expr::IsNotUnknown(Box::new(
parse_required_expr(msg.expr.as_deref(), ctx, "expr", codec)?,
))),
ExprType::Between(between) => Ok(Expr::Between(Between::new(
Box::new(parse_required_expr(
between.expr.as_deref(),
ctx,
"expr",
codec,
)?),
between.negated,
Box::new(parse_required_expr(
between.low.as_deref(),
ctx,
"expr",
codec,
)?),
Box::new(parse_required_expr(
between.high.as_deref(),
ctx,
"expr",
codec,
)?),
))),
ExprType::Like(like) => Ok(Expr::Like(Like::new(
like.negated,
Box::new(parse_required_expr(
like.expr.as_deref(),
ctx,
"expr",
codec,
)?),
Box::new(parse_required_expr(
like.pattern.as_deref(),
ctx,
"pattern",
codec,
)?),
parse_escape_char(&like.escape_char)?,
false,
))),
ExprType::Ilike(like) => Ok(Expr::Like(Like::new(
like.negated,
Box::new(parse_required_expr(
like.expr.as_deref(),
ctx,
"expr",
codec,
)?),
Box::new(parse_required_expr(
like.pattern.as_deref(),
ctx,
"pattern",
codec,
)?),
parse_escape_char(&like.escape_char)?,
true,
))),
ExprType::SimilarTo(like) => Ok(Expr::SimilarTo(Like::new(
like.negated,
Box::new(parse_required_expr(
like.expr.as_deref(),
ctx,
"expr",
codec,
)?),
Box::new(parse_required_expr(
like.pattern.as_deref(),
ctx,
"pattern",
codec,
)?),
parse_escape_char(&like.escape_char)?,
false,
))),
ExprType::Case(case) => {
let when_then_expr = case
.when_then_expr
.iter()
.map(|e| {
let when_expr = parse_required_expr(
e.when_expr.as_ref(),
ctx,
"when_expr",
codec,
)?;
let then_expr = parse_required_expr(
e.then_expr.as_ref(),
ctx,
"then_expr",
codec,
)?;
Ok((Box::new(when_expr), Box::new(then_expr)))
})
.collect::<Result<Vec<(Box<Expr>, Box<Expr>)>, Error>>()?;
Ok(Expr::Case(Case::new(
parse_optional_expr(case.expr.as_deref(), ctx, codec)?.map(Box::new),
when_then_expr,
parse_optional_expr(case.else_expr.as_deref(), ctx, codec)?.map(Box::new),
)))
}
ExprType::Cast(cast) => {
let expr = Box::new(parse_required_expr(
cast.expr.as_deref(),
ctx,
"expr",
codec,
)?);
let data_type: DataType = cast.arrow_type.as_ref().required("arrow_type")?;
let field = data_type
.into_nullable_field()
.with_nullable(cast.nullable.unwrap_or(true));
Ok(Expr::Cast(Cast::new_from_field(expr, Arc::new(field))))
}
ExprType::TryCast(cast) => {
let expr = Box::new(parse_required_expr(
cast.expr.as_deref(),
ctx,
"expr",
codec,
)?);
let data_type: DataType = cast.arrow_type.as_ref().required("arrow_type")?;
let field = data_type
.into_nullable_field()
.with_nullable(cast.nullable.unwrap_or(true));
Ok(Expr::TryCast(TryCast::new_from_field(
expr,
Arc::new(field),
)))
}
ExprType::Negative(negative) => Ok(Expr::Negative(Box::new(
parse_required_expr(negative.expr.as_deref(), ctx, "expr", codec)?,
))),
ExprType::Unnest(unnest) => {
let mut exprs = parse_exprs(&unnest.exprs, ctx, codec)?;
if exprs.len() != 1 {
return Err(proto_error("Unnest must have exactly one expression"));
}
Ok(Expr::Unnest(Unnest {
expr: Box::new(exprs.swap_remove(0)),
outer: unnest.outer,
}))
}
ExprType::InList(in_list) => Ok(Expr::InList(InList::new(
Box::new(parse_required_expr(
in_list.expr.as_deref(),
ctx,
"expr",
codec,
)?),
parse_exprs(&in_list.list, ctx, codec)?,
in_list.negated,
))),
ExprType::Wildcard(protobuf::Wildcard { qualifier }) => {
let qualifier = qualifier
.to_owned()
.map(TableReference::try_from)
.transpose()?;
#[expect(deprecated)]
Ok(Expr::Wildcard {
qualifier,
options: Box::new(WildcardOptions::default()),
})
}
ExprType::ScalarUdfExpr(protobuf::ScalarUdfExprNode {
fun_name,
args,
fun_definition,
}) => {
let scalar_fn = match fun_definition {
Some(buf) => codec.try_decode_udf(fun_name, buf)?,
None => ctx
.udf(fun_name.as_str())
.or_else(|_| codec.try_decode_udf(fun_name, &[]))?,
};
Ok(Expr::ScalarFunction(expr::ScalarFunction::new_udf(
scalar_fn,
parse_exprs(args, ctx, codec)?,
)))
}
ExprType::HigherOrderUdfExpr(protobuf::HigherOrderUdfExprNode {
fun_name,
args,
fun_definition,
}) => {
let hof_fn = match fun_definition {
Some(buf) => codec.try_decode_higher_order_function(fun_name, buf)?,
None => ctx
.higher_order_function(fun_name.as_str())
.or_else(|_| codec.try_decode_higher_order_function(fun_name, &[]))?,
};
Ok(Expr::HigherOrderFunction(expr::HigherOrderFunction::new(
hof_fn,
parse_exprs(args, ctx, codec)?,
)))
}
ExprType::AggregateUdfExpr(pb) => {
let agg_fn = match &pb.fun_definition {
Some(buf) => codec.try_decode_udaf(&pb.fun_name, buf)?,
None => ctx
.udaf(&pb.fun_name)
.or_else(|_| codec.try_decode_udaf(&pb.fun_name, &[]))?,
};
let null_treatment = match pb.null_treatment {
Some(null_treatment) => {
let null_treatment = protobuf::NullTreatment::try_from(null_treatment)
.map_err(|_| {
proto_error(format!(
"Received an AggregateUdfExprNode message with unknown NullTreatment {null_treatment}",
))
})?;
Some(NullTreatment::from(null_treatment))
}
None => None,
};
Ok(Expr::AggregateFunction(expr::AggregateFunction::new_udf(
agg_fn,
parse_exprs(&pb.args, ctx, codec)?,
pb.distinct,
parse_optional_expr(pb.filter.as_deref(), ctx, codec)?.map(Box::new),
parse_sorts(&pb.order_by, ctx, codec)?,
null_treatment,
)))
}
ExprType::GroupingSet(GroupingSetNode { expr }) => {
Ok(Expr::GroupingSet(GroupingSets(
expr.iter()
.map(|expr_list| parse_exprs(&expr_list.expr, ctx, codec))
.collect::<Result<Vec<_>, Error>>()?,
)))
}
ExprType::Cube(CubeNode { expr }) => Ok(Expr::GroupingSet(GroupingSet::Cube(
parse_exprs(expr, ctx, codec)?,
))),
ExprType::Rollup(RollupNode { expr }) => Ok(Expr::GroupingSet(
GroupingSet::Rollup(parse_exprs(expr, ctx, codec)?),
)),
ExprType::Placeholder(PlaceholderNode {
id,
data_type,
nullable,
metadata,
}) => match data_type {
None => Ok(Expr::Placeholder(Placeholder::new_with_field(
id.clone(),
None,
))),
Some(data_type) => {
let field =
Field::new("", data_type.try_into()?, nullable.unwrap_or(true))
.with_metadata(metadata.clone());
Ok(Expr::Placeholder(Placeholder::new_with_field(
id.clone(),
Some(field.into()),
)))
}
},
ExprType::ScalarSubqueryExpr(sq) => {
let subquery = parse_subquery(
sq.subquery
.as_deref()
.ok_or_else(|| Error::required("ScalarSubqueryExprNode.subquery"))?,
ctx,
codec,
)?;
Ok(Expr::ScalarSubquery(subquery))
}
ExprType::Lambda(lambda) => Ok(Expr::Lambda(Lambda::new(
lambda.params.clone(),
parse_required_expr(lambda.body.as_deref(), ctx, "body", codec)?,
))),
ExprType::LambdaVariable(lambda_variable) => {
Ok(Expr::LambdaVariable(LambdaVariable::new(
lambda_variable.name.clone(),
lambda_variable.field.as_ref().optional()?.map(Arc::new),
)))
}
}
}
fn parse_subquery(
proto: &protobuf::SubqueryNode,
ctx: &TaskContext,
codec: &dyn LogicalExtensionCodec,
) -> Result<Subquery, Error> {
let plan_node = proto
.subquery
.as_ref()
.ok_or_else(|| Error::required("SubqueryNode.subquery"))?;
let plan = plan_node.try_into_logical_plan(ctx, codec)?;
let outer_ref_columns = parse_exprs(&proto.outer_ref_columns, ctx, codec)?;
Ok(Subquery {
subquery: Arc::new(plan),
outer_ref_columns,
spans: Default::default(),
})
}
pub fn parse_exprs<'a, I>(
protos: I,
ctx: &TaskContext,
codec: &dyn LogicalExtensionCodec,
) -> Result<Vec<Expr>, Error>
where
I: IntoIterator<Item = &'a protobuf::LogicalExprNode>,
{
let res = protos
.into_iter()
.map(|elem| {
parse_expr(elem, ctx, codec).map_err(|e| plan_datafusion_err!("{}", e))
})
.collect::<Result<Vec<_>>>()?;
Ok(res)
}
pub fn parse_sorts<'a, I>(
protos: I,
ctx: &TaskContext,
codec: &dyn LogicalExtensionCodec,
) -> Result<Vec<Sort>, Error>
where
I: IntoIterator<Item = &'a protobuf::SortExprNode>,
{
protos
.into_iter()
.map(|sort| parse_sort(sort, ctx, codec))
.collect::<Result<Vec<Sort>, Error>>()
}
pub fn parse_sort(
sort: &protobuf::SortExprNode,
ctx: &TaskContext,
codec: &dyn LogicalExtensionCodec,
) -> Result<Sort, Error> {
Ok(Sort::new(
parse_required_expr(sort.expr.as_ref(), ctx, "expr", codec)?,
sort.asc,
sort.nulls_first,
))
}
fn parse_escape_char(s: &str) -> Result<Option<char>> {
match s.len() {
0 => Ok(None),
1 => Ok(s.chars().next()),
_ => internal_err!("Invalid length for escape char"),
}
}
pub fn from_proto_binary_op(op: &str) -> Result<Operator, Error> {
Operator::from_proto_name(op)
.ok_or_else(|| proto_error(format!("Unsupported binary operator '{op:?}'")))
}
fn parse_optional_expr(
p: Option<&protobuf::LogicalExprNode>,
ctx: &TaskContext,
codec: &dyn LogicalExtensionCodec,
) -> Result<Option<Expr>, Error> {
match p {
Some(expr) => parse_expr(expr, ctx, codec).map(Some),
None => Ok(None),
}
}
fn parse_required_expr(
p: Option<&protobuf::LogicalExprNode>,
ctx: &TaskContext,
field: impl Into<String>,
codec: &dyn LogicalExtensionCodec,
) -> Result<Expr, Error> {
match p {
Some(expr) => parse_expr(expr, ctx, codec),
None => Err(Error::required(field)),
}
}
fn proto_error<S: Into<String>>(message: S) -> Error {
Error::General(message.into())
}
pub(super) fn parse_protobuf_range_split_point(
split_point: &protobuf::RangeSplitPoint,
) -> Result<SplitPoint, Error> {
let values = split_point
.value
.iter()
.map(ScalarValue::try_from)
.collect::<Result<_, Error>>()?;
Ok(SplitPoint::new(values))
}