query/datafusion/
json_expr_planner.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::{Arc, LazyLock};
16
17use arrow_schema::Field;
18use common_function::scalars::json::json_get::JsonGetWithType;
19use common_function::scalars::udf::create_udf;
20use datafusion_common::arrow::datatypes::DataType;
21use datafusion_common::{Column, DFSchema, Result, ScalarValue, TableReference};
22use datafusion_expr::expr::{BinaryExpr, ScalarFunction};
23use datafusion_expr::planner::{ExprPlanner, PlannerResult, RawBinaryExpr};
24use datafusion_expr::{Expr, ExprSchemable, Operator, ScalarUDF};
25use datatypes::extension::json::is_json2_extension_type;
26use either::Either;
27use sqlparser::ast::BinaryOperator;
28
29/// Rewrites JSON-aware SQL expressions into DataFusion expressions.
30///
31/// This planner handles two cases:
32/// - Rewrites compound identifiers on JSON extension columns into `json_get` function.
33///   For example, `select a.b.c` => `select json_get(a, "b.c")`.
34/// - Pushes an "expected type" argument into the `json_get` function when it participates in a
35///   binary operator. So that `json_get` knows the wanted data type when dealing with variant
36///   JSON values.
37///   For example, `select json_get(a, "b.c") + 1` => `select json_get(a, "b.c", NULL::Int64) + 1`.
38#[derive(Debug)]
39pub(crate) struct JsonExprPlanner;
40
41impl ExprPlanner for JsonExprPlanner {
42    fn plan_binary_op(
43        &self,
44        expr: RawBinaryExpr,
45        schema: &DFSchema,
46    ) -> Result<PlannerResult<RawBinaryExpr>> {
47        let RawBinaryExpr {
48            op,
49            mut left,
50            mut right,
51        } = expr;
52
53        if extract_untyped_json_get(&mut left).is_none()
54            && extract_untyped_json_get(&mut right).is_none()
55        {
56            return Ok(PlannerResult::Original(RawBinaryExpr { op, left, right }));
57        }
58
59        let Some(expr_op) = parse_sql_op(&op) else {
60            return Ok(PlannerResult::Original(RawBinaryExpr { op, left, right }));
61        };
62
63        let left_type = left.get_type(schema)?;
64        let right_type = right.get_type(schema)?;
65        let left = push_json_get_type_arg(left, right_type)?;
66        let right = push_json_get_type_arg(right, left_type)?;
67        match (left, right) {
68            (Either::Left(left), Either::Left(right)) => {
69                Ok(PlannerResult::Original(RawBinaryExpr { op, left, right }))
70            }
71            (left, right) => Ok(PlannerResult::Planned(Expr::BinaryExpr(BinaryExpr::new(
72                Box::new(left.into_inner()),
73                expr_op,
74                Box::new(right.into_inner()),
75            )))),
76        }
77    }
78
79    fn plan_compound_identifier(
80        &self,
81        field: &Field,
82        qualifier: Option<&TableReference>,
83        nested_names: &[String],
84    ) -> Result<PlannerResult<Vec<Expr>>> {
85        if !is_json2_extension_type(field) {
86            return Ok(PlannerResult::Original(Vec::new()));
87        }
88
89        static JSON_GET_UDF: LazyLock<Arc<ScalarUDF>> =
90            LazyLock::new(|| Arc::new(create_udf(Arc::new(JsonGetWithType::default()))));
91
92        let json_get = JSON_GET_UDF.clone();
93        let path = nested_names.join(".");
94        Ok(PlannerResult::Planned(Expr::ScalarFunction(
95            ScalarFunction::new_udf(
96                json_get,
97                vec![
98                    Expr::Column(Column::from((qualifier, field))),
99                    Expr::Literal(ScalarValue::Utf8(Some(path)), None),
100                ],
101            ),
102        )))
103    }
104}
105
106fn extract_untyped_json_get(expr: &mut Expr) -> Option<&mut ScalarFunction> {
107    match expr {
108        Expr::ScalarFunction(f)
109            if f.func.name().eq_ignore_ascii_case(JsonGetWithType::NAME) && f.args.len() == 2 =>
110        {
111            Some(f)
112        }
113        _ => None,
114    }
115}
116
117fn push_json_get_type_arg(mut expr: Expr, mut data_type: DataType) -> Result<Either<Expr, Expr>> {
118    let Some(json_get) = extract_untyped_json_get(&mut expr) else {
119        return Ok(Either::Left(expr));
120    };
121
122    if data_type.is_string() {
123        data_type = DataType::Utf8View;
124    }
125    let with_type = ScalarValue::try_new_null(&data_type).map(|x| Expr::Literal(x, None))?;
126    json_get.args.push(with_type);
127
128    Ok(Either::Right(expr))
129}
130
131fn parse_sql_op(op: &BinaryOperator) -> Option<Operator> {
132    match *op {
133        BinaryOperator::Plus => Some(Operator::Plus),
134        BinaryOperator::Minus => Some(Operator::Minus),
135        BinaryOperator::Multiply => Some(Operator::Multiply),
136        BinaryOperator::Divide => Some(Operator::Divide),
137        BinaryOperator::Modulo => Some(Operator::Modulo),
138        BinaryOperator::Gt => Some(Operator::Gt),
139        BinaryOperator::GtEq => Some(Operator::GtEq),
140        BinaryOperator::Lt => Some(Operator::Lt),
141        BinaryOperator::LtEq => Some(Operator::LtEq),
142        BinaryOperator::Eq => Some(Operator::Eq),
143        BinaryOperator::NotEq => Some(Operator::NotEq),
144        BinaryOperator::And => Some(Operator::And),
145        BinaryOperator::Or => Some(Operator::Or),
146        BinaryOperator::BitwiseAnd => Some(Operator::BitwiseAnd),
147        BinaryOperator::BitwiseOr => Some(Operator::BitwiseOr),
148        BinaryOperator::BitwiseXor => Some(Operator::BitwiseXor),
149        _ => None,
150    }
151}
152
153#[cfg(test)]
154mod tests {
155    use arrow_schema::Fields;
156    use datatypes::extension::json::Json2ExtensionType;
157
158    use super::*;
159
160    fn json_get_expr(base: Expr, path: &str) -> Expr {
161        let json_get = Arc::new(create_udf(Arc::new(JsonGetWithType::default())));
162        Expr::ScalarFunction(ScalarFunction::new_udf(
163            json_get,
164            vec![
165                base,
166                Expr::Literal(ScalarValue::Utf8(Some(path.to_string())), None),
167            ],
168        ))
169    }
170
171    #[test]
172    fn test_plan_binary_op() -> Result<()> {
173        let planner = JsonExprPlanner;
174        let schema = DFSchema::from_unqualified_fields(
175            Fields::from(vec![Field::new("value", DataType::Int64, true)]),
176            Default::default(),
177        )?;
178
179        let planned = planner.plan_binary_op(
180            RawBinaryExpr {
181                op: BinaryOperator::Eq,
182                left: json_get_expr(
183                    Expr::Literal(ScalarValue::Binary(Some(b"{\"a\": 1}".to_vec())), None),
184                    "a",
185                ),
186                right: Expr::Column(Column::new_unqualified("value")),
187            },
188            &schema,
189        )?;
190
191        match planned {
192            PlannerResult::Planned(Expr::BinaryExpr(expr)) => {
193                assert_eq!(expr.op, Operator::Eq);
194
195                match expr.left.as_ref() {
196                    Expr::ScalarFunction(func) => {
197                        assert_eq!(func.func.name(), JsonGetWithType::NAME);
198                        assert_eq!(func.args.len(), 3);
199                        assert_eq!(func.args[2], Expr::Literal(ScalarValue::Int64(None), None));
200                    }
201                    other => panic!("expected json_get on left side, got {other:?}"),
202                }
203
204                assert_eq!(
205                    expr.right.as_ref(),
206                    &Expr::Column(Column::new_unqualified("value"))
207                );
208            }
209            other => panic!("expected planned binary expression, got {other:?}"),
210        }
211
212        let original = planner.plan_binary_op(
213            RawBinaryExpr {
214                op: BinaryOperator::StringConcat,
215                left: Expr::Column(Column::new_unqualified("value")),
216                right: Expr::Literal(ScalarValue::Utf8(Some("x".to_string())), None),
217            },
218            &schema,
219        )?;
220
221        match original {
222            PlannerResult::Original(expr) => {
223                assert!(matches!(expr.op, BinaryOperator::StringConcat));
224                assert_eq!(expr.left, Expr::Column(Column::new_unqualified("value")));
225                assert_eq!(
226                    expr.right,
227                    Expr::Literal(ScalarValue::Utf8(Some("x".to_string())), None)
228                );
229            }
230            other => panic!(
231                "expected original expression for unsupported operator, got {:?}",
232                other,
233            ),
234        }
235
236        Ok(())
237    }
238
239    #[test]
240    fn test_plan_compound_identifier() -> Result<()> {
241        let planner = JsonExprPlanner;
242        let qualifier = TableReference::bare("events");
243        let nested_names = vec!["payload".to_string(), "cpu".to_string()];
244
245        let planned = planner.plan_compound_identifier(
246            &Field::new("labels", DataType::Struct(Fields::empty()), true)
247                .with_extension_type(Json2ExtensionType::default()),
248            Some(&qualifier),
249            &nested_names,
250        )?;
251
252        match planned {
253            PlannerResult::Planned(Expr::ScalarFunction(func)) => {
254                assert_eq!(func.func.name(), JsonGetWithType::NAME);
255                assert_eq!(func.args.len(), 2);
256                assert_eq!(
257                    func.args[0],
258                    Expr::Column(Column::new(Some(qualifier.clone()), "labels"))
259                );
260                assert_eq!(
261                    func.args[1],
262                    Expr::Literal(ScalarValue::Utf8(Some("payload.cpu".to_string())), None)
263                );
264            }
265            other => panic!("expected json_get scalar function, got {other:?}"),
266        }
267
268        let original = planner.plan_compound_identifier(
269            &Field::new("plain", DataType::Utf8, true),
270            Some(&qualifier),
271            &nested_names,
272        )?;
273
274        match original {
275            PlannerResult::Original(exprs) => assert!(exprs.is_empty()),
276            other => panic!(
277                "expected original empty result for non-json field, got {:?}",
278                other,
279            ),
280        }
281
282        Ok(())
283    }
284}
query/datafusion/json_expr_planner.rs

query/datafusion/
json_expr_planner.rs
