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

datatypes/extension/
json.rs
