diff --git a/datafusion/expr-common/src/operator.rs b/datafusion/expr-common/src/operator.rs index 33512b0c354d6..427069b326f9d 100644 --- a/datafusion/expr-common/src/operator.rs +++ b/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}") } diff --git a/datafusion/expr-common/src/type_coercion/binary.rs b/datafusion/expr-common/src/type_coercion/binary.rs index c6ac86cd396c4..e696545ea6caf 100644 --- a/datafusion/expr-common/src/type_coercion/binary.rs +++ b/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) diff --git a/datafusion/physical-expr/src/expressions/binary.rs b/datafusion/physical-expr/src/expressions/binary.rs index 72eae396e68a6..02628b405ec6c 100644 --- a/datafusion/physical-expr/src/expressions/binary.rs +++ b/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 diff --git a/datafusion/sql/src/expr/binary_op.rs b/datafusion/sql/src/expr/binary_op.rs index edad5bbc6daad..4e9025e02e0c7 100644 --- a/datafusion/sql/src/expr/binary_op.rs +++ b/datafusion/sql/src/expr/binary_op.rs @@ -22,7 +22,7 @@ use sqlparser::ast::BinaryOperator; impl SqlToRel<'_, S> { pub(crate) fn parse_sql_binary_op(&self, op: &BinaryOperator) -> Result { - match *op { + match op { BinaryOperator::Gt => Ok(Operator::Gt), BinaryOperator::GtEq => Ok(Operator::GtEq), BinaryOperator::Lt => Ok(Operator::Lt), @@ -68,6 +68,7 @@ impl 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), } } diff --git a/datafusion/sql/src/expr/mod.rs b/datafusion/sql/src/expr/mod.rs index 9aa5be8131dcb..7902eed1e6922 100644 --- a/datafusion/sql/src/expr/mod.rs +++ b/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 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, + path: &JsonPath, + ) -> Result { + 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, diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs index 59a9207b51ef0..b82ab24adef71 100644 --- a/datafusion/sql/src/unparser/expr.rs +++ b/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())), } } diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs index 4717b843abb53..670046f164ed3 100644 --- a/datafusion/sql/tests/cases/plan_to_sql.rs +++ b/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"# + ); +} diff --git a/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs b/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs index bd8a9d9a99b53..9f70e903a0bd9 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs +++ b/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", } }