query/optimizer/
json_schema_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;
16use std::sync::Arc;
17
18use datafusion::config::ConfigOptions;
19use datafusion_common::tree_node::Transformed;
20use datafusion_common::{DFSchema, DFSchemaRef, Result};
21use datafusion_expr::{LogicalPlan, UserDefinedLogicalNodeCore};
22use datafusion_optimizer::analyzer::AnalyzerRule;
23use datatypes::extension::json::{
24    Json2ExtensionType, is_json2_extension_type, is_legacy_json2_extension_type,
25};
26use datatypes::types::json_type::JsonNativeType;
27
28use crate::dist_plan::MergeScanLogicalPlan;
29use crate::optimizer::json_type_concretize::deduce_json_types;
30
31/// Keeps JSON2 schemas consistent across distributed query boundaries.
32///
33/// An unresolved JSON2 column is represented as an empty `Struct`, while a remote stage always
34/// emits a concrete Arrow array. DataFusion expects the schema declared by a logical node, the
35/// schema used to build its physical plan, and the schema of its record batches to agree. This rule
36/// gives each [`MergeScanLogicalPlan`] the concrete schema emitted by its remote stage and propagates
37/// that schema through the local plan.
38///
39/// For example:
40///
41/// - `SELECT j FROM t` transfers the complete JSON2 value as `Binary` (`Variant`). A projection or
42///   window above the MergeScan must therefore also describe `j` as `Binary`, not an empty `Struct`.
43/// - `SELECT j.a FROM t` may transfer the complete `j` as `Binary` and extract `a` locally, or
44///   transfer only the extracted scalar when the expression runs remotely. The boundary schema
45///   must describe the remote output rather than the local expression that consumes it.
46/// - `SELECT l.j, r.j FROM l JOIN r ON l.k = r.k` has two independent boundaries. Each input and the
47///   join schema must agree on the concrete type of its JSON2 column.
48///
49/// Correcting only the physical MergeScan schema can make simple queries work because many
50/// operators access columns by position, but it leaves the logical plan describing a different
51/// type. Keeping the schemas consistent lets optimizers, physical planners, validators, and future
52/// type-aware operators rely on the normal DataFusion contract. It also keeps the generic physical
53/// MergeScan implementation independent of JSON2.
54///
55/// The boundary schema and the storage read layout answer different questions. For
56/// `SELECT j.a::BIGINT FROM t`, the storage scan may use a structured `{a: Int64}` layout, while the
57/// distributed boundary can still emit either the complete `j` as `Binary` or only `a` as `Int64`.
58/// Therefore this rule cannot replace the separate JSON2 scan-type inference.
59#[derive(Debug)]
60pub(crate) struct JsonSchemaConcretizeRule;
61
62impl AnalyzerRule for JsonSchemaConcretizeRule {
63    fn analyze(&self, plan: LogicalPlan, _config: &ConfigOptions) -> Result<LogicalPlan> {
64        let plan = plan.transform_up_with_subqueries(|plan| {
65            let LogicalPlan::Extension(mut extension) = plan else {
66                return Ok(Transformed::no(plan));
67            };
68            let Some(merge_scan) = extension
69                .node
70                .as_any()
71                .downcast_ref::<MergeScanLogicalPlan>()
72            else {
73                return Ok(Transformed::no(LogicalPlan::Extension(extension)));
74            };
75
76            // Infer the boundary schema from the hidden remote plan, not its local consumer.
77            let json_types = deduce_json_types(merge_scan.input())?;
78            if json_types.is_empty() {
79                return Ok(Transformed::no(LogicalPlan::Extension(extension)));
80            }
81            let schema = concretize_json2_schema(merge_scan.schema(), &json_types)?;
82            if schema.as_ref() == merge_scan.schema().as_ref() {
83                return Ok(Transformed::no(LogicalPlan::Extension(extension)));
84            }
85
86            extension.node = Arc::new(merge_scan.clone().with_output_schema(schema));
87            Ok(Transformed::yes(LogicalPlan::Extension(extension)))
88        })?;
89
90        if plan.transformed {
91            plan.data
92                .transform_up_with_subqueries(|plan| {
93                    if matches!(plan, LogicalPlan::Extension(_)) || plan.inputs().is_empty() {
94                        Ok(Transformed::no(plan))
95                    } else {
96                        plan.recompute_schema().map(Transformed::yes)
97                    }
98                })
99                .map(|x| x.data)
100        } else {
101            Ok(plan.data)
102        }
103    }
104
105    fn name(&self) -> &str {
106        "JsonSchemaConcretizeRule"
107    }
108}
109
110fn concretize_json2_schema(
111    schema: &DFSchemaRef,
112    json_types: &HashMap<String, JsonNativeType>,
113) -> Result<DFSchemaRef> {
114    if !schema
115        .iter()
116        .any(|(_, field)| json_types.contains_key(field.name()) && is_json2_extension_type(field))
117    {
118        return Ok(schema.clone());
119    }
120
121    let mut changed = false;
122    let fields = schema
123        .iter()
124        .map(|(qualifier, field)| {
125            let Some(json_type) = json_types
126                .get(field.name())
127                .filter(|_| is_json2_extension_type(field))
128            else {
129                return (qualifier.cloned(), field.clone());
130            };
131            let data_type = json_type.as_arrow_type();
132            if field.data_type() == &data_type {
133                return (qualifier.cloned(), field.clone());
134            }
135
136            changed = true;
137
138            // Before type hints, JSON2 used the `greptime.json` marker together with
139            // `json_structure_settings`. Once concretized to `Binary`, that field no longer
140            // matches the legacy JSON2 shape and could be mistaken for JSONB, so upgrade its
141            // marker. Do not replace modern markers because that would discard their JSON
142            // settings and layout version.
143            let legacy = is_legacy_json2_extension_type(field);
144            let mut field = field.as_ref().clone().with_data_type(data_type);
145            if legacy {
146                field = field.with_extension_type(Json2ExtensionType::default());
147            }
148            (qualifier.cloned(), Arc::new(field))
149        })
150        .collect();
151
152    if changed {
153        let schema = DFSchema::new_with_metadata(fields, schema.metadata().clone())?
154            .with_functional_dependencies(schema.functional_dependencies().clone())?;
155        Ok(Arc::new(schema))
156    } else {
157        Ok(schema.clone())
158    }
159}
160
161#[cfg(test)]
162mod tests {
163    use arrow_schema::extension::{
164        EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType,
165    };
166    use arrow_schema::{DataType, Field, Fields, Schema};
167    use datafusion_common::DFSchema;
168    use datafusion_expr::{LogicalPlanBuilder, col};
169    use datatypes::extension::json::{JsonExtensionType, is_json2_extension_type};
170
171    use super::*;
172
173    #[test]
174    fn test_json_schema_concretize_rule_updates_merge_scan() -> Result<()> {
175        let field = Field::new("j", DataType::Struct(Fields::empty()), true).with_metadata(
176            HashMap::from([
177                (
178                    EXTENSION_TYPE_NAME_KEY.to_string(),
179                    JsonExtensionType::NAME.to_string(),
180                ),
181                (
182                    EXTENSION_TYPE_METADATA_KEY.to_string(),
183                    serde_json::json!({
184                        "json_structure_settings": { "Structured": null }
185                    })
186                    .to_string(),
187                ),
188            ]),
189        );
190        let schema = Arc::new(DFSchema::try_from(Schema::new(vec![field]))?);
191        let input = LogicalPlan::EmptyRelation(datafusion_expr::logical_plan::EmptyRelation {
192            produce_one_row: false,
193            schema,
194        });
195        let merge_scan =
196            MergeScanLogicalPlan::new(input, false, Default::default()).into_logical_plan();
197        let plan = LogicalPlanBuilder::from(merge_scan)
198            .project(vec![col("j")])?
199            .build()?;
200
201        let plan = JsonSchemaConcretizeRule.analyze(plan, &ConfigOptions::default())?;
202        let field = plan.schema().field(0);
203        assert_eq!(&DataType::Binary, field.data_type());
204        assert_eq!(Some(Json2ExtensionType::NAME), field.extension_type_name());
205        assert!(is_json2_extension_type(field));
206        Ok(())
207    }
208}
query/optimizer/json_schema_concretize.rs

query/optimizer/
json_schema_concretize.rs
