From 32eb462065d9743eda11155b8cb831cebd16934f Mon Sep 17 00:00:00 2001 From: David Montague <35119617+dmontagu@users.noreply.github.com> Date: Mon, 17 Jun 2024 14:50:13 -0600 Subject: [PATCH] Add support for parsing common JSON operators --- datafusion/expr/src/operator.rs | 20 +++++++++++++++- datafusion/expr/src/type_coercion/binary.rs | 24 +++++++++++++++++++ .../physical-expr/src/expressions/binary.rs | 4 ++-- .../proto/src/logical_plan/from_proto.rs | 3 +++ .../tests/cases/roundtrip_logical_plan.rs | 3 +++ datafusion/sql/src/expr/binary_op.rs | 3 +++ datafusion/sql/src/unparser/expr.rs | 7 ++++-- .../substrait/src/logical_plan/consumer.rs | 3 +++ .../substrait/src/logical_plan/producer.rs | 3 +++ 9 files changed, 65 insertions(+), 5 deletions(-) diff --git a/datafusion/expr/src/operator.rs b/datafusion/expr/src/operator.rs index a10312e23446..a38a83db618a 100644 --- a/datafusion/expr/src/operator.rs +++ b/datafusion/expr/src/operator.rs @@ -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 `<@` @@ -129,6 +135,9 @@ impl Operator { | Operator::BitwiseShiftRight | Operator::BitwiseShiftLeft | Operator::StringConcat + | Operator::Question + | Operator::Arrow + | Operator::LongArrow | Operator::AtArrow | Operator::ArrowAt => None, } @@ -213,7 +222,10 @@ impl Operator { | Operator::BitwiseXor | Operator::BitwiseShiftRight | Operator::BitwiseShiftLeft - | Operator::StringConcat => None, + | Operator::StringConcat + | Operator::Question + | Operator::Arrow + | Operator::LongArrow => None, } } @@ -247,6 +259,9 @@ impl Operator { | Operator::BitwiseShiftRight | Operator::BitwiseXor | Operator::StringConcat + | Operator::Question + | Operator::Arrow + | Operator::LongArrow | Operator::AtArrow | Operator::ArrowAt => 0, } @@ -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 => "<@", }; diff --git a/datafusion/expr/src/type_coercion/binary.rs b/datafusion/expr/src/type_coercion/binary.rs index 615bb3ac568c..67c3b1635bd8 100644 --- a/datafusion/expr/src/type_coercion/binary.rs +++ b/datafusion/expr/src/type_coercion/binary.rs @@ -121,6 +121,30 @@ fn signature(lhs: &DataType, op: &Operator, rhs: &DataType) -> Result ) }) } + 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. diff --git a/datafusion/physical-expr/src/expressions/binary.rs b/datafusion/physical-expr/src/expressions/binary.rs index 98df0cba9f3e..621ee9836e85 100644 --- a/datafusion/physical-expr/src/expressions/binary.rs +++ b/datafusion/physical-expr/src/expressions/binary.rs @@ -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") } } } diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index 43cc352f98dd..85d542d38fbb 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -707,6 +707,9 @@ pub fn from_proto_binary_op(op: &str) -> Result { "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!( diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs index 52696a106183..963d99658b9b 100644 --- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs @@ -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); diff --git a/datafusion/sql/src/expr/binary_op.rs b/datafusion/sql/src/expr/binary_op.rs index fcb57e8a82e4..02603b0ab06c 100644 --- a/datafusion/sql/src/expr/binary_op.rs +++ b/datafusion/sql/src/expr/binary_op.rs @@ -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:?}"), diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs index 1d197f3a0d8a..1006f8d06f30 100644 --- a/datafusion/sql/src/unparser/expr.rs +++ b/datafusion/sql/src/unparser/expr.rs @@ -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), } } diff --git a/datafusion/substrait/src/logical_plan/consumer.rs b/datafusion/substrait/src/logical_plan/consumer.rs index 93f197885c0a..bba44d942d48 100644 --- a/datafusion/substrait/src/logical_plan/consumer.rs +++ b/datafusion/substrait/src/logical_plan/consumer.rs @@ -119,6 +119,9 @@ pub fn name_to_op(name: &str) -> Result { "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), diff --git a/datafusion/substrait/src/logical_plan/producer.rs b/datafusion/substrait/src/logical_plan/producer.rs index c0469d333164..8d118a82bc7b 100644 --- a/datafusion/substrait/src/logical_plan/producer.rs +++ b/datafusion/substrait/src/logical_plan/producer.rs @@ -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",