refactor: separate a json2 extension type (#8745)

Signed-off-by: luofucong <luofc@foxmail.com>

src/datatypes/src/extension/json.rs

@@ -12,32 +12,85 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashMap;
 use std::sync::Arc;
 
-use arrow_schema::extension::ExtensionType;
-use arrow_schema::{ArrowError, DataType, Field, FieldRef};
+use arrow_schema::extension::{
+    EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType,
+};
+use arrow_schema::{ArrowError, DataType, Field};
 use serde::{Deserialize, Serialize};
+use snafu::ResultExt;
 
 use crate::json::JsonSettings;
 
+const LEGACY_JSON_STRUCTURE_SETTINGS_KEY: &str = "json_structure_settings";
+
 #[derive(Debug, Clone, Serialize, Deserialize, Default)]
 pub struct JsonMetadata {
-    /// JSON2 settings stored in column schema metadata and represented through
-    /// Arrow extension metadata.
-    pub json_settings: Option<JsonSettings>,
+    /// JSON2 settings stored in Arrow extension metadata.
+    json_settings: JsonSettings,
 }
 
-#[derive(Debug, Clone)]
-pub struct JsonExtensionType(Arc<JsonMetadata>);
+impl JsonMetadata {
+    /// Creates JSON2 extension metadata.
+    pub fn new(json_settings: JsonSettings) -> Self {
+        Self { json_settings }
+    }
 
-impl JsonExtensionType {
-    pub fn new(metadata: Arc<JsonMetadata>) -> Self {
-        JsonExtensionType(metadata)
+    /// Returns the JSON2 settings.
+    pub fn json_settings(&self) -> &JsonSettings {
+        &self.json_settings
     }
 }
 
+/// Arrow extension type for legacy JSONB columns.
+#[derive(Debug, Clone, Default)]
+pub struct JsonExtensionType;
+
 impl ExtensionType for JsonExtensionType {
     const NAME: &'static str = "greptime.json";
+    type Metadata = ();
+
+    fn metadata(&self) -> &Self::Metadata {
+        &()
+    }
+
+    fn serialize_metadata(&self) -> Option<String> {
+        None
+    }
+
+    fn deserialize_metadata(_metadata: Option<&str>) -> Result<Self::Metadata, ArrowError> {
+        Ok(())
+    }
+
+    fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
+        match data_type {
+            DataType::Binary | DataType::Null => Ok(()),
+            t => Err(ArrowError::InvalidArgumentError(format!(
+                "Unexpected data type {t} for JsonExtensionType"
+            ))),
+        }
+    }
+
+    fn try_new(data_type: &DataType, _metadata: Self::Metadata) -> Result<Self, ArrowError> {
+        Self.supports_data_type(data_type).map(|_| Self)
+    }
+}
+
+/// Arrow extension type for JSON2 columns and concretized projections.
+#[derive(Debug, Clone, Default)]
+pub struct Json2ExtensionType(Arc<JsonMetadata>);
+
+impl Json2ExtensionType {
+    /// Creates a JSON2 extension type with the given metadata.
+    pub fn new(metadata: Arc<JsonMetadata>) -> Self {
+        Self(metadata)
+    }
+}
+
+impl ExtensionType for Json2ExtensionType {
+    const NAME: &'static str = "greptime.json2";
     type Metadata = Arc<JsonMetadata>;
 
     fn metadata(&self) -> &Self::Metadata {
@@ -102,29 +155,130 @@ impl ExtensionType for JsonExtensionType {
     }
 }
 
-/// Check if this field is to be treated as json extension type.
-pub fn is_json_extension_type<T: AsRef<Field>>(field: T) -> bool {
-    field.as_ref().extension_type_name() == Some(JsonExtensionType::NAME)
+/// Checks whether this field is either a legacy JSONB or JSON2 extension type.
+pub fn is_any_json_extension_type<T: AsRef<Field>>(field: T) -> bool {
+    let name = field.as_ref().extension_type_name();
+    name == Some(JsonExtensionType::NAME) || name == Some(Json2ExtensionType::NAME)
+}
+
+/// Parses JSON2 settings stored by the historical `greptime.json` extension.
+pub fn parse_legacy_json2_settings(
+    metadata: &HashMap<String, String>,
+) -> crate::error::Result<Option<JsonSettings>> {
+    #[derive(Deserialize)]
+    struct LegacyJsonMetadata {
+        #[serde(default)]
+        json_settings: Option<JsonSettings>,
+    }
+
+    if metadata.get(EXTENSION_TYPE_NAME_KEY).map(String::as_str) != Some(JsonExtensionType::NAME) {
+        return Ok(None);
+    }
+
+    metadata
+        .get(EXTENSION_TYPE_METADATA_KEY)
+        .map(|json| {
+            serde_json::from_str::<LegacyJsonMetadata>(json)
+                .map(|x| x.json_settings)
+                .context(crate::error::DeserializeSnafu { json })
+        })
+        .transpose()
+        .map(Option::flatten)
+}
+
+/// Checks whether this field uses the JSON2 extension layout from before type hints.
+///
+/// That layout used the same `greptime.json` extension name and
+/// `json_structure_settings` metadata as legacy JSONB. Its structured Arrow data type is
+/// therefore required to distinguish JSON2 from Binary JSONB.
+pub fn is_legacy_json2_extension_type<T: AsRef<Field>>(field: T) -> bool {
+    let field = field.as_ref();
+    if field.extension_type_name() != Some(JsonExtensionType::NAME)
+        || !matches!(field.data_type(), DataType::Struct(_))
+    {
+        return false;
+    }
+
+    field
+        .metadata()
+        .get(EXTENSION_TYPE_METADATA_KEY)
+        .and_then(|json| serde_json::from_str::<serde_json::Value>(json).ok())
+        .is_some_and(|metadata| metadata.get(LEGACY_JSON_STRUCTURE_SETTINGS_KEY).is_some())
 }
 
 /// Check if this field is a JSON2 extension type.
 ///
-/// Legacy JSONB and JSON2 share the same JSON extension name. The column schema construction
-/// invariant is that JSON2 always stores its settings as `Some`, including default settings,
-/// while legacy JSONB stores no JSON settings. Therefore, after checking the extension name,
-/// the presence of JSON settings distinguishes JSON2 from legacy JSONB.
+/// New schemas use [`Json2ExtensionType`]. For compatibility, old fields using
+/// [`JsonExtensionType`] with JSON settings or the pre-type-hint structured layout are also
+/// recognized as JSON2.
 pub fn is_json2_extension_type<T: AsRef<Field>>(field: T) -> bool {
     let field = field.as_ref();
-    is_json_extension_type(field)
-        && field
-            .try_extension_type::<JsonExtensionType>()
-            .is_ok_and(|x| x.metadata().json_settings.is_some())
+    field.extension_type_name() == Some(Json2ExtensionType::NAME)
+        || parse_legacy_json2_settings(field.metadata()).is_ok_and(|x| x.is_some())
+        || is_legacy_json2_extension_type(field)
 }
 
-/// Check if this field is a structured JSON field.
-///
-/// Legacy JSONB columns may carry JSON extension metadata due to old metadata versions, but their
-/// physical Arrow type is still Binary. They must not enter structured JSON alignment paths.
-pub fn is_structured_json_field(field: &FieldRef) -> bool {
-    is_json_extension_type(field) && matches!(field.data_type(), DataType::Struct(_))
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+
+    use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY};
+    use arrow_schema::{Field, Fields};
+
+    use super::*;
+
+    #[test]
+    fn test_json2_extension_type_detection() {
+        let extension = Json2ExtensionType::new(Arc::new(JsonMetadata::default()));
+        let json2 = Field::new("j", DataType::Struct(Fields::empty()), true)
+            .with_extension_type(extension.clone());
+        // "projection" is the special hack for selecting the whole column of json2
+        let projection = Field::new("j", DataType::Binary, true).with_extension_type(extension);
+        let legacy_json2 = Field::new("j", DataType::Struct(Fields::empty()), true).with_metadata(
+            HashMap::from([
+                (
+                    EXTENSION_TYPE_NAME_KEY.to_string(),
+                    JsonExtensionType::NAME.to_string(),
+                ),
+                (
+                    EXTENSION_TYPE_METADATA_KEY.to_string(),
+                    serde_json::json!({ "json_settings": JsonSettings::default() }).to_string(),
+                ),
+            ]),
+        );
+        // Before type hints, JSON2 and JSONB shared extension metadata and were distinguished by
+        // their physical Arrow data types.
+        let legacy_structure_metadata = HashMap::from([
+            (
+                EXTENSION_TYPE_NAME_KEY.to_string(),
+                JsonExtensionType::NAME.to_string(),
+            ),
+            (
+                EXTENSION_TYPE_METADATA_KEY.to_string(),
+                serde_json::json!({
+                    (LEGACY_JSON_STRUCTURE_SETTINGS_KEY): { "Structured": null }
+                })
+                .to_string(),
+            ),
+        ]);
+        let pre_type_hint_json2 = Field::new("j", DataType::Struct(Fields::empty()), true)
+            .with_metadata(legacy_structure_metadata.clone());
+        let legacy_jsonb =
+            Field::new("j", DataType::Binary, true).with_metadata(legacy_structure_metadata);
+
+        assert!(is_json2_extension_type(&json2));
+        assert!(is_json2_extension_type(&projection));
+        assert!(is_json2_extension_type(&legacy_json2));
+        assert!(is_legacy_json2_extension_type(&pre_type_hint_json2));
+        assert!(is_json2_extension_type(&pre_type_hint_json2));
+        assert_eq!(
+            Some(JsonSettings::default()),
+            parse_legacy_json2_settings(legacy_json2.metadata()).unwrap()
+        );
+        assert!(!is_legacy_json2_extension_type(&legacy_jsonb));
+        assert!(!is_json2_extension_type(&legacy_jsonb));
+        assert!(JsonExtensionType::try_new(&DataType::Binary, ()).is_ok());
+        assert!(JsonExtensionType::try_new(&DataType::Null, ()).is_ok());
+        assert!(JsonExtensionType::try_new(&DataType::Struct(Fields::empty()), ()).is_err());
+    }
 }

src/datatypes/src/schema/column_schema.rs

@@ -486,7 +486,8 @@ impl ColumnSchema {
         }
     }
 
-    pub fn with_extension_type<E>(&mut self, extension_type: &E) -> Result<()>
+    /// Sets the Arrow extension type metadata for this column.
+    pub fn with_extension_type<E>(&mut self, extension_type: &E)
     where
         E: ExtensionType,
     {
@@ -496,9 +497,10 @@ impl ColumnSchema {
         if let Some(extension_metadata) = extension_type.serialize_metadata() {
             self.metadata
                 .insert(EXTENSION_TYPE_METADATA_KEY.to_string(), extension_metadata);
+        } else {
+            // Replacing an extension must not retain metadata owned by the previous type.
+            self.metadata.remove(EXTENSION_TYPE_METADATA_KEY);
         }
-
-        Ok(())
     }
 
     pub fn is_indexed(&self) -> bool {
@@ -1230,6 +1232,7 @@ mod tests {
     use arrow::datatypes::{DataType as ArrowDataType, TimeUnit};
 
     use super::*;
+    use crate::extension::json::{Json2ExtensionType, JsonExtensionType};
     use crate::types::{StructField, StructType};
     use crate::value::Value;
     use crate::vectors::Int32Vector;
@@ -1246,6 +1249,31 @@ mod tests {
         assert_eq!(column_schema, new_column_schema);
     }
 
+    #[test]
+    fn test_with_extension_type_replaces_metadata() {
+        let mut schema = ColumnSchema::new("j", ConcreteDataType::json_datatype(), true);
+
+        schema.with_extension_type(&Json2ExtensionType::default());
+        assert_eq!(
+            Some(Json2ExtensionType::NAME),
+            schema
+                .metadata()
+                .get(EXTENSION_TYPE_NAME_KEY)
+                .map(String::as_str)
+        );
+        assert!(schema.metadata().contains_key(EXTENSION_TYPE_METADATA_KEY));
+
+        schema.with_extension_type(&JsonExtensionType);
+        assert_eq!(
+            Some(JsonExtensionType::NAME),
+            schema
+                .metadata()
+                .get(EXTENSION_TYPE_NAME_KEY)
+                .map(String::as_str)
+        );
+        assert!(!schema.metadata().contains_key(EXTENSION_TYPE_METADATA_KEY));
+    }
+
     #[test]
     fn test_column_schema_with_default_constraint() {
         let column_schema = ColumnSchema::new("test", ConcreteDataType::int32_datatype(), true)

src/datatypes/src/schema/ext.rs

@@ -22,6 +22,6 @@ pub trait ArrowSchemaExt {
 
 impl ArrowSchemaExt for arrow_schema::Schema {
     fn has_json_extension_field(&self) -> bool {
-        self.fields().iter().any(json::is_json_extension_type)
+        self.fields().iter().any(json::is_any_json_extension_type)
     }
 }