query/optimizer/
json_type_concretize.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::collections::HashMap;
16
17use arrow_schema::DataType;
18use common_function::scalars::json::json_get::JsonGetWithType;
19use datafusion::datasource::{DefaultTableSource, TableProvider};
20use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
21use datafusion_common::{Result, plan_datafusion_err, plan_err};
22use datafusion_expr::{Expr, LogicalPlan};
23use datafusion_optimizer::{OptimizerConfig, OptimizerRule};
24use datatypes::extension::json::is_json2_extension_type;
25use datatypes::types::json_type::{JsonNativeType, JsonObjectType};
26use table::table::adapter::DfTableProviderAdapter;
27
28use crate::dummy_catalog::DummyTableProvider;
29
30/// Concretize (deduce) the expected JSON type from query.
31/// For example, we can concretize a JSON type of `{ a: { b: Number } }` from `select j.a.b::Int64`.
32/// The JSON type will be later set into the scan request, for converting the JSON arrays.
33#[derive(Debug)]
34pub(crate) struct JsonTypeConcretizeRule;
35
36impl OptimizerRule for JsonTypeConcretizeRule {
37    fn name(&self) -> &str {
38        "JsonTypeConcretizeRule"
39    }
40
41    fn rewrite(
42        &self,
43        plan: LogicalPlan,
44        _config: &dyn OptimizerConfig,
45    ) -> Result<Transformed<LogicalPlan>> {
46        let json_types = deduce_json_types(&plan)?;
47        if json_types.is_empty() {
48            return Ok(Transformed::no(plan));
49        }
50
51        plan.transform_down(|plan| match &plan {
52            LogicalPlan::TableScan(table_scan) => {
53                let Some(source) = table_scan
54                    .source
55                    .as_any()
56                    .downcast_ref::<DefaultTableSource>()
57                else {
58                    return Ok(Transformed::no(plan));
59                };
60
61                if apply_json_type_hint(source.table_provider.as_ref(), &json_types) {
62                    Ok(Transformed::yes(plan))
63                } else {
64                    Ok(Transformed::no(plan))
65                }
66            }
67            _ => Ok(Transformed::no(plan)),
68        })
69    }
70}
71
72// FIXME: `json_types` is keyed only by unqualified column name. In joins with
73// same-named JSON2 columns, a hint deduced from one scan can be applied to
74// another scan. Carry the originating relation/scan when deducing hints.
75/// Applies JSON type hints to providers that can carry scan request hints.
76///
77/// Returns `true` if at least one JSON2 hint is retained and written to the provider.
78fn apply_json_type_hint(
79    provider: &dyn TableProvider,
80    json_types: &HashMap<String, JsonNativeType>,
81) -> bool {
82    let schema = provider.schema();
83    let json_types = json_types
84        .iter()
85        .filter(|(column, _)| {
86            schema
87                .fields()
88                .iter()
89                .any(|field| field.name() == *column && is_json2_extension_type(field))
90        })
91        .map(|(column, json_type)| (column.clone(), json_type.clone()))
92        .collect::<HashMap<_, _>>();
93
94    if json_types.is_empty() {
95        return false;
96    }
97
98    if let Some(adapter) = provider.as_any().downcast_ref::<DummyTableProvider>() {
99        adapter.with_json_type_hint(json_types);
100        return true;
101    }
102
103    if let Some(adapter) = provider.as_any().downcast_ref::<DfTableProviderAdapter>() {
104        adapter.with_json_type_hint(json_types);
105        return true;
106    }
107
108    false
109}
110
111pub(crate) fn deduce_json_types(plan: &LogicalPlan) -> Result<HashMap<String, JsonNativeType>> {
112    let mut json_types = HashMap::<String, JsonNativeType>::new();
113
114    // JSON2 columns in the final output must retain their complete values even when
115    // predicates or other expressions access only specific paths.
116    // For example, `SELECT j FROM t WHERE json_get(j, 'a') = 1`.
117    plan.schema()
118        .fields()
119        .iter()
120        .filter(|field| is_json2_extension_type(field))
121        .for_each(|field| {
122            json_types.insert(field.name().clone(), JsonNativeType::Variant);
123        });
124
125    plan.apply(|plan| {
126        for expr in plan.expressions() {
127            // Optimizer-generated projections may keep the JSON root only so later json_get
128            // expressions can access another path. A same-name pass-through does not require the
129            // complete root by itself; any real whole-column consumer above it is visited
130            // separately, and a whole root in the final output is captured from the plan schema.
131            if matches!(plan, LogicalPlan::Projection(_)) && is_same_name_column_projection(&expr) {
132                continue;
133            }
134            expr.apply(|expr| {
135                if let Some((column, json_type)) = deduce_json_type(expr)? {
136                    json_types.entry(column).or_default().merge(&json_type);
137                    Ok(TreeNodeRecursion::Jump)
138                } else {
139                    Ok(TreeNodeRecursion::Continue)
140                }
141            })?;
142        }
143        Ok(TreeNodeRecursion::Continue)
144    })?;
145    Ok(json_types)
146}
147
148fn is_same_name_column_projection(expr: &Expr) -> bool {
149    match expr {
150        Expr::Column(_) => true,
151        Expr::Alias(alias) => {
152            matches!(alias.expr.as_ref(), Expr::Column(column) if column.name == alias.name)
153        }
154        _ => false,
155    }
156}
157
158fn deduce_json_type(expr: &Expr) -> Result<Option<(String, JsonNativeType)>> {
159    let f = match expr {
160        Expr::ScalarFunction(f) if f.name().eq_ignore_ascii_case(JsonGetWithType::NAME) => f,
161        Expr::Column(c) => return Ok(Some((c.name.clone(), JsonNativeType::Variant))),
162        _ => return Ok(None),
163    };
164
165    let Some(Expr::Column(column)) = f.args.first() else {
166        return plan_err!(
167            "First argument of {} is expected to be a column expr, actual: {:?}",
168            JsonGetWithType::NAME,
169            f.args.first()
170        );
171    };
172
173    let Some(path) = f
174        .args
175        .get(1)
176        .and_then(|expr| expr.as_literal())
177        .and_then(|x| x.try_as_str())
178        .flatten()
179    else {
180        return plan_err!(
181            "Second argument of {} is expected to be a string literal, actual: {:?}",
182            JsonGetWithType::NAME,
183            f.args.get(1)
184        );
185    };
186
187    // Object-only type deduction cannot represent bracket JSONPath access, so preserve the
188    // full Variant and let json_get apply the expression.
189    if path.contains('[') {
190        return Ok(Some((column.name.clone(), JsonNativeType::Variant)));
191    }
192
193    let with_type = f
194        .args
195        .get(2)
196        .and_then(|expr| expr.as_literal())
197        .map(|x| x.data_type())
198        .unwrap_or(DataType::Utf8View);
199    let with_type =
200        JsonNativeType::try_from(&with_type).map_err(|e| plan_datafusion_err!("{e:?}"))?;
201
202    let mut split = path.rsplit(".");
203    let Some(leaf) = split.next().filter(|&x| !x.is_empty() && x != "$") else {
204        return Ok(Some((column.name.clone(), JsonNativeType::String)));
205    };
206
207    let mut object = JsonObjectType::new();
208    object.insert(leaf.to_string(), with_type);
209    let mut root = JsonNativeType::Object(object);
210
211    for s in split {
212        let mut object = JsonObjectType::new();
213        object.insert(s.to_string(), root);
214        root = JsonNativeType::Object(object);
215    }
216
217    Ok(Some((column.name.clone(), root)))
218}
219
220#[cfg(test)]
221mod tests {
222    use std::sync::Arc;
223
224    use api::v1::SemanticType;
225    use common_function::scalars::udf::create_udf;
226    use datafusion::datasource::provider_as_source;
227    use datafusion::functions_aggregate::expr_fn::count;
228    use datafusion_common::{Column, ScalarValue};
229    use datafusion_expr::expr::ScalarFunction;
230    use datafusion_expr::{LogicalPlanBuilder, col, lit};
231    use datafusion_optimizer::OptimizerContext;
232    use datatypes::extension::json::Json2ExtensionType;
233    use datatypes::schema::ColumnSchema;
234    use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
235    use store_api::storage::{ConcreteDataType, RegionId};
236
237    use super::*;
238    use crate::optimizer::test_util::{MetaRegionEngine, mock_table_provider};
239
240    fn json_get_expr(base: Expr, path: Expr, with_type: Option<DataType>) -> Result<Expr> {
241        let json_get = Arc::new(create_udf(Arc::new(JsonGetWithType::default())));
242        let mut args = vec![base, path];
243        if let Some(with_type) = with_type {
244            let with_type = ScalarValue::try_new_null(&with_type)?;
245            args.push(Expr::Literal(with_type, None));
246        }
247        Ok(Expr::ScalarFunction(ScalarFunction::new_udf(
248            json_get, args,
249        )))
250    }
251
252    fn path_expr(path: &str) -> Expr {
253        Expr::Literal(ScalarValue::Utf8(Some(path.to_string())), None)
254    }
255
256    fn build_plan(exprs: Vec<Expr>) -> Result<(Arc<DummyTableProvider>, LogicalPlan)> {
257        let provider = Arc::new(mock_table_provider(RegionId::new(1024, 1)));
258        let plan = LogicalPlanBuilder::scan("t", provider_as_source(provider.clone()), None)?
259            .project(exprs)?
260            .build()?;
261        Ok((provider, plan))
262    }
263
264    fn build_json2_scan() -> Result<(Arc<DummyTableProvider>, LogicalPlanBuilder)> {
265        let region_id = RegionId::new(1024, 2);
266        let mut builder = RegionMetadataBuilder::new(region_id);
267        let mut json_column = ColumnSchema::new(
268            "j",
269            ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())),
270            true,
271        );
272        json_column.with_extension_type(&Json2ExtensionType::default());
273        builder
274            .push_column_metadata(ColumnMetadata {
275                column_schema: json_column,
276                semantic_type: SemanticType::Field,
277                column_id: 1,
278            })
279            .push_column_metadata(ColumnMetadata {
280                column_schema: ColumnSchema::new(
281                    "ts",
282                    ConcreteDataType::timestamp_millisecond_datatype(),
283                    false,
284                ),
285                semantic_type: SemanticType::Timestamp,
286                column_id: 2,
287            });
288        let metadata = Arc::new(builder.build().unwrap());
289        let engine = Arc::new(MetaRegionEngine::with_metadata(metadata.clone()));
290        let provider = Arc::new(DummyTableProvider::new(region_id, engine, metadata));
291        let plan = LogicalPlanBuilder::scan("t", provider_as_source(provider.clone()), None)?;
292        Ok((provider, plan))
293    }
294
295    fn build_json2_plan(exprs: Vec<Expr>) -> Result<(Arc<DummyTableProvider>, LogicalPlan)> {
296        let (provider, plan) = build_json2_scan()?;
297        let plan = plan.project(exprs)?.build()?;
298        Ok((provider, plan))
299    }
300
301    #[test]
302    fn test_json_type_concretize_rule_rewrite() -> Result<()> {
303        let exprs = vec![
304            json_get_expr(col("j"), path_expr("a.b"), Some(DataType::Int64))?.alias("ab"),
305            json_get_expr(col("j"), path_expr("a.c"), None)?.alias("ac"),
306            json_get_expr(col("j"), path_expr("d"), Some(DataType::Boolean))?.alias("d"),
307        ];
308        let (provider, plan) = build_json2_plan(exprs)?;
309
310        assert!(
311            JsonTypeConcretizeRule
312                .rewrite(plan, &OptimizerContext::default())?
313                .transformed
314        );
315
316        let expected = JsonNativeType::Object(JsonObjectType::from([
317            (
318                "a".to_string(),
319                JsonNativeType::Object(JsonObjectType::from([
320                    ("b".to_string(), JsonNativeType::i64()),
321                    ("c".to_string(), JsonNativeType::String),
322                ])),
323            ),
324            ("d".to_string(), JsonNativeType::Bool),
325        ]));
326
327        let request = provider.scan_request();
328        assert_eq!(1, request.json_type_hint.len());
329        assert_eq!(Some(&expected), request.json_type_hint.get("j"));
330        Ok(())
331    }
332
333    #[test]
334    fn test_deduce_json_type_with_list_index() -> Result<()> {
335        let expr = json_get_expr(col("j"), path_expr("l[0]"), Some(DataType::Int64))?;
336
337        assert_eq!(
338            Some(("j".to_string(), JsonNativeType::Variant)),
339            deduce_json_type(&expr)?
340        );
341        Ok(())
342    }
343
344    #[test]
345    fn test_json_type_concretize_rule_conflict_to_variant() -> Result<()> {
346        let exprs = vec![
347            json_get_expr(col("j"), path_expr("a"), Some(DataType::Int64))?.alias("a_num"),
348            json_get_expr(col("j"), path_expr("a.b"), Some(DataType::Boolean))?.alias("a_obj"),
349        ];
350        let (provider, plan) = build_json2_plan(exprs)?;
351
352        assert!(
353            JsonTypeConcretizeRule
354                .rewrite(plan, &OptimizerContext::default())?
355                .transformed
356        );
357
358        let expected = JsonNativeType::Object(JsonObjectType::from([(
359            "a".to_string(),
360            JsonNativeType::Variant,
361        )]));
362        assert_eq!(
363            Some(&expected),
364            provider.scan_request().json_type_hint.get("j")
365        );
366        Ok(())
367    }
368
369    #[test]
370    fn test_json_type_concretize_rule_ignores_non_json2_columns() -> Result<()> {
371        let exprs =
372            vec![json_get_expr(col("k0"), path_expr("a.b"), Some(DataType::Int64))?.alias("ab")];
373        let (provider, plan) = build_plan(exprs)?;
374
375        assert!(
376            !JsonTypeConcretizeRule
377                .rewrite(plan, &OptimizerContext::default())?
378                .transformed
379        );
380        assert!(provider.scan_request().json_type_hint.is_empty());
381        Ok(())
382    }
383
384    #[test]
385    fn test_json_type_concretize_rule_no_json_get() -> Result<()> {
386        let (provider, plan) = build_plan(vec![col("k0"), col("v0")])?;
387
388        assert!(
389            !JsonTypeConcretizeRule
390                .rewrite(plan, &OptimizerContext::default())?
391                .transformed
392        );
393        assert!(provider.scan_request().json_type_hint.is_empty());
394        Ok(())
395    }
396
397    #[test]
398    fn test_allow_json2_path_use_in_intermediate_plan() -> Result<()> {
399        let json_get = json_get_expr(col("j"), path_expr("a"), Some(DataType::Int64))?;
400        let (provider, plan) = build_json2_scan()?;
401        let plan = plan
402            .aggregate(vec![json_get], Vec::<Expr>::new())?
403            .aggregate(Vec::<Expr>::new(), vec![count(lit(1))])?
404            .build()?;
405
406        assert!(
407            JsonTypeConcretizeRule
408                .rewrite(plan, &OptimizerContext::default())?
409                .transformed
410        );
411        assert_eq!(
412            Some(&JsonNativeType::Object(JsonObjectType::from([(
413                "a".to_string(),
414                JsonNativeType::i64(),
415            )]))),
416            provider.scan_request().json_type_hint.get("j")
417        );
418        Ok(())
419    }
420
421    #[test]
422    fn test_allow_json2_projection_by_path() -> Result<()> {
423        let expr = json_get_expr(col("j"), path_expr("a"), Some(DataType::Int64))?;
424        let (provider, plan) = build_json2_plan(vec![expr])?;
425
426        assert!(
427            JsonTypeConcretizeRule
428                .rewrite(plan, &OptimizerContext::default())?
429                .transformed
430        );
431        assert_eq!(
432            Some(&JsonNativeType::Object(JsonObjectType::from([(
433                "a".to_string(),
434                JsonNativeType::i64(),
435            )]))),
436            provider.scan_request().json_type_hint.get("j")
437        );
438        Ok(())
439    }
440
441    #[test]
442    fn test_allow_json2_filter_with_root_projection() -> Result<()> {
443        let predicate =
444            json_get_expr(col("j"), path_expr("a"), Some(DataType::Int64))?.eq(lit(1_i64));
445        let (provider, plan) = build_json2_scan()?;
446        let plan = plan.filter(predicate)?.build()?;
447
448        assert!(
449            JsonTypeConcretizeRule
450                .rewrite(plan, &OptimizerContext::default())?
451                .transformed
452        );
453        assert_eq!(
454            Some(&JsonNativeType::Variant),
455            provider.scan_request().json_type_hint.get("j")
456        );
457        Ok(())
458    }
459
460    #[test]
461    fn test_deduce_json_type_with_non_column_base() -> Result<()> {
462        let expr = json_get_expr(
463            Expr::Literal(ScalarValue::Utf8(Some("{}".to_string())), None),
464            path_expr("a"),
465            Some(DataType::Int64),
466        )?;
467
468        let err = deduce_json_type(&expr).unwrap_err();
469        assert!(
470            err.to_string()
471                .contains("First argument of json_get is expected to be a column expr")
472        );
473        Ok(())
474    }
475
476    #[test]
477    fn test_deduce_json_type_with_non_literal_path() -> Result<()> {
478        let expr = json_get_expr(
479            Expr::Column(Column::new_unqualified("k0")),
480            Expr::Column(Column::new_unqualified("path_col")),
481            Some(DataType::Int64),
482        )?;
483
484        let err = deduce_json_type(&expr).unwrap_err();
485        assert!(
486            err.to_string()
487                .contains("Second argument of json_get is expected to be a string literal")
488        );
489        Ok(())
490    }
491
492    #[test]
493    fn test_deduce_json_type_default_string() -> Result<()> {
494        let expr = json_get_expr(
495            Expr::Column(Column::new_unqualified("k0")),
496            path_expr("a.b"),
497            None,
498        )?;
499
500        let deduced = deduce_json_type(&expr)?;
501        let expected = JsonNativeType::Object(JsonObjectType::from([(
502            "a".to_string(),
503            JsonNativeType::Object(JsonObjectType::from([(
504                "b".to_string(),
505                JsonNativeType::String,
506            )])),
507        )]));
508
509        assert_eq!(Some(("k0".to_string(), expected)), deduced);
510        Ok(())
511    }
512}
query/optimizer/json_type_concretize.rs

query/optimizer/
json_type_concretize.rs
