datatypes/extension/
json.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 arrow_schema::extension::{
19    EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType,
20};
21use arrow_schema::{ArrowError, DataType, Field, FieldRef};
22use parquet_variant_compute::VariantType;
23use serde::{Deserialize, Serialize};
24use snafu::{ResultExt, ensure};
25
26use crate::error::InvalidJson2LayoutSnafu;
27pub use crate::json::JSON2_REMAINDER_FIELD_NAME;
28use crate::json::JsonSettings;
29
30const LEGACY_JSON_STRUCTURE_SETTINGS_KEY: &str = "json_structure_settings";
31const JSON2_LAYOUT_V1: u8 = 1;
32const JSON2_LAYOUT_V2: u8 = 2;
33
34/// Parsed physical layout of a JSON2 Arrow root field.
35#[derive(Debug, Clone, PartialEq, Eq)]
36pub struct Json2PhysicalLayout {
37    version: u8,
38}
39
40impl Json2PhysicalLayout {
41    /// Parses the JSON2 layout version from a root field.
42    ///
43    /// The version is read from the extension metadata; missing metadata
44    /// defaults to V1.
45    ///
46    /// Supported versions:
47    /// - V1
48    /// - V2
49    ///
50    /// Errors when the field is not a JSON2 extension or the version is not
51    /// one of the supported versions.
52    pub fn try_from_root(field: &Field) -> crate::error::Result<Self> {
53        ensure!(
54            is_json2_extension_type(field),
55            InvalidJson2LayoutSnafu {
56                reason: format!("field '{}' is not a JSON2 extension", field.name()),
57            }
58        );
59
60        let version = match field.extension_type_name() {
61            Some(Json2ExtensionType::NAME) => field
62                .metadata()
63                .get(EXTENSION_TYPE_METADATA_KEY)
64                .map(|x| parse_version(x))
65                .transpose()?
66                .flatten()
67                .unwrap_or(JSON2_LAYOUT_V1),
68            _ => JSON2_LAYOUT_V1,
69        };
70        ensure!(
71            matches!(version, JSON2_LAYOUT_V1 | JSON2_LAYOUT_V2),
72            InvalidJson2LayoutSnafu {
73                reason: format!("unsupported JSON2 layout version: {version}"),
74            }
75        );
76        Ok(Self { version })
77    }
78
79    /// Returns whether this is the JSON2 physical layout version 2.
80    pub fn is_version_2(&self) -> bool {
81        self.version == JSON2_LAYOUT_V2
82    }
83}
84
85fn parse_version(metadata: &str) -> crate::error::Result<Option<u8>> {
86    serde_json::from_str::<JsonMetadata>(metadata)
87        .map(|x| x.layout_version)
88        .map_err(|e| {
89            InvalidJson2LayoutSnafu {
90                reason: format!(r#"invalid extension metadata: "{metadata}", error: {e}"#),
91            }
92            .build()
93        })
94}
95
96/// Returns the remainder field of a JSON2 v2 root.
97pub(crate) fn json2_remainder_field(field: &Field) -> crate::error::Result<Option<&FieldRef>> {
98    let layout = Json2PhysicalLayout::try_from_root(field)?;
99    if !layout.is_version_2() {
100        return Ok(None);
101    }
102
103    let remainder = if let DataType::Struct(fields) = field.data_type() {
104        fields
105            .iter()
106            .find(|x| x.name() == JSON2_REMAINDER_FIELD_NAME)
107    } else {
108        None
109    };
110
111    if let Some(remainder) = remainder {
112        let _ = remainder.try_extension_type::<VariantType>().map_err(|e| {
113            InvalidJson2LayoutSnafu {
114                reason: e.to_string(),
115            }
116            .build()
117        })?;
118    }
119    Ok(remainder)
120}
121
122#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
123pub struct JsonMetadata {
124    /// JSON2 settings stored in Arrow extension metadata.
125    json_settings: JsonSettings,
126    /// Physical JSON2 layout used by this Arrow field.
127    #[serde(default, skip_serializing_if = "Option::is_none")]
128    layout_version: Option<u8>,
129}
130
131impl JsonMetadata {
132    /// Creates metadata for the JSON2 layout version 2.
133    pub fn new(json_settings: JsonSettings) -> Self {
134        Self {
135            json_settings,
136            layout_version: Some(JSON2_LAYOUT_V2),
137        }
138    }
139
140    /// Creates metadata for the legacy JSON2 layout.
141    pub fn new_v1(json_settings: JsonSettings) -> Self {
142        Self {
143            json_settings,
144            layout_version: None,
145        }
146    }
147
148    /// Returns the JSON2 settings.
149    pub fn json_settings(&self) -> &JsonSettings {
150        &self.json_settings
151    }
152
153    /// Consumes the metadata and returns its JSON2 settings.
154    pub fn into_json_settings(self) -> JsonSettings {
155        self.json_settings
156    }
157
158    /// Returns whether this metadata describes JSON2 layout version 2.
159    pub fn is_version_2(&self) -> bool {
160        self.layout_version == Some(JSON2_LAYOUT_V2)
161    }
162}
163
164impl Default for JsonMetadata {
165    fn default() -> Self {
166        Self::new(JsonSettings::default())
167    }
168}
169
170/// Arrow extension type for legacy JSONB columns.
171#[derive(Debug, Clone, Default)]
172pub struct JsonExtensionType;
173
174impl ExtensionType for JsonExtensionType {
175    const NAME: &'static str = "greptime.json";
176    type Metadata = ();
177
178    fn metadata(&self) -> &Self::Metadata {
179        &()
180    }
181
182    fn serialize_metadata(&self) -> Option<String> {
183        None
184    }
185
186    fn deserialize_metadata(_metadata: Option<&str>) -> Result<Self::Metadata, ArrowError> {
187        Ok(())
188    }
189
190    fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
191        match data_type {
192            DataType::Binary | DataType::Null => Ok(()),
193            t => Err(ArrowError::InvalidArgumentError(format!(
194                "Unexpected data type {t} for JsonExtensionType"
195            ))),
196        }
197    }
198
199    fn try_new(data_type: &DataType, _metadata: Self::Metadata) -> Result<Self, ArrowError> {
200        Self.supports_data_type(data_type).map(|_| Self)
201    }
202}
203
204/// Arrow extension type for JSON2 columns and concretized projections.
205#[derive(Debug, Clone, Default)]
206pub struct Json2ExtensionType(Arc<JsonMetadata>);
207
208impl Json2ExtensionType {
209    /// Creates a JSON2 extension type with the given metadata.
210    pub fn new(metadata: Arc<JsonMetadata>) -> Self {
211        Self(metadata)
212    }
213}
214
215impl ExtensionType for Json2ExtensionType {
216    const NAME: &'static str = "greptime.json2";
217    type Metadata = Arc<JsonMetadata>;
218
219    fn metadata(&self) -> &Self::Metadata {
220        &self.0
221    }
222
223    fn serialize_metadata(&self) -> Option<String> {
224        serde_json::to_string(self.metadata()).ok()
225    }
226
227    fn deserialize_metadata(metadata: Option<&str>) -> Result<Self::Metadata, ArrowError> {
228        if let Some(metadata) = metadata {
229            let metadata = serde_json::from_str(metadata).map_err(|e| {
230                ArrowError::ParseError(format!("Failed to deserialize JSON metadata: {}", e))
231            })?;
232            Ok(Arc::new(metadata))
233        } else {
234            Ok(Arc::new(JsonMetadata::new_v1(JsonSettings::default())))
235        }
236    }
237
238    fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
239        match data_type {
240            // object
241            DataType::Struct(_)
242            // array
243            | DataType::List(_)
244            | DataType::ListView(_)
245            | DataType::LargeList(_)
246            | DataType::LargeListView(_)
247            // string
248            | DataType::Utf8
249            | DataType::Utf8View
250            | DataType::LargeUtf8
251            // number
252            | DataType::Int8
253            | DataType::Int16
254            | DataType::Int32
255            | DataType::Int64
256            | DataType::UInt8
257            | DataType::UInt16
258            | DataType::UInt32
259            | DataType::UInt64
260            | DataType::Float32
261            | DataType::Float64
262            // boolean
263            | DataType::Boolean
264            // null
265            | DataType::Null
266            // legacy json type
267            | DataType::Binary => Ok(()),
268            dt => Err(ArrowError::SchemaError(format!(
269                "Unexpected data type {dt}"
270            ))),
271        }
272    }
273
274    fn try_new(data_type: &DataType, metadata: Self::Metadata) -> Result<Self, ArrowError> {
275        let json = Self(metadata);
276        json.supports_data_type(data_type)?;
277        Ok(json)
278    }
279}
280
281/// Checks whether this field is either a legacy JSONB or JSON2 extension type.
282pub fn is_any_json_extension_type<T: AsRef<Field>>(field: T) -> bool {
283    let name = field.as_ref().extension_type_name();
284    name == Some(JsonExtensionType::NAME) || name == Some(Json2ExtensionType::NAME)
285}
286
287/// Parses JSON2 settings stored by the historical `greptime.json` extension.
288pub fn parse_legacy_json2_settings(
289    metadata: &HashMap<String, String>,
290) -> crate::error::Result<Option<JsonSettings>> {
291    #[derive(Deserialize)]
292    struct LegacyJsonMetadata {
293        #[serde(default)]
294        json_settings: Option<JsonSettings>,
295    }
296
297    if metadata.get(EXTENSION_TYPE_NAME_KEY).map(String::as_str) != Some(JsonExtensionType::NAME) {
298        return Ok(None);
299    }
300
301    metadata
302        .get(EXTENSION_TYPE_METADATA_KEY)
303        .map(|json| {
304            serde_json::from_str::<LegacyJsonMetadata>(json)
305                .map(|x| x.json_settings)
306                .context(crate::error::DeserializeSnafu { json })
307        })
308        .transpose()
309        .map(Option::flatten)
310}
311
312/// Checks whether this field uses the JSON2 extension layout from before type hints.
313///
314/// That layout used the same `greptime.json` extension name and
315/// `json_structure_settings` metadata as legacy JSONB. Its structured Arrow data type is
316/// therefore required to distinguish JSON2 from Binary JSONB.
317pub fn is_legacy_json2_extension_type<T: AsRef<Field>>(field: T) -> bool {
318    let field = field.as_ref();
319    if field.extension_type_name() != Some(JsonExtensionType::NAME)
320        || !matches!(field.data_type(), DataType::Struct(_))
321    {
322        return false;
323    }
324
325    field
326        .metadata()
327        .get(EXTENSION_TYPE_METADATA_KEY)
328        .and_then(|json| serde_json::from_str::<serde_json::Value>(json).ok())
329        .is_some_and(|metadata| metadata.get(LEGACY_JSON_STRUCTURE_SETTINGS_KEY).is_some())
330}
331
332/// Check if this field is a JSON2 extension type.
333///
334/// New schemas use [`Json2ExtensionType`]. For compatibility, old fields using
335/// [`JsonExtensionType`] with JSON settings or the pre-type-hint structured layout are also
336/// recognized as JSON2.
337pub fn is_json2_extension_type<T: AsRef<Field>>(field: T) -> bool {
338    let field = field.as_ref();
339    field.extension_type_name() == Some(Json2ExtensionType::NAME)
340        || parse_legacy_json2_settings(field.metadata()).is_ok_and(|x| x.is_some())
341        || is_legacy_json2_extension_type(field)
342}
343
344#[cfg(test)]
345mod tests {
346    use std::collections::HashMap;
347
348    use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY};
349    use arrow_schema::{Field, Fields};
350
351    use super::*;
352    use crate::vectors::json::variant::variant_field;
353
354    #[test]
355    fn test_json2_extension_type_detection() {
356        let extension = Json2ExtensionType::new(Arc::new(JsonMetadata::default()));
357        let json2 = Field::new("j", DataType::Struct(Fields::empty()), true)
358            .with_extension_type(extension.clone());
359        // "projection" is the special hack for selecting the whole column of json2
360        let projection = Field::new("j", DataType::Binary, true).with_extension_type(extension);
361        let legacy_json2 = Field::new("j", DataType::Struct(Fields::empty()), true).with_metadata(
362            HashMap::from([
363                (
364                    EXTENSION_TYPE_NAME_KEY.to_string(),
365                    JsonExtensionType::NAME.to_string(),
366                ),
367                (
368                    EXTENSION_TYPE_METADATA_KEY.to_string(),
369                    serde_json::json!({ "json_settings": JsonSettings::default() }).to_string(),
370                ),
371            ]),
372        );
373        // Before type hints, JSON2 and JSONB shared extension metadata and were distinguished by
374        // their physical Arrow data types.
375        let legacy_structure_metadata = HashMap::from([
376            (
377                EXTENSION_TYPE_NAME_KEY.to_string(),
378                JsonExtensionType::NAME.to_string(),
379            ),
380            (
381                EXTENSION_TYPE_METADATA_KEY.to_string(),
382                serde_json::json!({
383                    (LEGACY_JSON_STRUCTURE_SETTINGS_KEY): { "Structured": null }
384                })
385                .to_string(),
386            ),
387        ]);
388        let pre_type_hint_json2 = Field::new("j", DataType::Struct(Fields::empty()), true)
389            .with_metadata(legacy_structure_metadata.clone());
390        let legacy_jsonb =
391            Field::new("j", DataType::Binary, true).with_metadata(legacy_structure_metadata);
392
393        assert!(is_json2_extension_type(&json2));
394        assert!(is_json2_extension_type(&projection));
395        assert!(is_json2_extension_type(&legacy_json2));
396        assert!(is_legacy_json2_extension_type(&pre_type_hint_json2));
397        assert!(is_json2_extension_type(&pre_type_hint_json2));
398        assert_eq!(
399            Some(JsonSettings::default()),
400            parse_legacy_json2_settings(legacy_json2.metadata()).unwrap()
401        );
402        assert!(!is_legacy_json2_extension_type(&legacy_jsonb));
403        assert!(!is_json2_extension_type(&legacy_jsonb));
404        assert!(JsonExtensionType::try_new(&DataType::Binary, ()).is_ok());
405        assert!(JsonExtensionType::try_new(&DataType::Null, ()).is_ok());
406        assert!(JsonExtensionType::try_new(&DataType::Struct(Fields::empty()), ()).is_err());
407    }
408
409    #[test]
410    fn test_json_metadata_layout_version_compatibility() -> serde_json::Result<()> {
411        let legacy: JsonMetadata = serde_json::from_str(r#"{"json_settings":{}}"#)?;
412        assert!(!legacy.is_version_2());
413
414        let metadata = JsonMetadata::new(JsonSettings::default());
415        assert!(metadata.is_version_2());
416        let serialized = serde_json::to_string(&metadata)?;
417        let deserialized: JsonMetadata = serde_json::from_str(&serialized)?;
418        assert!(deserialized.is_version_2());
419        assert_eq!(deserialized, metadata);
420        Ok(())
421    }
422
423    #[test]
424    fn test_parse_json2_physical_layout() -> crate::error::Result<()> {
425        let legacy = Field::new("data", DataType::Struct(Fields::empty()), true)
426            .with_extension_type(Json2ExtensionType::new(Arc::new(JsonMetadata::new_v1(
427                JsonSettings::default(),
428            ))));
429        assert!(!Json2PhysicalLayout::try_from_root(&legacy)?.is_version_2());
430
431        let v2 = Field::new(
432            "data",
433            DataType::Struct(
434                vec![
435                    Arc::new(variant_field(JSON2_REMAINDER_FIELD_NAME, true)),
436                    Arc::new(Field::new("count", DataType::Int64, true)),
437                ]
438                .into(),
439            ),
440            true,
441        )
442        .with_extension_type(Json2ExtensionType::new(Arc::new(JsonMetadata::new(
443            JsonSettings::default(),
444        ))));
445        assert!(Json2PhysicalLayout::try_from_root(&v2)?.is_version_2());
446        let remainder = json2_remainder_field(&v2)?.unwrap();
447        assert_eq!(JSON2_REMAINDER_FIELD_NAME, remainder.name());
448        Ok(())
449    }
450
451    #[test]
452    fn test_reject_invalid_json2_physical_layout() {
453        let field = Field::new("data", DataType::Struct(Fields::empty()), true);
454        assert!(Json2PhysicalLayout::try_from_root(&field).is_err());
455
456        let metadata =
457            Json2ExtensionType::new(Arc::new(JsonMetadata::new(JsonSettings::default())));
458        let missing = Field::new("data", DataType::Struct(Fields::empty()), true)
459            .with_extension_type(metadata.clone());
460        assert!(json2_remainder_field(&missing).is_ok_and(|x| x.is_none()));
461
462        let invalid = Field::new(
463            "data",
464            DataType::Struct(
465                vec![Arc::new(Field::new(
466                    JSON2_REMAINDER_FIELD_NAME,
467                    DataType::Binary,
468                    true,
469                ))]
470                .into(),
471            ),
472            true,
473        )
474        .with_extension_type(metadata);
475        assert!(json2_remainder_field(&invalid).is_err());
476
477        let future = Field::new("data", DataType::Struct(Fields::empty()), true)
478            .with_extension_type(Json2ExtensionType::new(Arc::new(JsonMetadata {
479                json_settings: JsonSettings::default(),
480                layout_version: Some(JSON2_LAYOUT_V2 + 1),
481            })));
482        assert!(Json2PhysicalLayout::try_from_root(&future).is_err());
483    }
484}
datatypes/extension/json.rs

datatypes/extension/
json.rs
