Add support for parsing common JSON operators

Author: dmontagu

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 => "<@",
         };

datafusion/expr/src/type_coercion/binary.rs

@@ -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.

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")
             }
         }
     }

datafusion/proto/src/logical_plan/from_proto.rs

@@ -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!(

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);

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:?}"),

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),
         }
     }
 

datafusion/substrait/src/logical_plan/consumer.rs

@@ -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),

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",