[branch-53] feat: parse JsonAccess as a binary operator, add Operator::Colon (apache#20628)

Backport of apache#20628 on v53
branch.

- Convert `JsonAccess` from sqlparser as a normal binary operator
- Add `Operator::Colon`
- Needed for datafusion-contrib/datafusion-variant#26

Author: Samyak2

datafusion/expr-common/src/operator.rs

@@ -140,6 +140,10 @@ pub enum Operator {
     ///
     /// Not implemented in DataFusion yet.
     QuestionPipe,
+    /// Colon operator, like `:`
+    ///
+    /// Not implemented in DataFusion yet.
+    Colon,
 }
 
 impl Operator {
@@ -188,7 +192,8 @@ impl Operator {
             | Operator::AtQuestion
             | Operator::Question
             | Operator::QuestionAnd
-            | Operator::QuestionPipe => None,
+            | Operator::QuestionPipe
+            | Operator::Colon => None,
         }
     }
 
@@ -283,7 +288,8 @@ impl Operator {
             | Operator::AtQuestion
             | Operator::Question
             | Operator::QuestionAnd
-            | Operator::QuestionPipe => None,
+            | Operator::QuestionPipe
+            | Operator::Colon => None,
         }
     }
 
@@ -323,7 +329,8 @@ impl Operator {
             | Operator::AtQuestion
             | Operator::Question
             | Operator::QuestionAnd
-            | Operator::QuestionPipe => 30,
+            | Operator::QuestionPipe
+            | Operator::Colon => 30,
             Operator::Plus | Operator::Minus => 40,
             Operator::Multiply | Operator::Divide | Operator::Modulo => 45,
         }
@@ -369,7 +376,8 @@ impl Operator {
             | Operator::AtQuestion
             | Operator::Question
             | Operator::QuestionAnd
-            | Operator::QuestionPipe => true,
+            | Operator::QuestionPipe
+            | Operator::Colon => true,
 
             // E.g. `TRUE OR NULL` is `TRUE`
             Operator::Or
@@ -429,6 +437,7 @@ impl fmt::Display for Operator {
             Operator::Question => "?",
             Operator::QuestionAnd => "?&",
             Operator::QuestionPipe => "?|",
+            Operator::Colon => ":",
         };
         write!(f, "{display}")
     }

datafusion/expr-common/src/type_coercion/binary.rs

@@ -324,6 +324,9 @@ impl<'a> BinaryTypeCoercer<'a> {
                 )
             }
         },
+        Colon => {
+            Ok(Signature { lhs: lhs.clone(), rhs: rhs.clone(), ret: lhs.clone() })
+        },
         IntegerDivide | Arrow | LongArrow | HashArrow | HashLongArrow
         | HashMinus | AtQuestion | Question | QuestionAnd | QuestionPipe => {
             not_impl_err!("Operator {} is not yet supported", self.op)

datafusion/physical-expr/src/expressions/binary.rs

@@ -715,7 +715,7 @@ impl BinaryExpr {
             StringConcat => concat_elements(&left, &right),
             AtArrow | ArrowAt | Arrow | LongArrow | HashArrow | HashLongArrow | AtAt
             | HashMinus | AtQuestion | Question | QuestionAnd | QuestionPipe
-            | IntegerDivide => {
+            | IntegerDivide | Colon => {
                 not_impl_err!(
                     "Binary operator '{:?}' is not supported in the physical expr",
                     self.op

datafusion/sql/src/expr/binary_op.rs

@@ -22,7 +22,7 @@ use sqlparser::ast::BinaryOperator;
 
 impl<S: ContextProvider> SqlToRel<'_, S> {
     pub(crate) fn parse_sql_binary_op(&self, op: &BinaryOperator) -> Result<Operator> {
-        match *op {
+        match op {
             BinaryOperator::Gt => Ok(Operator::Gt),
             BinaryOperator::GtEq => Ok(Operator::GtEq),
             BinaryOperator::Lt => Ok(Operator::Lt),
@@ -68,6 +68,7 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
             BinaryOperator::Question => Ok(Operator::Question),
             BinaryOperator::QuestionAnd => Ok(Operator::QuestionAnd),
             BinaryOperator::QuestionPipe => Ok(Operator::QuestionPipe),
+            BinaryOperator::Custom(s) if s == ":" => Ok(Operator::Colon),
             _ => not_impl_err!("Unsupported binary operator: {:?}", op),
         }
     }

datafusion/sql/src/expr/mod.rs

@@ -22,8 +22,8 @@ use datafusion_expr::planner::{
 use sqlparser::ast::{
     AccessExpr, BinaryOperator, CastFormat, CastKind, CeilFloorKind,
     DataType as SQLDataType, DateTimeField, DictionaryField, Expr as SQLExpr,
-    ExprWithAlias as SQLExprWithAlias, MapEntry, StructField, Subscript, TrimWhereField,
-    TypedString, Value, ValueWithSpan,
+    ExprWithAlias as SQLExprWithAlias, JsonPath, MapEntry, StructField, Subscript,
+    TrimWhereField, TypedString, Value, ValueWithSpan,
 };
 
 use datafusion_common::{
@@ -651,10 +651,36 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                 options: Box::new(WildcardOptions::default()),
             }),
             SQLExpr::Tuple(values) => self.parse_tuple(schema, planner_context, values),
+            SQLExpr::JsonAccess { value, path } => {
+                self.parse_json_access(schema, planner_context, value, &path)
+            }
             _ => not_impl_err!("Unsupported ast node in sqltorel: {sql:?}"),
         }
     }
 
+    fn parse_json_access(
+        &self,
+        schema: &DFSchema,
+        planner_context: &mut PlannerContext,
+        value: Box<SQLExpr>,
+        path: &JsonPath,
+    ) -> Result<Expr> {
+        let json_path = path.to_string();
+        let json_path = if let Some(json_path) = json_path.strip_prefix(":") {
+            // sqlparser's JsonPath display adds an extra `:` at the beginning.
+            json_path.to_owned()
+        } else {
+            json_path
+        };
+        self.build_logical_expr(
+            BinaryOperator::Custom(":".to_owned()),
+            self.sql_to_expr(*value, schema, planner_context)?,
+            // pass json path as a string literal, let the impl parse it when needed.
+            Expr::Literal(ScalarValue::Utf8(Some(json_path)), None),
+            schema,
+        )
+    }
+
     /// Parses a struct(..) expression and plans it creation
     fn parse_struct(
         &self,

datafusion/sql/src/unparser/expr.rs

@@ -1094,6 +1094,7 @@ impl Unparser<'_> {
             Operator::Question => Ok(BinaryOperator::Question),
             Operator::QuestionAnd => Ok(BinaryOperator::QuestionAnd),
             Operator::QuestionPipe => Ok(BinaryOperator::QuestionPipe),
+            Operator::Colon => Ok(BinaryOperator::Custom(":".to_owned())),
         }
     }
 

datafusion/sql/tests/cases/plan_to_sql.rs

@@ -2821,3 +2821,39 @@ fn test_struct_expr3() {
         @r#"SELECT test.c1."metadata".product."name" FROM (SELECT {"metadata": {product: {"name": 'Product Name'}}} AS c1) AS test"#
     );
 }
+
+#[test]
+fn test_json_access_1() {
+    let statement = generate_round_trip_statement(
+        GenericDialect {},
+        r#"SELECT j1_string:field FROM j1"#,
+    );
+    assert_snapshot!(
+        statement,
+        @r#"SELECT (j1.j1_string : 'field') FROM j1"#
+    );
+}
+
+#[test]
+fn test_json_access_2() {
+    let statement = generate_round_trip_statement(
+        GenericDialect {},
+        r#"SELECT j1_string:field[0] FROM j1"#,
+    );
+    assert_snapshot!(
+        statement,
+        @r#"SELECT (j1.j1_string : 'field[0]') FROM j1"#
+    );
+}
+
+#[test]
+fn test_json_access_3() {
+    let statement = generate_round_trip_statement(
+        GenericDialect {},
+        r#"SELECT j1_string:field.inner1['inner2'] FROM j1"#,
+    );
+    assert_snapshot!(
+        statement,
+        @r#"SELECT (j1.j1_string : 'field.inner1[''inner2'']') FROM j1"#
+    );
+}

datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs

@@ -344,5 +344,6 @@ pub fn operator_to_name(op: Operator) -> &'static str {
         Operator::BitwiseXor => "bitwise_xor",
         Operator::BitwiseShiftRight => "bitwise_shift_right",
         Operator::BitwiseShiftLeft => "bitwise_shift_left",
+        Operator::Colon => "colon",
     }
 }