From 95d9d92e42bfaa3338cf460bdb031803e29368bf Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:15:05 +0800 Subject: [PATCH] refactor: separate a json2 extension type (#8745) Signed-off-by: luofucong --- .../function/src/scalars/json/json_get.rs | 2 +- src/common/recordbatch/src/lib.rs | 6 +- src/common/recordbatch/src/recordbatch.rs | 5 +- src/common/sql/src/convert.rs | 21 +- src/datatypes/src/extension/json.rs | 208 +++++++++++++++--- src/datatypes/src/schema/column_schema.rs | 34 ++- src/datatypes/src/schema/ext.rs | 2 +- src/mito2/src/compaction/reader.rs | 6 +- src/mito2/src/flush.rs | 4 +- src/mito2/src/memtable/bulk.rs | 5 +- src/mito2/src/memtable/bulk/json_align.rs | 17 +- src/mito2/src/memtable/bulk/part.rs | 8 +- .../src/memtable/bulk/row_group_reader.rs | 4 +- src/mito2/src/read/compat.rs | 6 +- src/mito2/src/read/flat_projection.rs | 6 +- src/mito2/src/read/scan_region.rs | 4 +- .../src/sst/parquet/json_align/schema.rs | 10 +- .../src/sst/parquet/json_align/stream.rs | 10 +- src/mito2/src/sst/parquet/reader.rs | 8 +- src/mito2/src/sst/parquet/writer.rs | 11 +- src/mito2/src/test_util.rs | 21 +- .../src/etl/transform/transformer/greptime.rs | 28 ++- src/query/src/datafusion/json_expr_planner.rs | 15 +- src/query/src/dist_plan/merge_scan.rs | 47 +++- .../src/optimizer/json_type_concretize.rs | 6 +- src/query/src/sql/show_create_table.rs | 21 +- src/sql/src/error.rs | 12 +- src/sql/src/statements.rs | 16 +- .../common/types/json/json2_limit.result | 24 +- .../common/types/json/json2_limit.sql | 12 +- 30 files changed, 388 insertions(+), 191 deletions(-) diff --git a/src/common/function/src/scalars/json/json_get.rs b/src/common/function/src/scalars/json/json_get.rs index aabe6d08c8..ce49f9d59f 100644 --- a/src/common/function/src/scalars/json/json_get.rs +++ b/src/common/function/src/scalars/json/json_get.rs @@ -404,7 +404,6 @@ impl Function for JsonGetWithType { let result = match arg0.data_type() { DataType::Binary | DataType::LargeBinary | DataType::BinaryView => { let arg0 = compute::cast(&arg0, &DataType::BinaryView)?; - let jsons = arg0.as_binary_view(); if args.arg_fields.first().is_some_and(is_json2_extension_type) { // Query concretization projects nested JSON2 paths as Struct arrays. A binary @@ -414,6 +413,7 @@ impl Function for JsonGetWithType { .project_to(&with_type) .map_err(|e| exec_datafusion_err!("{e:?}"))? } else { + let jsons = arg0.as_binary_view(); let mut builder = result_builder(len, &with_type)?; jsonb_get(jsons, path, builder.as_mut())?; builder.build() diff --git a/src/common/recordbatch/src/lib.rs b/src/common/recordbatch/src/lib.rs index ed3bda72cc..72918e9e5e 100644 --- a/src/common/recordbatch/src/lib.rs +++ b/src/common/recordbatch/src/lib.rs @@ -45,7 +45,7 @@ use datatypes::arrow::util::display::{ ArrayFormatter, ArrayFormatterFactory, DisplayIndex, FormatOptions, FormatResult, }; use datatypes::arrow::util::pretty::pretty_format_batches_with_options; -use datatypes::extension::json::is_json_extension_type; +use datatypes::extension::json::is_any_json_extension_type; use datatypes::prelude::{ConcreteDataType, DataType, VectorRef}; use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; use datatypes::types::{JsonFormat, StructField, StructType, jsonb_to_string}; @@ -455,7 +455,7 @@ impl ArrayFormatterFactory for BinaryFormatterFactory { Ok(Some(ArrayFormatter::new( Box::new(BinaryFormatter { array, - is_json: field.is_some_and(is_json_extension_type), + is_json: field.is_some_and(is_any_json_extension_type), default: ArrayFormatter::try_new(array, options)?, null: options.null(), }), @@ -485,7 +485,7 @@ impl DisplayIndex for BinaryFormatter<'_> { ArrowDataType::Binary => self.array.as_binary::().value(idx), ArrowDataType::LargeBinary => self.array.as_binary::().value(idx), ArrowDataType::BinaryView => self.array.as_binary_view().value(idx), - _ => unreachable!(), + _ => return Ok(self.default.value(idx).write(f)?), }; let value = jsonb_to_string(bytes).map_err(|e| ArrowError::ExternalError(Box::new(e)))?; diff --git a/src/common/recordbatch/src/recordbatch.rs b/src/common/recordbatch/src/recordbatch.rs index 6257aabf82..48a2757489 100644 --- a/src/common/recordbatch/src/recordbatch.rs +++ b/src/common/recordbatch/src/recordbatch.rs @@ -378,7 +378,7 @@ mod tests { use datatypes::arrow::array::{AsArray, StringArray, StringViewArray, UInt32Array}; use datatypes::arrow::datatypes::{DataType, Field, Schema as ArrowSchema, UInt32Type}; use datatypes::data_type::ConcreteDataType; - use datatypes::extension::json::{JsonExtensionType, JsonMetadata}; + use datatypes::extension::json::JsonExtensionType; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::{BinaryVector, StringVector, UInt32Vector}; @@ -594,8 +594,7 @@ mod tests { #[test] fn test_legacy_json_with_extension_does_not_align_as_structured_json() { - let field = Field::new("j", DataType::Binary, true) - .with_extension_type(JsonExtensionType::new(Arc::new(JsonMetadata::default()))); + let field = Field::new("j", DataType::Binary, true).with_extension_type(JsonExtensionType); let arrow_schema = Arc::new(ArrowSchema::new(vec![field])); let schema = Arc::new(Schema::try_from(arrow_schema).unwrap()); let columns: Vec = vec![Arc::new(BinaryVector::from(vec![Some( diff --git a/src/common/sql/src/convert.rs b/src/common/sql/src/convert.rs index 3a10c5476a..fdd90ebe3d 100644 --- a/src/common/sql/src/convert.rs +++ b/src/common/sql/src/convert.rs @@ -17,7 +17,7 @@ use std::str::FromStr; use arrow_schema::extension::ExtensionType; use common_time::Timestamp; use common_time::timezone::Timezone; -use datatypes::extension::json::JsonExtensionType; +use datatypes::extension::json::{Json2ExtensionType, parse_legacy_json2_settings}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema}; use datatypes::types::{JsonFormat, parse_string_to_jsonb, parse_string_to_vector_type_value}; @@ -307,13 +307,20 @@ pub(crate) fn parse_string_to_value( Ok(Value::Binary(v.into())) } JsonFormat::Json2(_) => { - let extension_type: Option = - column_schema.extension_type().context(DatatypeSnafu)?; - let json_settings = extension_type - .and_then(|x| x.metadata().json_settings.clone()) - .unwrap_or_default(); let v = serde_json::from_str(&s).context(DeserializeSnafu { json: s })?; - json_settings.encode(v).context(DatatypeSnafu) + + if let Some(extension) = column_schema + .extension_type::() + .context(DatatypeSnafu)? + { + extension.metadata().json_settings().encode(v) + } else { + parse_legacy_json2_settings(column_schema.metadata()) + .context(DatatypeSnafu)? + .unwrap_or_default() + .encode(v) + } + .context(DatatypeSnafu) } }, ConcreteDataType::Vector(d) => { diff --git a/src/datatypes/src/extension/json.rs b/src/datatypes/src/extension/json.rs index b78a9df49b..26d470641f 100644 --- a/src/datatypes/src/extension/json.rs +++ b/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, + /// JSON2 settings stored in Arrow extension metadata. + json_settings: JsonSettings, } -#[derive(Debug, Clone)] -pub struct JsonExtensionType(Arc); +impl JsonMetadata { + /// Creates JSON2 extension metadata. + pub fn new(json_settings: JsonSettings) -> Self { + Self { json_settings } + } -impl JsonExtensionType { - pub fn new(metadata: Arc) -> 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 { + None + } + + fn deserialize_metadata(_metadata: Option<&str>) -> Result { + 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.supports_data_type(data_type).map(|_| Self) + } +} + +/// Arrow extension type for JSON2 columns and concretized projections. +#[derive(Debug, Clone, Default)] +pub struct Json2ExtensionType(Arc); + +impl Json2ExtensionType { + /// Creates a JSON2 extension type with the given metadata. + pub fn new(metadata: Arc) -> Self { + Self(metadata) + } +} + +impl ExtensionType for Json2ExtensionType { + const NAME: &'static str = "greptime.json2"; type Metadata = Arc; 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>(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>(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, +) -> crate::error::Result> { + #[derive(Deserialize)] + struct LegacyJsonMetadata { + #[serde(default)] + json_settings: Option, + } + + 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::(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>(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::(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>(field: T) -> bool { let field = field.as_ref(); - is_json_extension_type(field) - && field - .try_extension_type::() - .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()); + } } diff --git a/src/datatypes/src/schema/column_schema.rs b/src/datatypes/src/schema/column_schema.rs index d3ecd90116..6efab380e0 100644 --- a/src/datatypes/src/schema/column_schema.rs +++ b/src/datatypes/src/schema/column_schema.rs @@ -486,7 +486,8 @@ impl ColumnSchema { } } - pub fn with_extension_type(&mut self, extension_type: &E) -> Result<()> + /// Sets the Arrow extension type metadata for this column. + pub fn with_extension_type(&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) diff --git a/src/datatypes/src/schema/ext.rs b/src/datatypes/src/schema/ext.rs index 16ffcc6901..9da93f6b47 100644 --- a/src/datatypes/src/schema/ext.rs +++ b/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) } } diff --git a/src/mito2/src/compaction/reader.rs b/src/mito2/src/compaction/reader.rs index 04055ed9c1..89b70f89c0 100644 --- a/src/mito2/src/compaction/reader.rs +++ b/src/mito2/src/compaction/reader.rs @@ -20,7 +20,7 @@ use common_time::range::TimestampRange; use common_time::timestamp::TimeUnit; use datafusion_common::ScalarValue; use datafusion_expr::Expr; -use datatypes::extension::json::is_structured_json_field; +use datatypes::extension::json::is_json2_extension_type; use datatypes::types::json_type::JsonNativeType; use parquet::arrow::parquet_to_arrow_schema; use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData}; @@ -84,11 +84,11 @@ impl CompactionSstReaderBuilder<'_> { (row_group.num_rows() as u64, uncompressed_bytes) }), ); - let json_type_hint = if schema.fields().iter().any(is_structured_json_field) { + let json_type_hint = if schema.fields().iter().any(is_json2_extension_type) { let mut json_type_hint = schema .fields() .iter() - .filter(|&field| is_structured_json_field(field)) + .filter(|&field| is_json2_extension_type(field)) .map(|field| (field.name().clone(), JsonNativeType::Null)) .collect::>(); diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index a317f73fdc..a4a42970e4 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -24,7 +24,7 @@ use bytes::Bytes; use common_base::cancellation::CancellableFuture; use common_telemetry::{debug, error, info}; use datatypes::arrow::datatypes::SchemaRef; -use datatypes::extension::json::is_structured_json_field; +use datatypes::extension::json::is_json2_extension_type; use partition::expr::PartitionExpr; use smallvec::{SmallVec, smallvec}; use snafu::ResultExt; @@ -951,7 +951,7 @@ fn memtable_flat_sources( let mut input_iters = Vec::with_capacity(num_ranges); let mut current_ranges = Vec::new(); - let has_json2 = schema.fields().iter().any(is_structured_json_field); + let has_json2 = schema.fields().iter().any(is_json2_extension_type); let mut json_align_schemas = if has_json2 { Some(Vec::with_capacity(num_ranges)) } else { diff --git a/src/mito2/src/memtable/bulk.rs b/src/mito2/src/memtable/bulk.rs index dc9e97524a..596db7aa34 100644 --- a/src/mito2/src/memtable/bulk.rs +++ b/src/mito2/src/memtable/bulk.rs @@ -1533,7 +1533,7 @@ mod tests { use api::v1::value::ValueData; use api::v1::{Mutation, Row, Rows, SemanticType}; use datatypes::data_type::ConcreteDataType; - use datatypes::extension::json::{JsonExtensionType, JsonMetadata}; + use datatypes::extension::json::Json2ExtensionType; use datatypes::json::value::JsonValue; use datatypes::schema::ColumnSchema; use datatypes::types::json_type::{JsonNativeType, JsonObjectType}; @@ -1730,8 +1730,7 @@ mod tests { let data_type = ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())); let mut col_schema = ColumnSchema::new("data", data_type, true); - let extension = JsonExtensionType::new(Arc::new(JsonMetadata::default())); - col_schema.with_extension_type(&extension).unwrap(); + col_schema.with_extension_type(&Json2ExtensionType::default()); let col_meta_2 = ColumnMetadata { column_schema: col_schema, diff --git a/src/mito2/src/memtable/bulk/json_align.rs b/src/mito2/src/memtable/bulk/json_align.rs index 030de4cb5b..3a5c6e7fbc 100644 --- a/src/mito2/src/memtable/bulk/json_align.rs +++ b/src/mito2/src/memtable/bulk/json_align.rs @@ -17,7 +17,7 @@ use std::sync::Arc; use datatypes::arrow::datatypes::{DataType as ArrowDataType, Schema, SchemaRef}; use datatypes::arrow::record_batch::RecordBatch; -use datatypes::extension::json::is_structured_json_field; +use datatypes::extension::json::is_json2_extension_type; use datatypes::types::json_type::JsonNativeType; use datatypes::vectors::json::array::JsonArray; use snafu::{OptionExt, ResultExt}; @@ -60,7 +60,7 @@ impl Json2Aligner { .fields() .iter() .enumerate() - .filter(|&(_idx, field)| is_structured_json_field(field)) + .filter(|&(_idx, field)| is_json2_extension_type(field)) .map(|(idx, field)| { let json_type = JsonNativeType::try_from(field.data_type()).context(DataTypeMismatchSnafu)?; @@ -169,8 +169,8 @@ fn assert_columns_match_except_json2(base_schema: &Schema, schema: &Schema) { "input schemas for Json2Aligner must have the same column count" ); for (idx, (base_field, field)) in base_schema.fields().iter().zip(schema.fields()).enumerate() { - let base_is_json2 = is_structured_json_field(base_field); - let is_json2 = is_structured_json_field(field); + let base_is_json2 = is_json2_extension_type(base_field); + let is_json2 = is_json2_extension_type(field); debug_assert_eq!( base_is_json2, is_json2, "column {idx} must be JSON2 in all input schemas or none" @@ -192,7 +192,7 @@ mod tests { Array, ArrayRef, AsArray, Int64Array, StringViewArray, StructArray, UInt64Array, }; use datatypes::arrow::datatypes::{DataType, Field, Fields, Schema}; - use datatypes::extension::json::{JsonExtensionType, JsonMetadata}; + use datatypes::extension::json::{Json2ExtensionType, JsonExtensionType}; use serde_json::json; use super::*; @@ -234,8 +234,7 @@ mod tests { #[test] fn test_try_new_ignores_legacy_jsonb_extension_field() { let legacy_jsonb_field = Arc::new( - Field::new("data", DataType::Binary, true) - .with_extension_type(JsonExtensionType::new(Arc::new(JsonMetadata::default()))), + Field::new("data", DataType::Binary, true).with_extension_type(JsonExtensionType), ); let schema = Arc::new(Schema::new(vec![ Arc::new(Field::new("ts", DataType::Int64, false)), @@ -266,7 +265,7 @@ mod tests { assert_eq!(&DataType::Int64, fields[0].data_type()); assert_eq!("name", fields[1].name()); assert_eq!(&DataType::Utf8View, fields[1].data_type()); - assert!(is_structured_json_field(&aligner.schema().fields()[1])); + assert!(is_json2_extension_type(&aligner.schema().fields()[1])); } #[test] @@ -427,7 +426,7 @@ mod tests { fn json_field(name: &str, fields: Fields) -> Arc { Arc::new( Field::new(name, DataType::Struct(fields), true) - .with_extension_type(JsonExtensionType::new(Arc::new(JsonMetadata::default()))), + .with_extension_type(Json2ExtensionType::default()), ) } diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index a328ff8ab1..fcba451947 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -37,7 +37,7 @@ use datatypes::arrow::datatypes::{ DataType as ArrowDataType, Field, Schema, SchemaRef, UInt32Type, }; use datatypes::data_type::DataType; -use datatypes::extension::json::is_structured_json_field; +use datatypes::extension::json::is_json2_extension_type; use datatypes::prelude::{MutableVector, Vector}; use datatypes::value::ValueRef; use datatypes::vectors::Helper; @@ -446,7 +446,7 @@ impl UnorderedPart { // Get the schema from the first part let schema = self.parts[0].batch.schema(); - let concatenated = if schema.fields().iter().any(is_structured_json_field) { + let concatenated = if schema.fields().iter().any(is_json2_extension_type) { let aligner = Json2Aligner::try_new(self.parts.iter().map(|part| part.batch.schema()))?; let aligned_batches = aligner.align_batches(self.parts.iter().map(|part| part.batch.clone()))?; @@ -722,13 +722,13 @@ impl BulkPartConverter { } fn align_schema_with_json_array(schema: SchemaRef, columns: &[ArrayRef]) -> SchemaRef { - if schema.fields().iter().all(|f| !is_structured_json_field(f)) { + if schema.fields().iter().all(|f| !is_json2_extension_type(f)) { return schema; } let mut fields = Vec::with_capacity(schema.fields().len()); for (field, array) in schema.fields().iter().zip(columns) { - if !is_structured_json_field(field) { + if !is_json2_extension_type(field) { fields.push(field.clone()); continue; } diff --git a/src/mito2/src/memtable/bulk/row_group_reader.rs b/src/mito2/src/memtable/bulk/row_group_reader.rs index cab11dd674..0788485bcd 100644 --- a/src/mito2/src/memtable/bulk/row_group_reader.rs +++ b/src/mito2/src/memtable/bulk/row_group_reader.rs @@ -15,7 +15,7 @@ use std::sync::Arc; use bytes::Bytes; -use datatypes::extension::json::is_structured_json_field; +use datatypes::extension::json::is_json2_extension_type; use parquet::arrow::ProjectionMask; use parquet::arrow::arrow_reader::{ ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReader, @@ -57,7 +57,7 @@ impl MemtableRowGroupReaderBuilder { .arrow_schema() .fields() .iter() - .any(is_structured_json_field) + .any(is_json2_extension_type) { arrow_reader_options = arrow_reader_options.with_schema(context.read_format().arrow_schema().clone()); diff --git a/src/mito2/src/read/compat.rs b/src/mito2/src/read/compat.rs index 561a34a52e..af0bdc58c5 100644 --- a/src/mito2/src/read/compat.rs +++ b/src/mito2/src/read/compat.rs @@ -25,7 +25,7 @@ use datatypes::arrow::compute::{TakeOptions, take}; use datatypes::arrow::datatypes::{FieldRef, Schema, SchemaRef}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::data_type::ConcreteDataType; -use datatypes::extension::json::is_structured_json_field; +use datatypes::extension::json::is_json2_extension_type; use datatypes::prelude::DataType; use datatypes::value::Value; use datatypes::vectors::VectorRef; @@ -103,10 +103,10 @@ impl FlatCompatBatch { .arrow_schema() .fields() .iter() - .any(is_structured_json_field) + .any(is_json2_extension_type) { for field in read_format.arrow_schema().fields() { - if is_structured_json_field(field) + if is_json2_extension_type(field) && let Some(column_id) = actual.column_by_name(field.name()).map(|x| x.column_id) && let Some(i) = actual_schema.iter().position(|x| x.0 == column_id) diff --git a/src/mito2/src/read/flat_projection.rs b/src/mito2/src/read/flat_projection.rs index ed5e0fa36a..fbd6088f67 100644 --- a/src/mito2/src/read/flat_projection.rs +++ b/src/mito2/src/read/flat_projection.rs @@ -25,7 +25,7 @@ use common_recordbatch::error::{ use common_recordbatch::{DfRecordBatch, RecordBatch}; use datatypes::arrow::array::Array; use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field}; -use datatypes::extension::json::{is_json_extension_type, is_structured_json_field}; +use datatypes::extension::json::is_json2_extension_type; use datatypes::prelude::{ConcreteDataType, DataType}; use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; use datatypes::types::JsonType; @@ -292,7 +292,7 @@ impl FlatProjectionMapper { .input_arrow_schema .fields() .iter() - .filter(|&field| is_structured_json_field(field)) + .filter(|&field| is_json2_extension_type(field)) .map(|field| (field.name().clone(), field.data_type().clone())) .collect(); to_flat_sst_arrow_schema(&self.metadata, &options) @@ -358,7 +358,7 @@ impl FlatProjectionMapper { } let field = &self.output_schema.arrow_schema().fields()[output_idx]; - if is_json_extension_type(field) { + if is_json2_extension_type(field) { array = JsonArray::from(&array) .project_to(field.data_type()) .context(DataTypesSnafu)?; diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index cc6d7b8d98..482c2d4a33 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -34,7 +34,7 @@ use datafusion_common::{Column, ScalarValue}; use datafusion_expr::Expr; use datafusion_expr::utils::expr_to_columns; use datatypes::arrow::array::{ArrayRef, BooleanArray, UInt64Array}; -use datatypes::extension::json::is_structured_json_field; +use datatypes::extension::json::is_json2_extension_type; use datatypes::types::json_type::JsonNativeType; use datatypes::value::timestamp_to_scalar_value; use futures::StreamExt; @@ -431,7 +431,7 @@ impl ScanRegion { .arrow_schema() .fields() .iter() - .any(is_structured_json_field); + .any(is_json2_extension_type); let read_cols = if has_structured_json { self.read_columns_with_json_type_hint(&read_col_ids) } else { diff --git a/src/mito2/src/sst/parquet/json_align/schema.rs b/src/mito2/src/sst/parquet/json_align/schema.rs index 81a04c0067..86bd4d4d0d 100644 --- a/src/mito2/src/sst/parquet/json_align/schema.rs +++ b/src/mito2/src/sst/parquet/json_align/schema.rs @@ -17,7 +17,7 @@ use std::sync::Arc; use arrow_schema::{DataType as ArrowDataType, FieldRef}; use datatypes::arrow::datatypes::Schema; -use datatypes::extension::json::is_structured_json_field; +use datatypes::extension::json::is_json2_extension_type; use store_api::storage::NestedPath; /// Aligns nested struct fields according to the requested nested paths. @@ -25,7 +25,7 @@ use store_api::storage::NestedPath; /// For each root field: /// - An empty path list keeps the whole field unchanged. /// - Non-JSON root fields ignore nested paths and keep the whole field unchanged. -/// - Structured JSON root fields are rebuilt from `nested_paths`. +/// - JSON2 root fields are rebuilt from `nested_paths`. /// - Existing schema fields are preserved only when they are the requested leaf. /// - Requested paths missing from the schema are synthesized with JSONB (`Binary`) /// leaves. @@ -42,7 +42,7 @@ where .into_iter() .zip(nested_paths) .map(|(field, paths)| { - if !paths.is_empty() && is_structured_json_field(field) { + if !paths.is_empty() && is_json2_extension_type(field) { let child_paths = paths .iter() .map(|path| { @@ -131,7 +131,7 @@ fn new_jsonb_field(name: &str) -> FieldRef { #[cfg(test)] mod tests { use arrow_schema::Field; - use datatypes::extension::json::{JsonExtensionType, JsonMetadata}; + use datatypes::extension::json::Json2ExtensionType; use super::*; @@ -152,7 +152,7 @@ mod tests { ArrowDataType::Struct(fields.into_iter().collect()), true, ) - .with_extension_type(JsonExtensionType::new(Arc::new(JsonMetadata::default()))), + .with_extension_type(Json2ExtensionType::default()), ) } diff --git a/src/mito2/src/sst/parquet/json_align/stream.rs b/src/mito2/src/sst/parquet/json_align/stream.rs index 62d3c67ec0..fabb3c807b 100644 --- a/src/mito2/src/sst/parquet/json_align/stream.rs +++ b/src/mito2/src/sst/parquet/json_align/stream.rs @@ -20,7 +20,7 @@ use datafusion_common::format::DEFAULT_CAST_OPTIONS; use datatypes::arrow::array::{ArrayRef, new_null_array}; use datatypes::arrow::datatypes::{DataType, FieldRef, SchemaRef}; use datatypes::arrow::record_batch::RecordBatch; -use datatypes::extension::json::is_structured_json_field; +use datatypes::extension::json::is_json2_extension_type; use datatypes::vectors::json::array::JsonArray; use futures::Stream; use snafu::{ResultExt, ensure}; @@ -173,7 +173,7 @@ fn align_array(array: &ArrayRef, field: &FieldRef) -> Result { return Ok(array.clone()); } - if is_structured_json_field(field) { + if is_json2_extension_type(field) { return JsonArray::from(array) .project_to(field.data_type()) .context(DataTypeMismatchSnafu); @@ -194,7 +194,7 @@ mod tests { Array, ArrayRef, BinaryArray, Int64Array, StringArray, StringViewArray, StructArray, }; use datatypes::arrow::datatypes::{DataType, Field, Fields, Schema}; - use datatypes::extension::json::{JsonExtensionType, JsonMetadata}; + use datatypes::extension::json::Json2ExtensionType; use datatypes::types::parse_string_to_jsonb; use futures::{StreamExt, stream}; @@ -388,7 +388,7 @@ mod tests { ))])), true, ) - .with_extension_type(JsonExtensionType::new(Arc::new(JsonMetadata::default())))]); + .with_extension_type(Json2ExtensionType::default())]); let mut aligner = NestedSchemaAligner::new(stream::iter([Ok(input)]), vec![true], output_schema.clone()) .unwrap(); @@ -477,7 +477,7 @@ mod tests { ))])), true, ) - .with_extension_type(JsonExtensionType::new(Arc::new(JsonMetadata::default())))]); + .with_extension_type(Json2ExtensionType::default())]); let mut aligner = NestedSchemaAligner::new(stream::iter([Ok(input)]), vec![true], output_schema.clone()) .unwrap(); diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index d70dd50195..b4f0accf9f 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -31,7 +31,7 @@ use datatypes::arrow::array::ArrayRef; use datatypes::arrow::datatypes::{Field, Schema as ArrowSchema, SchemaRef}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::data_type::ConcreteDataType; -use datatypes::extension::json::is_structured_json_field; +use datatypes::extension::json::is_json2_extension_type; use datatypes::prelude::DataType; use futures::StreamExt; use mito_codec::row_converter::build_primary_key_codec; @@ -541,7 +541,7 @@ impl ParquetReaderBuilder { .arrow_schema() .fields() .iter() - .any(is_structured_json_field) + .any(is_json2_extension_type) { // Read `__primary_key` as Binary when it's too large for dictionary // encoding; convert_batch wraps it back to a DictionaryArray. @@ -2443,7 +2443,7 @@ mod tests { use datatypes::arrow::array::{ArrayRef, Int64Array, StringArray, StructArray}; use datatypes::arrow::datatypes::{Fields, Schema}; use datatypes::arrow::record_batch::RecordBatch; - use datatypes::extension::json::{JsonExtensionType, JsonMetadata}; + use datatypes::extension::json::Json2ExtensionType; use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnSchema; use object_store::services::Memory; @@ -2520,7 +2520,7 @@ mod tests { let b_field = Arc::new(Field::new("b", DataType::Utf8, true)); let json_fields = Fields::from(vec![a_field, b_field]); let json_field = Field::new("j", DataType::Struct(json_fields.clone()), true) - .with_extension_type(JsonExtensionType::new(Arc::new(JsonMetadata::default()))); + .with_extension_type(Json2ExtensionType::default()); let schema = Arc::new(Schema::new(vec![json_field])); let a_array = Arc::new(StructArray::new( diff --git a/src/mito2/src/sst/parquet/writer.rs b/src/mito2/src/sst/parquet/writer.rs index a98c2856d5..a6a4cd5cd0 100644 --- a/src/mito2/src/sst/parquet/writer.rs +++ b/src/mito2/src/sst/parquet/writer.rs @@ -32,7 +32,7 @@ use datatypes::arrow::array::{ use datatypes::arrow::compute::{max, min}; use datatypes::arrow::datatypes::{DataType, SchemaRef, TimeUnit}; use datatypes::arrow::record_batch::RecordBatch; -use datatypes::extension::json::is_structured_json_field; +use datatypes::extension::json::is_json2_extension_type; use object_store::{FuturesAsyncWriter, ObjectStore}; use parquet::arrow::AsyncArrowWriter; use parquet::basic::{Compression, Encoding, ZstdLevel}; @@ -278,17 +278,12 @@ where ) -> Result { let mut options = FlatSchemaOptions::from_encoding(self.metadata.primary_key_encoding); - if source - .schema() - .fields() - .iter() - .any(is_structured_json_field) - { + if source.schema().fields().iter().any(is_json2_extension_type) { options.concretized_json_types = source .schema() .fields() .iter() - .filter(|&field| is_structured_json_field(field)) + .filter(|&field| is_json2_extension_type(field)) .map(|field| (field.name().clone(), field.data_type().clone())) .collect::>(); } diff --git a/src/mito2/src/test_util.rs b/src/mito2/src/test_util.rs index 5e02788498..8a63d749b5 100644 --- a/src/mito2/src/test_util.rs +++ b/src/mito2/src/test_util.rs @@ -33,7 +33,6 @@ use api::v1::column_def::options_from_column_schema; use api::v1::helper::row; use api::v1::value::ValueData; use api::v1::{OpType, Row, Rows, SemanticType}; -use arrow_schema::extension::{EXTENSION_TYPE_NAME_KEY, ExtensionType}; use common_base::Plugins; use common_base::readable_size::ReadableSize; use common_datasource::compression::CompressionType; @@ -45,7 +44,7 @@ use common_telemetry::{debug, warn}; use common_test_util::temp_dir::{TempDir, create_temp_dir}; use common_wal::options::{KafkaWalOptions, WAL_OPTIONS_KEY, WalOptions}; use datatypes::arrow::array::{TimestampMillisecondArray, UInt8Array, UInt64Array}; -use datatypes::extension::json::JsonExtensionType; +use datatypes::extension::json::{Json2ExtensionType, JsonExtensionType}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnSchema; use log_store::kafka::log_store::KafkaLogStore; @@ -866,11 +865,10 @@ impl CreateRequestBuilder { for i in 0..self.field_num { let mut column_schema = ColumnSchema::new(format!("field_{i}"), self.field_datatype.clone(), nullable); - if self.field_datatype.is_json() { - column_schema.mut_metadata().insert( - EXTENSION_TYPE_NAME_KEY.to_string(), - JsonExtensionType::NAME.to_string(), - ); + if self.field_datatype.is_json2() { + column_schema.with_extension_type(&Json2ExtensionType::default()); + } else if self.field_datatype.is_json() { + column_schema.with_extension_type(&JsonExtensionType); } column_metadatas.push(ColumnMetadata { column_schema, @@ -931,11 +929,10 @@ impl CreateRequestBuilder { for i in 0..self.field_num { let mut column_schema = ColumnSchema::new(format!("field_{i}"), self.field_datatype.clone(), nullable); - if self.field_datatype.is_json() { - column_schema.mut_metadata().insert( - EXTENSION_TYPE_NAME_KEY.to_string(), - JsonExtensionType::NAME.to_string(), - ); + if self.field_datatype.is_json2() { + column_schema.with_extension_type(&Json2ExtensionType::default()); + } else if self.field_datatype.is_json() { + column_schema.with_extension_type(&JsonExtensionType); } column_metadatas.push(ColumnMetadata { column_schema, diff --git a/src/pipeline/src/etl/transform/transformer/greptime.rs b/src/pipeline/src/etl/transform/transformer/greptime.rs index 814e3b36c0..3292de749a 100644 --- a/src/pipeline/src/etl/transform/transformer/greptime.rs +++ b/src/pipeline/src/etl/transform/transformer/greptime.rs @@ -28,7 +28,8 @@ use coerce::{coerce_columns, coerce_value}; use common_query::prelude::{greptime_timestamp, greptime_value}; use common_telemetry::warn; use datatypes::data_type::ConcreteDataType; -use datatypes::extension::json::JsonExtensionType; +use datatypes::extension::json::{Json2ExtensionType, parse_legacy_json2_settings}; +use datatypes::json::JsonSettings; use datatypes::value::Value; use greptime_proto::v1::{ColumnSchema, Row, Rows, Value as GreptimeValue}; use itertools::Itertools; @@ -701,19 +702,24 @@ fn resolve_value( }); let value = if is_json2 { - let json_extension_type: Option = - if let Some(x) = schema_info.find_column_schema_in_table(&column_name) { - x.column_schema.extension_type()? - } else { - None - }; - let settings = json_extension_type - .and_then(|x| x.metadata().json_settings.clone()) - .unwrap_or_default(); let value: serde_json::Value = value.try_into().map_err(|e: StdError| { CoerceIncompatibleTypesSnafu { msg: e.to_string() }.build() })?; - let value = settings.encode(value)?; + let value = + if let Some(column) = schema_info.find_column_schema_in_table(&column_name) { + if let Some(extension) = column + .column_schema + .extension_type::()? + { + extension.metadata().json_settings().encode(value)? + } else { + parse_legacy_json2_settings(column.column_schema.metadata())? + .unwrap_or_default() + .encode(value)? + } + } else { + JsonSettings::default().encode(value)? + }; resolve_schema( index, diff --git a/src/query/src/datafusion/json_expr_planner.rs b/src/query/src/datafusion/json_expr_planner.rs index 786561db2b..e650ac102c 100644 --- a/src/query/src/datafusion/json_expr_planner.rs +++ b/src/query/src/datafusion/json_expr_planner.rs @@ -15,7 +15,6 @@ use std::sync::{Arc, LazyLock}; use arrow_schema::Field; -use arrow_schema::extension::ExtensionType; use common_function::scalars::json::json_get::JsonGetWithType; use common_function::scalars::udf::create_udf; use datafusion_common::arrow::datatypes::DataType; @@ -23,7 +22,7 @@ use datafusion_common::{Column, DFSchema, Result, ScalarValue, TableReference}; use datafusion_expr::expr::{BinaryExpr, ScalarFunction}; use datafusion_expr::planner::{ExprPlanner, PlannerResult, RawBinaryExpr}; use datafusion_expr::{Expr, ExprSchemable, Operator, ScalarUDF}; -use datatypes::extension::json::JsonExtensionType; +use datatypes::extension::json::is_json2_extension_type; use either::Either; use sqlparser::ast::BinaryOperator; @@ -83,7 +82,7 @@ impl ExprPlanner for JsonExprPlanner { qualifier: Option<&TableReference>, nested_names: &[String], ) -> Result>> { - if field.extension_type_name() != Some(JsonExtensionType::NAME) { + if !is_json2_extension_type(field) { return Ok(PlannerResult::Original(Vec::new())); } @@ -154,7 +153,7 @@ fn parse_sql_op(op: &BinaryOperator) -> Option { #[cfg(test)] mod tests { use arrow_schema::Fields; - use datatypes::extension::json::JsonMetadata; + use datatypes::extension::json::Json2ExtensionType; use super::*; @@ -169,11 +168,6 @@ mod tests { )) } - fn json_field(name: &str) -> Field { - Field::new(name, DataType::Binary, true) - .with_extension_type(JsonExtensionType::new(Arc::new(JsonMetadata::default()))) - } - #[test] fn test_plan_binary_op() -> Result<()> { let planner = JsonExprPlanner; @@ -249,7 +243,8 @@ mod tests { let nested_names = vec!["payload".to_string(), "cpu".to_string()]; let planned = planner.plan_compound_identifier( - &json_field("labels"), + &Field::new("labels", DataType::Struct(Fields::empty()), true) + .with_extension_type(Json2ExtensionType::default()), Some(&qualifier), &nested_names, )?; diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index 0de66bb069..0245f98ddc 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -46,7 +46,10 @@ use datafusion_common::{Column as ColumnExpr, DataFusionError, Result}; use datafusion_expr::{Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore}; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr::{Distribution, EquivalenceProperties, PhysicalSortExpr}; -use datatypes::extension::json::is_json_extension_type; +use datatypes::extension::json::{ + Json2ExtensionType, is_any_json_extension_type, is_json2_extension_type, + is_legacy_json2_extension_type, +}; use futures_util::StreamExt; use greptime_proto::v1::region::RegionRequestHeader; use meter_core::data::ReadItem; @@ -122,7 +125,7 @@ fn merge_scan_schema_error_count_for_test() -> u64 { /// parts of the field and must match. fn json_fields_compatible(expected_field: &Field, actual_field: &Field) -> bool { let is_json = |field: &Field| { - is_json_extension_type(field) + is_any_json_extension_type(field) || field .metadata() .get(datatypes::schema::TYPE_KEY) @@ -913,11 +916,18 @@ fn maybe_amend_json2_field(schema: &ArrowSchema) -> ArrowSchemaRef { let schema = schema.clone(); let mut new_fields = Vec::with_capacity(schema.fields().len()); for field in schema.fields().iter() { - let new_field = if is_json_extension_type(field) + let new_field = if is_json2_extension_type(field) && matches!(field.data_type(), DataType::Struct(fields) if fields.is_empty()) { + let is_legacy_json2 = is_legacy_json2_extension_type(field); let mut new_field = field.as_ref().clone(); new_field.set_data_type(DataType::Binary); + if is_legacy_json2 { + // Pre-type-hint JSON2 is identified partly by its Struct data type. Promote the + // ephemeral field before rewriting it to Binary so later checks retain its JSON2 + // identity. + new_field = new_field.with_extension_type(Json2ExtensionType::default()); + } Arc::new(new_field) } else { field.clone() @@ -1291,7 +1301,10 @@ mod tests { use std::task::{Context, Poll}; use arrow::array::{Int64Array, TimestampMillisecondArray}; - use arrow_schema::{DataType as TestArrowDataType, Field, TimeUnit}; + use arrow_schema::extension::{ + EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType, + }; + use arrow_schema::{DataType as TestArrowDataType, Field, Fields, TimeUnit}; use async_trait::async_trait; use common_base::Plugins; use common_meta::peer::Peer; @@ -1311,6 +1324,7 @@ mod tests { use datafusion_physical_expr::expressions::{ Column, DynamicFilterPhysicalExpr, lit as physical_lit, }; + use datatypes::extension::json::JsonExtensionType; use datatypes::prelude::{ConcreteDataType, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::{Int64Vector, StringVector, TimestampMillisecondVector}; @@ -1340,6 +1354,31 @@ mod tests { QueryId::from(Uuid::from_u128(value)) } + #[test] + fn test_amend_legacy_json2_field_preserves_json2_identity() { + let field = Field::new("j", DataType::Struct(Fields::empty()), true).with_metadata( + StdHashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + JsonExtensionType::NAME.to_string(), + ), + ( + EXTENSION_TYPE_METADATA_KEY.to_string(), + serde_json::json!({ + "json_structure_settings": { "Structured": null } + }) + .to_string(), + ), + ]), + ); + + let schema = maybe_amend_json2_field(&ArrowSchema::new(vec![field])); + let field = schema.field(0); + assert_eq!(&DataType::Binary, field.data_type()); + assert_eq!(Some(Json2ExtensionType::NAME), field.extension_type_name()); + assert!(is_json2_extension_type(field)); + } + fn merge_scan_exec_with_sorted_input( region_count: u64, target_partition: usize, diff --git a/src/query/src/optimizer/json_type_concretize.rs b/src/query/src/optimizer/json_type_concretize.rs index 2687402583..4210433eb3 100644 --- a/src/query/src/optimizer/json_type_concretize.rs +++ b/src/query/src/optimizer/json_type_concretize.rs @@ -156,7 +156,7 @@ mod tests { use datafusion_expr::expr::ScalarFunction; use datafusion_expr::{LogicalPlanBuilder, col, lit}; use datafusion_optimizer::OptimizerContext; - use datatypes::extension::json::{JsonExtensionType, JsonMetadata}; + use datatypes::extension::json::Json2ExtensionType; use datatypes::schema::ColumnSchema; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; use store_api::storage::{ConcreteDataType, RegionId}; @@ -196,9 +196,7 @@ mod tests { ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())), true, ); - json_column - .with_extension_type(&JsonExtensionType::new(Arc::new(JsonMetadata::default()))) - .unwrap(); + json_column.with_extension_type(&Json2ExtensionType::default()); builder .push_column_metadata(ColumnMetadata { column_schema: json_column, diff --git a/src/query/src/sql/show_create_table.rs b/src/query/src/sql/show_create_table.rs index 1c37558d4b..0c66ae8738 100644 --- a/src/query/src/sql/show_create_table.rs +++ b/src/query/src/sql/show_create_table.rs @@ -18,7 +18,7 @@ use std::collections::HashMap; use arrow_schema::extension::ExtensionType; use common_meta::SchemaOptions; -use datatypes::extension::json::JsonExtensionType; +use datatypes::extension::json::{Json2ExtensionType, parse_legacy_json2_settings}; use datatypes::schema::{ COLUMN_FULLTEXT_OPT_KEY_ANALYZER, COLUMN_FULLTEXT_OPT_KEY_BACKEND, COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE, COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE, @@ -211,12 +211,12 @@ fn create_column(column_schema: &ColumnSchema, quote_style: char) -> Result()? { - let settings = json_extension - .metadata() - .json_settings - .clone() - .unwrap_or_default(); + let settings = if let Some(extension) = column_schema.extension_type::()? { + Some(extension.metadata().json_settings().clone()) + } else { + parse_legacy_json2_settings(column_schema.metadata())? + }; + if let Some(settings) = settings { extensions.set_json_settings(settings).context(SqlSnafu)?; } @@ -321,6 +321,7 @@ mod tests { use std::time::Duration; use common_time::timestamp::TimeUnit; + use datatypes::extension::json::JsonExtensionType; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ FulltextOptions, Schema, SchemaRef, SkippingIndexOptions, VectorIndexOptions, @@ -433,11 +434,7 @@ WITH( #[test] fn test_show_create_legacy_json_with_json_extension() { let mut json_column = ColumnSchema::new("j", ConcreteDataType::json_datatype(), true); - json_column - .with_extension_type(&JsonExtensionType::new(Arc::new( - datatypes::extension::json::JsonMetadata::default(), - ))) - .unwrap(); + json_column.with_extension_type(&JsonExtensionType); let table_schema = SchemaRef::new(Schema::new(vec![ json_column, diff --git a/src/sql/src/error.rs b/src/sql/src/error.rs index e4323bb661..cc0b9f1950 100644 --- a/src/sql/src/error.rs +++ b/src/sql/src/error.rs @@ -339,14 +339,6 @@ pub enum Error { #[snafu(implicit)] location: Location, }, - - #[snafu(display("Failed to set JSON structure settings: {value}"))] - SetJsonSettings { - value: String, - source: datatypes::error::Error, - #[snafu(implicit)] - location: Location, - }, } impl ErrorExt for Error { @@ -392,9 +384,7 @@ impl ErrorExt for Error { #[cfg(feature = "enterprise")] InvalidTriggerWebhookOption { .. } => StatusCode::InvalidArguments, - SerializeColumnDefaultConstraint { source, .. } | SetJsonSettings { source, .. } => { - source.status_code() - } + SerializeColumnDefaultConstraint { source, .. } => source.status_code(), ConvertToGrpcDataType { source, .. } => source.status_code(), SqlCommon { source, .. } => source.status_code(), diff --git a/src/sql/src/statements.rs b/src/sql/src/statements.rs index 3b6f6ad4f9..aafe62d2e6 100644 --- a/src/sql/src/statements.rs +++ b/src/sql/src/statements.rs @@ -39,7 +39,7 @@ use api::helper::ColumnDataTypeWrapper; use api::v1::SemanticType; use common_sql::default_constraint::parse_column_default_constraint; use common_time::timezone::Timezone; -use datatypes::extension::json::{JsonExtensionType, JsonMetadata}; +use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{COMMENT_KEY, ColumnDefaultConstraint, ColumnSchema}; use datatypes::types::json_type::JsonNativeType; @@ -54,8 +54,8 @@ use crate::ast::{ }; use crate::error::{ self, ConvertToGrpcDataTypeSnafu, ConvertValueSnafu, Result, - SerializeColumnDefaultConstraintSnafu, SetFulltextOptionSnafu, SetJsonSettingsSnafu, - SetSkippingIndexOptionSnafu, SetVectorIndexOptionSnafu, SqlCommonSnafu, + SerializeColumnDefaultConstraintSnafu, SetFulltextOptionSnafu, SetSkippingIndexOptionSnafu, + SetVectorIndexOptionSnafu, SqlCommonSnafu, }; use crate::statements::create::Column; pub use crate::statements::option_map::OptionMap; @@ -164,14 +164,8 @@ pub fn column_to_schema( }; if is_json2_column { let settings = column.extensions.build_json_settings()?.unwrap_or_default(); - let extension = JsonExtensionType::new(Arc::new(JsonMetadata { - json_settings: Some(settings.clone()), - })); - column_schema - .with_extension_type(&extension) - .with_context(|_| SetJsonSettingsSnafu { - value: format!("{settings:?}"), - })?; + let extension = Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings))); + column_schema.with_extension_type(&extension); } Ok(column_schema) diff --git a/tests/cases/standalone/common/types/json/json2_limit.result b/tests/cases/standalone/common/types/json/json2_limit.result index 838717ba2d..3cb462a8a7 100644 --- a/tests/cases/standalone/common/types/json/json2_limit.result +++ b/tests/cases/standalone/common/types/json/json2_limit.result @@ -36,7 +36,7 @@ drop table json2_disable_non_object_insert; Affected Rows: 0 -create table json2_disable_whole_column_read ( +create table json2_whole_and_path_read ( ts timestamp time index, j json2 ) @@ -46,7 +46,7 @@ with ( Affected Rows: 0 -insert into json2_disable_whole_column_read values +insert into json2_whole_and_path_read values (1, '{"a": {"b": 1}}'), (2, '{"a": {"b": 2}}'); @@ -54,26 +54,26 @@ Affected Rows: 2 -- JSON2 field projection remains supported (case 5): use in an intermediate plan node. select json_get(j, 'a.b'), count(*) -from json2_disable_whole_column_read +from json2_whole_and_path_read group by json_get(j, 'a.b') order by json_get(j, 'a.b'); -+---------------------------------------------------------+----------+ -| json_get(json2_disable_whole_column_read.j,Utf8("a.b")) | count(*) | -+---------------------------------------------------------+----------+ -| 1 | 1 | -| 2 | 1 | -+---------------------------------------------------------+----------+ ++---------------------------------------------------+----------+ +| json_get(json2_whole_and_path_read.j,Utf8("a.b")) | count(*) | ++---------------------------------------------------+----------+ +| 1 | 1 | +| 2 | 1 | ++---------------------------------------------------+----------+ -select j, j.a from json2_disable_whole_column_read; +select j, j.a from json2_whole_and_path_read; Error: 3001(EngineExecuteQuery), Invalid argument error: column types must match schema types, expected Binary but found Struct("a": Utf8View) at column index 0 -select j from json2_disable_whole_column_read where j.a.b = 1; +select j from json2_whole_and_path_read where j.a.b = 1; Error: 3001(EngineExecuteQuery), Invalid argument error: column types must match schema types, expected Binary but found Struct("a": Struct("b": Int64)) at column index 0 -drop table json2_disable_whole_column_read; +drop table json2_whole_and_path_read; Affected Rows: 0 diff --git a/tests/cases/standalone/common/types/json/json2_limit.sql b/tests/cases/standalone/common/types/json/json2_limit.sql index 42fac25362..4917b9690e 100644 --- a/tests/cases/standalone/common/types/json/json2_limit.sql +++ b/tests/cases/standalone/common/types/json/json2_limit.sql @@ -20,7 +20,7 @@ insert into json2_disable_non_object_insert values (6, '{}'); drop table json2_disable_non_object_insert; -create table json2_disable_whole_column_read ( +create table json2_whole_and_path_read ( ts timestamp time index, j json2 ) @@ -28,21 +28,21 @@ with ( 'append_mode' = 'true' ); -insert into json2_disable_whole_column_read values +insert into json2_whole_and_path_read values (1, '{"a": {"b": 1}}'), (2, '{"a": {"b": 2}}'); -- JSON2 field projection remains supported (case 5): use in an intermediate plan node. select json_get(j, 'a.b'), count(*) -from json2_disable_whole_column_read +from json2_whole_and_path_read group by json_get(j, 'a.b') order by json_get(j, 'a.b'); -select j, j.a from json2_disable_whole_column_read; +select j, j.a from json2_whole_and_path_read; -select j from json2_disable_whole_column_read where j.a.b = 1; +select j from json2_whole_and_path_read where j.a.b = 1; -drop table json2_disable_whole_column_read; +drop table json2_whole_and_path_read; create table json2_without_append_mode ( ts timestamp time index,