mito2/sst/parquet/json_align/
schema.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::BTreeMap;
16use std::sync::Arc;
17
18use arrow_schema::{DataType as ArrowDataType, FieldRef};
19use datatypes::arrow::datatypes::Schema;
20use datatypes::extension::json::is_json2_extension_type;
21use store_api::storage::NestedPath;
22
23/// Aligns nested struct fields according to the requested nested paths.
24///
25/// For each root field:
26/// - An empty path list keeps the whole field unchanged.
27/// - Non-JSON root fields ignore nested paths and keep the whole field unchanged.
28/// - JSON2 root fields are rebuilt from `nested_paths`.
29/// - Existing schema fields are preserved only when they are the requested leaf.
30/// - Requested paths missing from the schema are synthesized with JSONB (`Binary`)
31///   leaves.
32///
33/// For example, if the schema has `j: struct<a: struct<x: Int64, y: Utf8>>`
34/// and `nested_paths` requests `j.a.x` and `j.a.z`, the result is
35/// `j: struct<a: struct<x: Int64, z: Binary>>`.
36pub(crate) fn align_schema_by_nested_paths<'a, I>(schema: &mut Schema, nested_paths: I)
37where
38    I: IntoIterator<Item = &'a [NestedPath]>,
39{
40    let fields = schema
41        .fields
42        .into_iter()
43        .zip(nested_paths)
44        .map(|(field, paths)| {
45            if !paths.is_empty() && is_json2_extension_type(field) {
46                let child_paths = paths
47                    .iter()
48                    .map(|path| {
49                        if path.first().is_some_and(|root| root == field.name()) {
50                            &path[1..]
51                        } else {
52                            path
53                        }
54                    })
55                    .collect::<Vec<_>>();
56                rebuild_field_by_nested_paths(field, &child_paths)
57            } else {
58                field.clone()
59            }
60        })
61        .collect::<Vec<_>>();
62    schema.fields = fields.into()
63}
64
65fn rebuild_field_by_nested_paths(field: &FieldRef, nested_paths: &[&[String]]) -> FieldRef {
66    if nested_paths.iter().any(|path| path.is_empty()) {
67        return field.clone();
68    };
69
70    let fields = group_child_paths(nested_paths)
71        .into_iter()
72        .map(|(name, paths)| {
73            let existing = find_struct_child(field, &name);
74            build_field_from_nested_paths(&name, existing, &paths)
75        })
76        .collect::<Vec<_>>();
77
78    Arc::new(
79        field
80            .as_ref()
81            .clone()
82            .with_data_type(ArrowDataType::Struct(fields.into())),
83    )
84}
85
86fn group_child_paths<'a>(nested_paths: &[&'a [String]]) -> BTreeMap<String, Vec<&'a [String]>> {
87    let mut child_paths = BTreeMap::<String, Vec<&'a [String]>>::new();
88    for path in nested_paths {
89        let Some((name, remaining)) = path.split_first() else {
90            continue;
91        };
92        child_paths.entry(name.clone()).or_default().push(remaining);
93    }
94    child_paths
95}
96
97fn find_struct_child<'a>(field: &'a FieldRef, name: &str) -> Option<&'a FieldRef> {
98    let ArrowDataType::Struct(fields) = field.data_type() else {
99        return None;
100    };
101    fields.iter().find(|field| field.name() == name)
102}
103
104fn build_field_from_nested_paths(
105    name: &str,
106    existing: Option<&FieldRef>,
107    nested_paths: &[&[String]],
108) -> FieldRef {
109    if nested_paths.iter().any(|path| path.is_empty()) {
110        return existing.cloned().unwrap_or_else(|| new_jsonb_field(name));
111    }
112
113    let fields = group_child_paths(nested_paths)
114        .into_iter()
115        .map(|(name, paths)| {
116            let existing_child = existing.and_then(|field| find_struct_child(field, &name));
117            build_field_from_nested_paths(&name, existing_child, &paths)
118        })
119        .collect::<Vec<_>>();
120
121    let field = existing
122        .map(|field| field.as_ref().clone())
123        .unwrap_or_else(|| arrow_schema::Field::new(name, ArrowDataType::Binary, true));
124    Arc::new(field.with_data_type(ArrowDataType::Struct(fields.into())))
125}
126
127fn new_jsonb_field(name: &str) -> FieldRef {
128    Arc::new(arrow_schema::Field::new(name, ArrowDataType::Binary, true))
129}
130
131#[cfg(test)]
132mod tests {
133    use arrow_schema::Field;
134    use datatypes::extension::json::Json2ExtensionType;
135
136    use super::*;
137
138    #[test]
139    fn test_align_schema_by_nested_paths() {
140        fn new_field(name: &str, data_type: ArrowDataType) -> FieldRef {
141            Arc::new(Field::new(name, data_type, true))
142        }
143
144        fn struct_field(name: &str, fields: impl IntoIterator<Item = FieldRef>) -> FieldRef {
145            new_field(name, ArrowDataType::Struct(fields.into_iter().collect()))
146        }
147
148        fn json_struct_field(name: &str, fields: impl IntoIterator<Item = FieldRef>) -> FieldRef {
149            Arc::new(
150                Field::new(
151                    name,
152                    ArrowDataType::Struct(fields.into_iter().collect()),
153                    true,
154                )
155                .with_extension_type(Json2ExtensionType::default()),
156            )
157        }
158
159        let mut schema = Schema::new([
160            json_struct_field(
161                "j",
162                [
163                    struct_field(
164                        "a",
165                        [
166                            new_field("x", ArrowDataType::Int64),
167                            new_field("y", ArrowDataType::Utf8),
168                            struct_field(
169                                "z",
170                                [
171                                    new_field("q", ArrowDataType::Boolean),
172                                    new_field("r", ArrowDataType::Float64),
173                                ],
174                            ),
175                        ],
176                    ),
177                    new_field("b", ArrowDataType::Utf8),
178                    struct_field(
179                        "c",
180                        vec![
181                            new_field("d", ArrowDataType::Int64),
182                            new_field("e", ArrowDataType::Utf8),
183                        ],
184                    ),
185                ],
186            ),
187            new_field("tag", ArrowDataType::Utf8),
188            struct_field(
189                "k",
190                [
191                    new_field("k_0", ArrowDataType::Int64),
192                    new_field("k_1", ArrowDataType::Utf8),
193                ],
194            ),
195        ]);
196
197        let nested_paths = [
198            vec![
199                ["j", "a", "x"].iter().map(|x| x.to_string()).collect(),
200                ["j", "a", "z", "q"].iter().map(|x| x.to_string()).collect(),
201                ["j", "a", "m"].iter().map(|x| x.to_string()).collect(),
202                ["j", "b", "x"].iter().map(|x| x.to_string()).collect(),
203                ["j", "c"].iter().map(|x| x.to_string()).collect(),
204                ["j", "d", "e"].iter().map(|x| x.to_string()).collect(),
205            ],
206            vec![["tag", "ignored"].iter().map(|x| x.to_string()).collect()],
207            vec![],
208        ];
209
210        align_schema_by_nested_paths(
211            &mut schema,
212            nested_paths.iter().map(|paths| paths.as_slice()),
213        );
214
215        let expected = Schema::new([
216            json_struct_field(
217                "j",
218                [
219                    struct_field(
220                        "a",
221                        [
222                            new_field("m", ArrowDataType::Binary),
223                            new_field("x", ArrowDataType::Int64),
224                            struct_field("z", vec![new_field("q", ArrowDataType::Boolean)]),
225                        ],
226                    ),
227                    struct_field("b", [new_field("x", ArrowDataType::Binary)]),
228                    struct_field(
229                        "c",
230                        [
231                            new_field("d", ArrowDataType::Int64),
232                            new_field("e", ArrowDataType::Utf8),
233                        ],
234                    ),
235                    struct_field("d", [new_field("e", ArrowDataType::Binary)]),
236                ],
237            ),
238            new_field("tag", ArrowDataType::Utf8),
239            struct_field(
240                "k",
241                [
242                    new_field("k_0", ArrowDataType::Int64),
243                    new_field("k_1", ArrowDataType::Utf8),
244                ],
245            ),
246        ]);
247
248        assert_eq!(schema, expected);
249    }
250}
mito2/sst/parquet/json_align/schema.rs

mito2/sst/parquet/json_align/
schema.rs
