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, DataFusionError, Result, ScalarValue, TableReference};
22use datafusion_expr::expr::{BinaryExpr, ScalarFunction};
23use datafusion_expr::planner::{
24    ExprPlanner, PlannerResult, RawAggregateExpr, RawBinaryExpr, RawFieldAccessExpr, RawScalarExpr,
25    RawWindowExpr,
26};
27use datafusion_expr::type_coercion::functions::{UDFCoercionExt, fields_with_udf};
28use datafusion_expr::{
29    Expr, ExprSchemable, GetFieldAccess, Operator, ScalarUDF, WindowFunctionDefinition,
30};
31use datatypes::extension::json::is_json2_extension_type;
32use sqlparser::ast::BinaryOperator;
33
34/// Rewrites JSON-aware SQL expressions into DataFusion expressions.
35///
36/// This planner handles three cases:
37/// - Rewrites compound identifiers on JSON extension columns into `json_get` function.
38///   For example, `select a.b.c` => `select json_get(a, "b.c")`.
39/// - Extends a JSON path with list indexes and fields following an index.
40///   For example, `select a.b[0].c` => `select json_get(a, "b[0][\"c\"]")`.
41/// - Pushes an "expected type" argument into the `json_get` function when it participates in a
42///   binary operator. So that `json_get` knows the wanted data type when dealing with variant
43///   JSON values.
44///   For example, `select json_get(a, "b.c") + 1` => `select json_get(a, "b.c", NULL::Int64) + 1`.
45/// - Infers the expected type from scalar, aggregate, and window function signatures.
46///   For example, `select abs(a.b.c)` => `select abs(json_get(a, "b.c", NULL::Float64))`.
47#[derive(Debug)]
48pub(crate) struct JsonExprPlanner;
49
50impl ExprPlanner for JsonExprPlanner {
51    fn plan_binary_op(
52        &self,
53        expr: RawBinaryExpr,
54        schema: &DFSchema,
55    ) -> Result<PlannerResult<RawBinaryExpr>> {
56        let RawBinaryExpr {
57            op,
58            mut left,
59            mut right,
60        } = expr;
61
62        if !is_untyped_json_get(&left) && !is_untyped_json_get(&right) {
63            return Ok(PlannerResult::Original(RawBinaryExpr { op, left, right }));
64        }
65
66        let Some(expr_op) = parse_sql_op(&op) else {
67            return Ok(PlannerResult::Original(RawBinaryExpr { op, left, right }));
68        };
69
70        let left_type = left.get_type(schema)?;
71        let right_type = right.get_type(schema)?;
72        let left_changed = push_json_get_type_arg(&mut left, &right_type)?;
73        let right_changed = push_json_get_type_arg(&mut right, &left_type)?;
74        if left_changed || right_changed {
75            Ok(PlannerResult::Planned(Expr::BinaryExpr(BinaryExpr::new(
76                Box::new(left),
77                expr_op,
78                Box::new(right),
79            ))))
80        } else {
81            Ok(PlannerResult::Original(RawBinaryExpr { op, left, right }))
82        }
83    }
84
85    /// Extends the path of an untyped `json_get` with one field access.
86    ///
87    /// For `j.o.l[1].inner.l[2]`, `plan_compound_identifier` first produces
88    /// `json_get(j, "o.l")`. DataFusion then calls this method successively
89    /// with a list index, two named fields, and another list index, producing
90    /// the final path `o.l[1]["inner"]["l"][2]`.
91    fn plan_field_access(
92        &self,
93        mut expr: RawFieldAccessExpr,
94        _schema: &DFSchema,
95    ) -> Result<PlannerResult<RawFieldAccessExpr>> {
96        // See `normalize_field_access_after_subscript` for the reason why we construct the
97        // "suffix" like this.
98        let suffix = match &expr.field_access {
99            GetFieldAccess::ListIndex { key } => {
100                // DataFusion parses ordinary integer literals within the i64 range as Int64.
101                let Expr::Literal(ScalarValue::Int64(Some(index)), _) = key.as_ref() else {
102                    return Ok(PlannerResult::Original(expr));
103                };
104                format!("[{index}]")
105            }
106            GetFieldAccess::NamedStructField { name } => {
107                let Some(name) = name.try_as_str().flatten() else {
108                    return Ok(PlannerResult::Original(expr));
109                };
110                // Encode the field name as a JSON string before embedding it in the
111                // bracket accessor. This preserves dots as literal field-name characters
112                // and escapes quotes, backslashes, and control characters correctly.
113                let name = serde_json::to_string(name)
114                    .map_err(|e| DataFusionError::External(Box::new(e)))?;
115                format!("[{name}]")
116            }
117            GetFieldAccess::ListRange { .. } => return Ok(PlannerResult::Original(expr)),
118        };
119        let Some(json_get) = extract_untyped_json_get(&mut expr.expr) else {
120            return Ok(PlannerResult::Original(expr));
121        };
122        let Some(Expr::Literal(ScalarValue::Utf8(Some(path)), _)) = json_get.args.get_mut(1) else {
123            return Ok(PlannerResult::Original(expr));
124        };
125
126        path.push_str(&suffix);
127        Ok(PlannerResult::Planned(expr.expr))
128    }
129
130    fn plan_compound_identifier(
131        &self,
132        field: &Field,
133        qualifier: Option<&TableReference>,
134        nested_names: &[String],
135    ) -> Result<PlannerResult<Vec<Expr>>> {
136        if !is_json2_extension_type(field) {
137            return Ok(PlannerResult::Original(Vec::new()));
138        }
139
140        static JSON_GET_UDF: LazyLock<Arc<ScalarUDF>> =
141            LazyLock::new(|| Arc::new(create_udf(Arc::new(JsonGetWithType::default()))));
142
143        let json_get = JSON_GET_UDF.clone();
144        let path = nested_names.join(".");
145        Ok(PlannerResult::Planned(Expr::ScalarFunction(
146            ScalarFunction::new_udf(
147                json_get,
148                vec![
149                    Expr::Column(Column::from((qualifier, field))),
150                    Expr::Literal(ScalarValue::Utf8(Some(path)), None),
151                ],
152            ),
153        )))
154    }
155
156    /// Rewrites JSON2 arguments without taking over the final function planning.
157    ///
158    /// `Original` carries the possibly modified raw expression to subsequent planners and then
159    /// DataFusion's default function construction. Returning `Planned` would short-circuit both.
160    fn plan_scalar(&self, mut expr: RawScalarExpr) -> Result<PlannerResult<RawScalarExpr>> {
161        push_function_arg_types(expr.func.as_ref(), &mut expr.args)?;
162        Ok(PlannerResult::Original(expr))
163    }
164
165    /// Rewrites JSON2 arguments while preserving subsequent aggregate planning.
166    fn plan_aggregate(
167        &self,
168        mut expr: RawAggregateExpr,
169    ) -> Result<PlannerResult<RawAggregateExpr>> {
170        push_function_arg_types(expr.func.as_ref(), &mut expr.args)?;
171        Ok(PlannerResult::Original(expr))
172    }
173
174    /// Rewrites JSON2 arguments while preserving subsequent window planning.
175    fn plan_window(&self, mut expr: RawWindowExpr) -> Result<PlannerResult<RawWindowExpr>> {
176        match &expr.func_def {
177            WindowFunctionDefinition::AggregateUDF(func) => {
178                push_function_arg_types(func.as_ref(), &mut expr.args)?;
179            }
180            WindowFunctionDefinition::WindowUDF(func) => {
181                push_function_arg_types(func.as_ref(), &mut expr.args)?;
182            }
183        }
184        Ok(PlannerResult::Original(expr))
185    }
186}
187
188enum JsonGetTypeResolution {
189    Fallback,
190    Typed(Vec<(usize, DataType)>),
191}
192
193/// Infers static output types for untyped `json_get` arguments from a function signature.
194///
195/// DataFusion requires every expression to have one Arrow data type during planning. A JSON path
196/// may contain heterogeneous values across rows, but it cannot expose those values as different
197/// Arrow types in one result column. Preserving their runtime types would require a single
198/// Variant-like data type and Variant-aware functions instead. Maybe we can wait for
199/// https://github.com/apache/datafusion/issues/16116
200///
201/// This helper uses the function's coercion rules to select a supported output type, then appends
202/// a typed NULL argument to each relevant `json_get`. The typed argument makes `json_get` project
203/// compatible JSON values to that type and return NULL for incompatible values. Functions that
204/// accept json_get's default `Utf8View` output keep the two-argument form so later rewrites can
205/// still push down an outer cast.
206fn push_function_arg_types<F>(func: &F, args: &mut [Expr]) -> Result<()>
207where
208    F: UDFCoercionExt,
209{
210    if !args.iter().any(is_untyped_json_get) {
211        return Ok(());
212    }
213
214    let fields = args.iter().map(function_arg_field).collect::<Vec<_>>();
215    match infer_json_get_types(func, args, &fields) {
216        JsonGetTypeResolution::Fallback => {
217            let Some(data_type) = fallback_json_get_type(func, args, &fields) else {
218                return Ok(());
219            };
220            for arg in args.iter_mut() {
221                if is_untyped_json_get(arg) {
222                    let _ = push_json_get_type_arg(arg, &data_type)?;
223                }
224            }
225        }
226        JsonGetTypeResolution::Typed(types) => {
227            for (index, data_type) in types {
228                let _ = push_json_get_type_arg(&mut args[index], &data_type)?;
229            }
230        }
231    }
232    Ok(())
233}
234
235fn infer_json_get_types<F>(func: &F, args: &[Expr], fields: &[Arc<Field>]) -> JsonGetTypeResolution
236where
237    F: UDFCoercionExt,
238{
239    // Only untyped json_get arguments use Null placeholders; preserve every other known argument
240    // type. fields_with_udf performs contextual coercion rather than reverse inference from a
241    // signature alone. Numeric signatures may preserve all-Null inputs, while Comparable
242    // signatures may default them to Utf8. For example, retaining the Float64 peer in
243    // coalesce(json_get(...), 1.0) lets DataFusion resolve json_get to Float64 instead of Utf8.
244    //
245    // This is a best-effort probe: a failure does not mean the actual function call is invalid, so
246    // try concrete JSON types before leaving final validation to DataFusion's default planner.
247    let Ok(coerced) = fields_with_udf(fields, func) else {
248        return JsonGetTypeResolution::Fallback;
249    };
250
251    let mut inferred_types = Vec::with_capacity(coerced.len());
252    for (index, (arg, field)) in args.iter().zip(coerced).enumerate() {
253        if !is_untyped_json_get(arg) || field.data_type().is_null() {
254            continue;
255        }
256        let Some(data_type) = json_get_output_type(field.data_type()) else {
257            return JsonGetTypeResolution::Fallback;
258        };
259        inferred_types.push((index, data_type));
260    }
261    if inferred_types.is_empty() {
262        JsonGetTypeResolution::Fallback
263    } else {
264        JsonGetTypeResolution::Typed(inferred_types)
265    }
266}
267
268fn fallback_json_get_type<F>(func: &F, args: &[Expr], fields: &[Arc<Field>]) -> Option<DataType>
269where
270    F: UDFCoercionExt,
271{
272    // Prefer json_get's default Utf8View type. If the function rejects strings but accepts numeric
273    // values, prefer Float64 so both integers and fractions remain usable.
274    let mut candidate_fields = fields.to_vec();
275    for data_type in [
276        DataType::Utf8View,
277        DataType::Float64,
278        DataType::Int64,
279        DataType::Boolean,
280    ] {
281        for (index, arg) in args.iter().enumerate() {
282            if is_untyped_json_get(arg) {
283                candidate_fields[index] = Arc::new(
284                    fields[index]
285                        .as_ref()
286                        .clone()
287                        .with_data_type(data_type.clone()),
288                );
289            }
290        }
291        if fields_with_udf(&candidate_fields, func).is_ok() {
292            return Some(data_type);
293        }
294    }
295    None
296}
297
298fn function_arg_field(expr: &Expr) -> Arc<Field> {
299    let data_type = if is_untyped_json_get(expr) {
300        DataType::Null
301    } else if let Some(data_type) = extract_json_get_type(expr) {
302        data_type
303    } else {
304        // Treat unresolved expressions as untyped NULL. This lets signatures such as `power`
305        // infer a JSON type, while functions such as `coalesce` can leave it untyped for default
306        // planning. This is only best-effort: overloaded or user-defined functions may select a
307        // different signature for NULL than for the expression's actual type.
308        // TODO(LFC): Use the input schema once DataFusion passes it to ExprPlanner::plan_*().
309        expr.get_type(&DFSchema::empty()).unwrap_or(DataType::Null)
310    };
311    Arc::new(Field::new("", data_type, true))
312}
313
314fn json_get_output_type(data_type: &DataType) -> Option<DataType> {
315    let output_type = match data_type {
316        DataType::Boolean => DataType::Boolean,
317        data_type if data_type.is_integer() => DataType::Int64,
318        data_type if data_type.is_floating() => DataType::Float64,
319        DataType::Decimal128(_, _) | DataType::Decimal256(_, _) => DataType::Float64,
320        data_type if data_type.is_string() => DataType::Utf8View,
321        _ => return None,
322    };
323    Some(output_type)
324}
325
326macro_rules! is_untyped_json_get_func {
327    ($func:expr) => {
328        $func
329            .func
330            .name()
331            .eq_ignore_ascii_case(JsonGetWithType::NAME)
332            && $func.args.len() == 2
333    };
334}
335
336macro_rules! is_typed_json_get_func {
337    ($func:expr) => {
338        $func
339            .func
340            .name()
341            .eq_ignore_ascii_case(JsonGetWithType::NAME)
342            && $func.args.len() == 3
343    };
344}
345
346fn extract_untyped_json_get(expr: &mut Expr) -> Option<&mut ScalarFunction> {
347    match expr {
348        Expr::ScalarFunction(f) if is_untyped_json_get_func!(f) => Some(f),
349        _ => None,
350    }
351}
352
353fn extract_json_get_type(expr: &Expr) -> Option<DataType> {
354    match expr {
355        Expr::ScalarFunction(f) if is_typed_json_get_func!(f) => f
356            .args
357            .get(2)
358            .and_then(|x| x.as_literal())
359            .map(|x| x.data_type()),
360        _ => None,
361    }
362}
363
364fn is_untyped_json_get(expr: &Expr) -> bool {
365    matches!(
366        expr,
367        Expr::ScalarFunction(f) if is_untyped_json_get_func!(f)
368    )
369}
370
371fn push_json_get_type_arg(expr: &mut Expr, data_type: &DataType) -> Result<bool> {
372    let Some(json_get) = extract_untyped_json_get(expr) else {
373        return Ok(false);
374    };
375
376    // The two-argument form already returns Utf8View. Keep it so JsonGetRewriter can still absorb
377    // a cast added by subsequent function coercion.
378    if data_type.is_string() {
379        return Ok(false);
380    }
381    let with_type = ScalarValue::try_new_null(data_type).map(|x| Expr::Literal(x, None))?;
382    json_get.args.push(with_type);
383    Ok(true)
384}
385
386fn parse_sql_op(op: &BinaryOperator) -> Option<Operator> {
387    match *op {
388        BinaryOperator::Plus => Some(Operator::Plus),
389        BinaryOperator::Minus => Some(Operator::Minus),
390        BinaryOperator::Multiply => Some(Operator::Multiply),
391        BinaryOperator::Divide => Some(Operator::Divide),
392        BinaryOperator::Modulo => Some(Operator::Modulo),
393        BinaryOperator::Gt => Some(Operator::Gt),
394        BinaryOperator::GtEq => Some(Operator::GtEq),
395        BinaryOperator::Lt => Some(Operator::Lt),
396        BinaryOperator::LtEq => Some(Operator::LtEq),
397        BinaryOperator::Eq => Some(Operator::Eq),
398        BinaryOperator::NotEq => Some(Operator::NotEq),
399        BinaryOperator::And => Some(Operator::And),
400        BinaryOperator::Or => Some(Operator::Or),
401        BinaryOperator::BitwiseAnd => Some(Operator::BitwiseAnd),
402        BinaryOperator::BitwiseOr => Some(Operator::BitwiseOr),
403        BinaryOperator::BitwiseXor => Some(Operator::BitwiseXor),
404        _ => None,
405    }
406}
407
408#[cfg(test)]
409mod tests {
410    use arrow_schema::Fields;
411    use datafusion::functions_aggregate::count::count_udaf;
412    use datafusion::functions_aggregate::sum::sum_udaf;
413    use datafusion_expr::WindowFrame;
414    use datafusion_functions::core::coalesce;
415    use datafusion_functions::math::{abs, power};
416    use datatypes::extension::json::Json2ExtensionType;
417
418    use super::*;
419
420    fn json_get_expr(base: Expr, path: &str) -> Expr {
421        let json_get = Arc::new(create_udf(Arc::new(JsonGetWithType::default())));
422        Expr::ScalarFunction(ScalarFunction::new_udf(
423            json_get,
424            vec![
425                base,
426                Expr::Literal(ScalarValue::Utf8(Some(path.to_string())), None),
427            ],
428        ))
429    }
430
431    #[test]
432    fn test_plan_binary_op() -> Result<()> {
433        let planner = JsonExprPlanner;
434        let schema = DFSchema::from_unqualified_fields(
435            Fields::from(vec![Field::new("value", DataType::Int64, true)]),
436            Default::default(),
437        )?;
438
439        let planned = planner.plan_binary_op(
440            RawBinaryExpr {
441                op: BinaryOperator::Eq,
442                left: json_get_expr(
443                    Expr::Literal(ScalarValue::Binary(Some(b"{\"a\": 1}".to_vec())), None),
444                    "a",
445                ),
446                right: Expr::Column(Column::new_unqualified("value")),
447            },
448            &schema,
449        )?;
450
451        match planned {
452            PlannerResult::Planned(Expr::BinaryExpr(expr)) => {
453                assert_eq!(expr.op, Operator::Eq);
454
455                match expr.left.as_ref() {
456                    Expr::ScalarFunction(func) => {
457                        assert_eq!(func.func.name(), JsonGetWithType::NAME);
458                        assert_eq!(func.args.len(), 3);
459                        assert_eq!(func.args[2], Expr::Literal(ScalarValue::Int64(None), None));
460                    }
461                    other => panic!("expected json_get on left side, got {other:?}"),
462                }
463
464                assert_eq!(
465                    expr.right.as_ref(),
466                    &Expr::Column(Column::new_unqualified("value"))
467                );
468            }
469            other => panic!("expected planned binary expression, got {other:?}"),
470        }
471
472        let original = planner.plan_binary_op(
473            RawBinaryExpr {
474                op: BinaryOperator::StringConcat,
475                left: Expr::Column(Column::new_unqualified("value")),
476                right: Expr::Literal(ScalarValue::Utf8(Some("x".to_string())), None),
477            },
478            &schema,
479        )?;
480
481        match original {
482            PlannerResult::Original(expr) => {
483                assert!(matches!(expr.op, BinaryOperator::StringConcat));
484                assert_eq!(expr.left, Expr::Column(Column::new_unqualified("value")));
485                assert_eq!(
486                    expr.right,
487                    Expr::Literal(ScalarValue::Utf8(Some("x".to_string())), None)
488                );
489            }
490            other => panic!(
491                "expected original expression for unsupported operator, got {:?}",
492                other,
493            ),
494        }
495
496        Ok(())
497    }
498
499    #[test]
500    fn test_plan_list_index() -> Result<()> {
501        let planner = JsonExprPlanner;
502        let planned = planner.plan_field_access(
503            RawFieldAccessExpr {
504                field_access: GetFieldAccess::ListIndex {
505                    key: Box::new(Expr::Literal(ScalarValue::Int64(Some(0)), None)),
506                },
507                expr: json_get_expr(Expr::Column(Column::new_unqualified("j")), "list"),
508            },
509            &DFSchema::empty(),
510        )?;
511        let PlannerResult::Planned(Expr::ScalarFunction(func)) = planned else {
512            unreachable!()
513        };
514        assert_eq!(func.func.name(), JsonGetWithType::NAME);
515        assert_eq!(func.args.len(), 2);
516        assert_eq!(
517            func.args[1],
518            Expr::Literal(ScalarValue::Utf8(Some("list[0]".to_string())), None)
519        );
520        Ok(())
521    }
522
523    #[test]
524    fn test_plan_field_after_list_index() -> Result<()> {
525        let planner = JsonExprPlanner;
526        let planned = planner.plan_field_access(
527            RawFieldAccessExpr {
528                field_access: GetFieldAccess::NamedStructField {
529                    name: ScalarValue::Utf8(Some("a.b".to_string())),
530                },
531                expr: json_get_expr(Expr::Column(Column::new_unqualified("j")), "list[0]"),
532            },
533            &DFSchema::empty(),
534        )?;
535        let PlannerResult::Planned(Expr::ScalarFunction(func)) = planned else {
536            unreachable!()
537        };
538        assert_eq!(
539            func.args[1],
540            Expr::Literal(
541                ScalarValue::Utf8(Some("list[0][\"a.b\"]".to_string())),
542                None
543            )
544        );
545        Ok(())
546    }
547
548    #[test]
549    fn test_plan_compound_identifier() -> Result<()> {
550        let planner = JsonExprPlanner;
551        let qualifier = TableReference::bare("events");
552        let nested_names = vec!["payload".to_string(), "cpu".to_string()];
553
554        let planned = planner.plan_compound_identifier(
555            &Field::new("labels", DataType::Struct(Fields::empty()), true)
556                .with_extension_type(Json2ExtensionType::default()),
557            Some(&qualifier),
558            &nested_names,
559        )?;
560
561        match planned {
562            PlannerResult::Planned(Expr::ScalarFunction(func)) => {
563                assert_eq!(func.func.name(), JsonGetWithType::NAME);
564                assert_eq!(func.args.len(), 2);
565                assert_eq!(
566                    func.args[0],
567                    Expr::Column(Column::new(Some(qualifier.clone()), "labels"))
568                );
569                assert_eq!(
570                    func.args[1],
571                    Expr::Literal(ScalarValue::Utf8(Some("payload.cpu".to_string())), None)
572                );
573            }
574            other => panic!("expected json_get scalar function, got {other:?}"),
575        }
576
577        let original = planner.plan_compound_identifier(
578            &Field::new("plain", DataType::Utf8, true),
579            Some(&qualifier),
580            &nested_names,
581        )?;
582
583        match original {
584            PlannerResult::Original(exprs) => assert!(exprs.is_empty()),
585            other => panic!(
586                "expected original empty result for non-json field, got {:?}",
587                other,
588            ),
589        }
590
591        Ok(())
592    }
593
594    #[test]
595    fn test_plan_functions() -> Result<()> {
596        let planner = JsonExprPlanner;
597        let json_get = || json_get_expr(Expr::Column(Column::new_unqualified("j")), "a.b");
598
599        let PlannerResult::Original(scalar) = planner.plan_scalar(RawScalarExpr {
600            func: abs(),
601            args: vec![json_get()],
602        })?
603        else {
604            unreachable!();
605        };
606        assert_eq!(
607            Some(DataType::Float64),
608            extract_json_get_type(&scalar.args[0])
609        );
610
611        let PlannerResult::Original(scalar) = planner.plan_scalar(RawScalarExpr {
612            func: power(),
613            args: vec![
614                json_get(),
615                Expr::Column(Column::new_unqualified("exponent")),
616            ],
617        })?
618        else {
619            unreachable!();
620        };
621        assert_eq!(
622            Some(DataType::Float64),
623            extract_json_get_type(&scalar.args[0])
624        );
625
626        let PlannerResult::Original(aggregate) = planner.plan_aggregate(RawAggregateExpr {
627            func: sum_udaf(),
628            args: vec![json_get()],
629            distinct: false,
630            filter: None,
631            order_by: vec![],
632            null_treatment: None,
633        })?
634        else {
635            unreachable!();
636        };
637        assert_eq!(
638            Some(DataType::Float64),
639            extract_json_get_type(&aggregate.args[0])
640        );
641
642        let PlannerResult::Original(count) = planner.plan_aggregate(RawAggregateExpr {
643            func: count_udaf(),
644            args: vec![json_get()],
645            distinct: false,
646            filter: None,
647            order_by: vec![],
648            null_treatment: None,
649        })?
650        else {
651            unreachable!();
652        };
653        assert_eq!(None, extract_json_get_type(&count.args[0]));
654
655        let PlannerResult::Original(window) = planner.plan_window(RawWindowExpr {
656            func_def: WindowFunctionDefinition::AggregateUDF(sum_udaf()),
657            args: vec![json_get()],
658            partition_by: vec![],
659            order_by: vec![],
660            window_frame: WindowFrame::new(None),
661            filter: None,
662            null_treatment: None,
663            distinct: false,
664        })?
665        else {
666            unreachable!();
667        };
668        assert_eq!(
669            Some(DataType::Float64),
670            extract_json_get_type(&window.args[0])
671        );
672        Ok(())
673    }
674
675    #[test]
676    fn test_plan_function_with_mixed_json_get_types() -> Result<()> {
677        let planner = JsonExprPlanner;
678        let json_get = || json_get_expr(Expr::Column(Column::new_unqualified("j")), "a.b");
679        let mut typed = json_get();
680        push_json_get_type_arg(&mut typed, &DataType::Float64)?;
681
682        let PlannerResult::Original(scalar) = planner.plan_scalar(RawScalarExpr {
683            func: coalesce(),
684            args: vec![json_get(), typed],
685        })?
686        else {
687            unreachable!();
688        };
689        assert_eq!(
690            Some(DataType::Float64),
691            extract_json_get_type(&scalar.args[0])
692        );
693        assert_eq!(
694            Some(DataType::Float64),
695            extract_json_get_type(&scalar.args[1])
696        );
697        Ok(())
698    }
699}
query/datafusion/json_expr_planner.rs

query/datafusion/
json_expr_planner.rs
