Skip to content

Commit

Permalink
Add support for parsing common JSON operators
Browse files Browse the repository at this point in the history
  • Loading branch information
dmontagu committed Jun 17, 2024
1 parent f373a86 commit 32eb462
Show file tree
Hide file tree
Showing 9 changed files with 65 additions and 5 deletions.
20 changes: 19 additions & 1 deletion datafusion/expr/src/operator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,12 @@ pub enum Operator {
BitwiseShiftLeft,
/// String concat
StringConcat,
/// Question, like `?`
Question,
/// Arrow, like `->`
Arrow,
/// Long arrow, like `->>`
LongArrow,
/// At arrow, like `@>`
AtArrow,
/// Arrow at, like `<@`
Expand Down Expand Up @@ -129,6 +135,9 @@ impl Operator {
| Operator::BitwiseShiftRight
| Operator::BitwiseShiftLeft
| Operator::StringConcat
| Operator::Question
| Operator::Arrow
| Operator::LongArrow
| Operator::AtArrow
| Operator::ArrowAt => None,
}
Expand Down Expand Up @@ -213,7 +222,10 @@ impl Operator {
| Operator::BitwiseXor
| Operator::BitwiseShiftRight
| Operator::BitwiseShiftLeft
| Operator::StringConcat => None,
| Operator::StringConcat
| Operator::Question
| Operator::Arrow
| Operator::LongArrow => None,
}
}

Expand Down Expand Up @@ -247,6 +259,9 @@ impl Operator {
| Operator::BitwiseShiftRight
| Operator::BitwiseXor
| Operator::StringConcat
| Operator::Question
| Operator::Arrow
| Operator::LongArrow
| Operator::AtArrow
| Operator::ArrowAt => 0,
}
Expand Down Expand Up @@ -285,6 +300,9 @@ impl fmt::Display for Operator {
Operator::BitwiseShiftRight => ">>",
Operator::BitwiseShiftLeft => "<<",
Operator::StringConcat => "||",
Operator::Question => "?",
Operator::Arrow => "->",
Operator::LongArrow => "->>",
Operator::AtArrow => "@>",
Operator::ArrowAt => "<@",
};
Expand Down
24 changes: 24 additions & 0 deletions datafusion/expr/src/type_coercion/binary.rs
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,30 @@ fn signature(lhs: &DataType, op: &Operator, rhs: &DataType) -> Result<Signature>
)
})
}
Question => {
match (lhs, rhs) {
(Utf8 | LargeUtf8, Utf8 | LargeUtf8) => Ok(Signature{
lhs: lhs.clone(),
rhs: rhs.clone(),
ret: Boolean,
}),
_ => Err(plan_datafusion_err!(
"Cannot coerce question operation {lhs} {op} {rhs} to vaild types"
))
}
}
Arrow | LongArrow => {
match (lhs, rhs) {
(Utf8 | LargeUtf8, Utf8 | LargeUtf8) => Ok(Signature{
lhs: lhs.clone(),
rhs: rhs.clone(),
ret: lhs.clone(),
}),
_ => Err(plan_datafusion_err!(
"Cannot coerce arrow operation {lhs} {op} {rhs} to vaild types"
))
}
}
AtArrow | ArrowAt => {
// ArrowAt and AtArrow check for whether one array is contained in another.
// The result type is boolean. Signature::comparison defines this signature.
Expand Down
4 changes: 2 additions & 2 deletions datafusion/physical-expr/src/expressions/binary.rs
Original file line number Diff line number Diff line change
Expand Up @@ -628,8 +628,8 @@ impl BinaryExpr {
BitwiseShiftRight => bitwise_shift_right_dyn(left, right),
BitwiseShiftLeft => bitwise_shift_left_dyn(left, right),
StringConcat => binary_string_array_op!(left, right, concat_elements),
AtArrow | ArrowAt => {
unreachable!("ArrowAt and AtArrow should be rewritten to function")
Question | Arrow | LongArrow | AtArrow | ArrowAt => {
unreachable!("Question, Arrow, LongArrow, ArrowAt, and AtArrow should be rewritten to function")
}
}
}
Expand Down
3 changes: 3 additions & 0 deletions datafusion/proto/src/logical_plan/from_proto.rs
Original file line number Diff line number Diff line change
Expand Up @@ -707,6 +707,9 @@ pub fn from_proto_binary_op(op: &str) -> Result<Operator, Error> {
"RegexNotIMatch" => Ok(Operator::RegexNotIMatch),
"RegexNotMatch" => Ok(Operator::RegexNotMatch),
"StringConcat" => Ok(Operator::StringConcat),
"Question" => Ok(Operator::Question),
"Arrow" => Ok(Operator::Arrow),
"LongArrow" => Ok(Operator::LongArrow),
"AtArrow" => Ok(Operator::AtArrow),
"ArrowAt" => Ok(Operator::ArrowAt),
other => Err(proto_error(format!(
Expand Down
3 changes: 3 additions & 0 deletions datafusion/proto/tests/cases/roundtrip_logical_plan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1595,6 +1595,9 @@ fn roundtrip_binary_op() {
let ctx = SessionContext::new();
roundtrip_expr_test(test_expr, ctx);
}
test(Operator::Question);
test(Operator::Arrow);
test(Operator::LongArrow);
test(Operator::ArrowAt);
test(Operator::AtArrow);
test(Operator::StringConcat);
Expand Down
3 changes: 3 additions & 0 deletions datafusion/sql/src/expr/binary_op.rs
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,9 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
BinaryOperator::PGBitwiseShiftRight => Ok(Operator::BitwiseShiftRight),
BinaryOperator::PGBitwiseShiftLeft => Ok(Operator::BitwiseShiftLeft),
BinaryOperator::StringConcat => Ok(Operator::StringConcat),
BinaryOperator::Question => Ok(Operator::Question),
BinaryOperator::Arrow => Ok(Operator::Arrow),
BinaryOperator::LongArrow => Ok(Operator::LongArrow),
BinaryOperator::ArrowAt => Ok(Operator::ArrowAt),
BinaryOperator::AtArrow => Ok(Operator::AtArrow),
_ => not_impl_err!("Unsupported SQL binary operator {op:?}"),
Expand Down
7 changes: 5 additions & 2 deletions datafusion/sql/src/unparser/expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -639,8 +639,11 @@ impl Unparser<'_> {
Operator::BitwiseShiftRight => Ok(ast::BinaryOperator::PGBitwiseShiftRight),
Operator::BitwiseShiftLeft => Ok(ast::BinaryOperator::PGBitwiseShiftLeft),
Operator::StringConcat => Ok(ast::BinaryOperator::StringConcat),
Operator::AtArrow => not_impl_err!("unsupported operation: {op:?}"),
Operator::ArrowAt => not_impl_err!("unsupported operation: {op:?}"),
Operator::Question => Ok(ast::BinaryOperator::Question),
Operator::Arrow => Ok(ast::BinaryOperator::Arrow),
Operator::LongArrow => Ok(ast::BinaryOperator::LongArrow),
Operator::AtArrow => Ok(ast::BinaryOperator::AtArrow),
Operator::ArrowAt => Ok(ast::BinaryOperator::ArrowAt),
}
}

Expand Down
3 changes: 3 additions & 0 deletions datafusion/substrait/src/logical_plan/consumer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,9 @@ pub fn name_to_op(name: &str) -> Result<Operator> {
"bitwise_and" => Ok(Operator::BitwiseAnd),
"bitwise_or" => Ok(Operator::BitwiseOr),
"str_concat" => Ok(Operator::StringConcat),
"question" => Ok(Operator::Question),
"arrow" => Ok(Operator::Arrow),
"long_arrow" => Ok(Operator::LongArrow),
"at_arrow" => Ok(Operator::AtArrow),
"arrow_at" => Ok(Operator::ArrowAt),
"bitwise_xor" => Ok(Operator::BitwiseXor),
Expand Down
3 changes: 3 additions & 0 deletions datafusion/substrait/src/logical_plan/producer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -713,6 +713,9 @@ pub fn operator_to_name(op: Operator) -> &'static str {
Operator::BitwiseAnd => "bitwise_and",
Operator::BitwiseOr => "bitwise_or",
Operator::StringConcat => "str_concat",
Operator::Question => "question",
Operator::Arrow => "arrow",
Operator::LongArrow => "long_arrow",
Operator::AtArrow => "at_arrow",
Operator::ArrowAt => "arrow_at",
Operator::BitwiseXor => "bitwise_xor",
Expand Down

0 comments on commit 32eb462

Please sign in to comment.