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::sync::Arc;
16
17use arrow_schema::extension::ExtensionType;
18use arrow_schema::{ArrowError, DataType};
19use serde::{Deserialize, Serialize};
20
21use crate::json::JsonStructureSettings;
22
23#[derive(Debug, Clone, Serialize, Deserialize, Default)]
24pub struct JsonMetadata {
25    /// Indicates how to handle JSON is stored in underlying data type
26    ///
27    /// This field can be `None` for data is converted to complete structured in-memory form.
28    pub json_structure_settings: Option<JsonStructureSettings>,
29}
30
31#[derive(Debug, Clone)]
32pub struct JsonExtensionType(Arc<JsonMetadata>);
33
34impl JsonExtensionType {
35    pub fn new(metadata: Arc<JsonMetadata>) -> Self {
36        JsonExtensionType(metadata)
37    }
38}
39
40impl ExtensionType for JsonExtensionType {
41    const NAME: &'static str = "greptime.json";
42    type Metadata = Arc<JsonMetadata>;
43
44    fn metadata(&self) -> &Self::Metadata {
45        &self.0
46    }
47
48    fn serialize_metadata(&self) -> Option<String> {
49        serde_json::to_string(self.metadata()).ok()
50    }
51
52    fn deserialize_metadata(metadata: Option<&str>) -> Result<Self::Metadata, ArrowError> {
53        if let Some(metadata) = metadata {
54            let metadata = serde_json::from_str(metadata).map_err(|e| {
55                ArrowError::ParseError(format!("Failed to deserialize JSON metadata: {}", e))
56            })?;
57            Ok(Arc::new(metadata))
58        } else {
59            Ok(Arc::new(JsonMetadata::default()))
60        }
61    }
62
63    fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
64        match data_type {
65            // object
66            DataType::Struct(_)
67            // array
68            | DataType::List(_)
69            | DataType::ListView(_)
70            | DataType::LargeList(_)
71            | DataType::LargeListView(_)
72            // string
73            | DataType::Utf8
74            | DataType::Utf8View
75            | DataType::LargeUtf8
76            // number
77            | DataType::Int8
78            | DataType::Int16
79            | DataType::Int32
80            | DataType::Int64
81            | DataType::UInt8
82            | DataType::UInt16
83            | DataType::UInt32
84            | DataType::UInt64
85            | DataType::Float32
86            | DataType::Float64
87            // boolean
88            | DataType::Boolean
89            // null
90            | DataType::Null
91            // legacy json type
92            | DataType::Binary => Ok(()),
93            dt => Err(ArrowError::SchemaError(format!(
94                "Unexpected data type {dt}"
95            ))),
96        }
97    }
98
99    fn try_new(data_type: &DataType, metadata: Self::Metadata) -> Result<Self, ArrowError> {
100        let json = Self(metadata);
101        json.supports_data_type(data_type)?;
102        Ok(json)
103    }
104}