diff --git a/Cargo.lock b/Cargo.lock index 94428d8b13..ff2ce93d77 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10283,6 +10283,9 @@ dependencies = [ "num-integer", "num-traits", "object_store", + "parquet-variant", + "parquet-variant-compute", + "parquet-variant-json", "paste", "seq-macro", "simdutf8", diff --git a/src/datatypes/src/extension/json.rs b/src/datatypes/src/extension/json.rs index ab4a138420..b912a0a3f4 100644 --- a/src/datatypes/src/extension/json.rs +++ b/src/datatypes/src/extension/json.rs @@ -129,27 +129,32 @@ pub struct JsonMetadata { } impl JsonMetadata { - /// Creates metadata for the legacy JSON2 layout. + /// Creates metadata for the JSON2 layout version 2. pub fn new(json_settings: JsonSettings) -> Self { - Self { - json_settings, - layout_version: None, - } - } - - /// Creates metadata for the new JSON2 physical v2 layout. - pub fn new_v2(json_settings: JsonSettings) -> Self { Self { json_settings, layout_version: Some(JSON2_LAYOUT_V2), } } + /// Creates metadata for the legacy JSON2 layout. + pub fn new_v1(json_settings: JsonSettings) -> Self { + Self { + json_settings, + layout_version: None, + } + } + /// Returns the JSON2 settings. pub fn json_settings(&self) -> &JsonSettings { &self.json_settings } + /// Consumes the metadata and returns its JSON2 settings. + pub fn into_json_settings(self) -> JsonSettings { + self.json_settings + } + /// Returns whether this metadata describes JSON2 layout version 2. pub fn is_version_2(&self) -> bool { self.layout_version == Some(JSON2_LAYOUT_V2) @@ -226,7 +231,7 @@ impl ExtensionType for Json2ExtensionType { })?; Ok(Arc::new(metadata)) } else { - Ok(Arc::new(JsonMetadata::default())) + Ok(Arc::new(JsonMetadata::new_v1(JsonSettings::default()))) } } @@ -406,7 +411,7 @@ mod tests { let legacy: JsonMetadata = serde_json::from_str(r#"{"json_settings":{}}"#)?; assert!(!legacy.is_version_2()); - let metadata = JsonMetadata::new_v2(JsonSettings::default()); + let metadata = JsonMetadata::new(JsonSettings::default()); assert!(metadata.is_version_2()); let serialized = serde_json::to_string(&metadata)?; let deserialized: JsonMetadata = serde_json::from_str(&serialized)?; @@ -418,7 +423,9 @@ mod tests { #[test] fn test_parse_json2_physical_layout() -> crate::error::Result<()> { let legacy = Field::new("data", DataType::Struct(Fields::empty()), true) - .with_extension_type(Json2ExtensionType::default()); + .with_extension_type(Json2ExtensionType::new(Arc::new(JsonMetadata::new_v1( + JsonSettings::default(), + )))); assert!(!Json2PhysicalLayout::try_from_root(&legacy)?.is_version_2()); let v2 = Field::new( @@ -432,7 +439,7 @@ mod tests { ), true, ) - .with_extension_type(Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2( + .with_extension_type(Json2ExtensionType::new(Arc::new(JsonMetadata::new( JsonSettings::default(), )))); assert!(Json2PhysicalLayout::try_from_root(&v2)?.is_version_2()); @@ -447,7 +454,7 @@ mod tests { assert!(Json2PhysicalLayout::try_from_root(&field).is_err()); let metadata = - Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(JsonSettings::default()))); + Json2ExtensionType::new(Arc::new(JsonMetadata::new(JsonSettings::default()))); let missing = Field::new("data", DataType::Struct(Fields::empty()), true) .with_extension_type(metadata.clone()); assert!(json2_remainder_field(&missing).is_ok_and(|x| x.is_none())); diff --git a/src/datatypes/src/json.rs b/src/datatypes/src/json.rs index 1ef8c85a22..0afe563734 100644 --- a/src/datatypes/src/json.rs +++ b/src/datatypes/src/json.rs @@ -39,6 +39,8 @@ use crate::value::{ListValue, StructValue, Value}; pub const JSON2_MAX_STRUCTURED_DEPTH: usize = 50; /// Reserved physical field containing unexpanded JSON2 paths. pub const JSON2_REMAINDER_FIELD_NAME: &str = "!__remainder__!"; +/// Default maximum number of unhinted JSON leaf paths expanded into Arrow fields. +pub const JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS: u32 = 100; /// JSON2 settings stored in column schema metadata and represented through /// Arrow extension metadata. @@ -105,6 +107,14 @@ pub struct JsonContext<'a> { } impl JsonSettings { + /// Creates default v2 settings for newly created JSON2 columns. + pub fn new_v2() -> Self { + Self { + type_hints: vec![], + max_auto_expanded_paths: Some(JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS), + } + } + /// Creates and validates JSON2 settings. pub fn try_new( type_hints: Vec, @@ -159,6 +169,10 @@ impl JsonSettings { } fn validate_type_hints(type_hints: &[JsonTypeHint]) -> Result<()> { + if type_hints.is_empty() { + return Ok(()); + } + let mut object = JsonObjectType::new(); for hint in type_hints { if hint.path.len() > JSON2_MAX_STRUCTURED_DEPTH { @@ -193,7 +207,9 @@ fn validate_type_hints(type_hints: &[JsonTypeHint]) -> Result<()> { | ConcreteDataType::Int64(_) | ConcreteDataType::Float32(_) | ConcreteDataType::Float64(_) - | ConcreteDataType::String(_) => (&hint.data_type).into(), + | ConcreteDataType::String(_) + | ConcreteDataType::List(_) + | ConcreteDataType::Struct(_) => (&hint.data_type).into(), data_type => { return InvalidJson2SettingsSnafu { reason: format!("unsupported JSON2 type hint data type: {data_type}"), diff --git a/src/datatypes/src/json/value.rs b/src/datatypes/src/json/value.rs index 728aadeb28..67bfff4bc6 100644 --- a/src/datatypes/src/json/value.rs +++ b/src/datatypes/src/json/value.rs @@ -30,6 +30,8 @@ use crate::types::json_type::{JsonNativeType, JsonNumberType, is_include}; use crate::types::{StructField, StructType}; use crate::value::{ListValue, StructValue, Value}; +pub type JsonObjectVariant = BTreeMap; + /// Number in json, can be a positive integer, a negative integer, or a floating number. /// Each of which is represented as `u64`, `i64` and `f64`. /// @@ -124,7 +126,7 @@ pub enum JsonVariant { Number(JsonNumber), String(String), Array(Vec), - Object(BTreeMap), + Object(JsonObjectVariant), /// A special "variant" value of JSON, to represent a union result of conflict JSON type values. Variant(Vec), } diff --git a/src/datatypes/src/schema/column_schema.rs b/src/datatypes/src/schema/column_schema.rs index 6efab380e0..a71605afdb 100644 --- a/src/datatypes/src/schema/column_schema.rs +++ b/src/datatypes/src/schema/column_schema.rs @@ -28,10 +28,12 @@ use crate::data_type::{ConcreteDataType, DataType}; use crate::error::{ self, ArrowMetadataSnafu, Error, InvalidFulltextOptionSnafu, ParseExtendedTypeSnafu, Result, }; +use crate::extension::json::Json2ExtensionType; use crate::schema::TYPE_KEY; use crate::schema::constraint::ColumnDefaultConstraint; use crate::value::Value; -use crate::vectors::VectorRef; +use crate::vectors::json::builder::JsonVectorBuilder; +use crate::vectors::{MutableVector, VectorRef}; pub type Metadata = HashMap; @@ -128,6 +130,21 @@ impl ColumnSchema { } } + /// Creates a mutable vector using this column's extension metadata. + pub fn create_mutable_vector(&self, capacity: usize) -> Box { + if self.data_type.is_json2() + && let Some(extension) = self.extension_type::().ok().flatten() + && extension.metadata().is_version_2() + { + Box::new(JsonVectorBuilder::with_settings( + extension.metadata().json_settings(), + capacity, + )) + } else { + self.data_type.create_mutable_vector(capacity) + } + } + #[inline] pub fn is_time_index(&self) -> bool { self.is_time_index diff --git a/src/datatypes/src/types/json_type.rs b/src/datatypes/src/types/json_type.rs index 7905fce76b..bf00236aef 100644 --- a/src/datatypes/src/types/json_type.rs +++ b/src/datatypes/src/types/json_type.rs @@ -80,7 +80,7 @@ impl JsonNativeType { Self::Number(JsonNumberType::F64) } - fn object() -> Self { + pub fn object() -> Self { Self::Object(JsonObjectType::new()) } @@ -143,6 +143,14 @@ impl JsonNativeType { JsonNativeType::Variant => ArrowDataType::Binary, } } + + /// Returns whether this type is a boolean, number, or string scalar. + pub fn is_primitive(&self) -> bool { + matches!( + self, + JsonNativeType::Bool | JsonNativeType::Number(_) | JsonNativeType::String + ) + } } impl From<&ConcreteDataType> for JsonNativeType { @@ -360,6 +368,7 @@ impl DataType for JsonType { fn create_mutable_vector(&self, capacity: usize) -> Box { match &self.format { JsonFormat::Jsonb => Box::new(BinaryVectorBuilder::with_capacity(capacity)), + // TODO(LFC): Carry JsonSettings in JsonFormat::Json2 and use with_settings here. JsonFormat::Json2(x) => Box::new(JsonVectorBuilder::new(x.as_ref().clone(), capacity)), } } diff --git a/src/datatypes/src/vectors/helper.rs b/src/datatypes/src/vectors/helper.rs index b2689eee2c..685b07419f 100644 --- a/src/datatypes/src/vectors/helper.rs +++ b/src/datatypes/src/vectors/helper.rs @@ -87,22 +87,6 @@ impl Helper { }) } - pub fn check_get_mutable_vector( - vector: &mut dyn MutableVector, - ) -> Result<&mut T> { - let ty = vector.data_type(); - vector - .as_mut_any() - .downcast_mut() - .with_context(|| error::UnknownVectorSnafu { - msg: format!( - "downcast vector error, vector type: {:?}, expected vector: {:?}", - ty, - std::any::type_name::(), - ), - }) - } - pub fn check_get_scalar_vector( vector: &VectorRef, ) -> Result<&::VectorType> { diff --git a/src/datatypes/src/vectors/json.rs b/src/datatypes/src/vectors/json.rs index c68e6cf3d9..958a98d9eb 100644 --- a/src/datatypes/src/vectors/json.rs +++ b/src/datatypes/src/vectors/json.rs @@ -15,3 +15,5 @@ pub mod array; pub(crate) mod builder; pub mod variant; + +pub use builder::json2_physical_data_type; diff --git a/src/datatypes/src/vectors/json/array.rs b/src/datatypes/src/vectors/json/array.rs index 0426334c3d..153d6b490e 100644 --- a/src/datatypes/src/vectors/json/array.rs +++ b/src/datatypes/src/vectors/json/array.rs @@ -33,9 +33,12 @@ use crate::error::{ AlignJsonArraySnafu, ArrowComputeSnafu, InvalidJsonSnafu, InvalidJsonbSnafu, Result, }; use crate::extension::json::{JSON2_REMAINDER_FIELD_NAME, json2_remainder_field}; +use crate::json::JsonSettings; use crate::json::value::{decode_json_variant, encode_serde_json_as_jsonb}; use crate::prelude::{DataType as _, Value as GreptimeValue}; use crate::value::{ListValue, StructValue}; +use crate::vectors::MutableVector; +use crate::vectors::json::builder::{JsonVectorBuilder, json2_physical_data_type}; use crate::vectors::json::variant::variant_to_json_values; pub struct JsonArray<'a> { @@ -115,6 +118,38 @@ impl JsonArray<'_> { } } + /// Rewrites a JSON2 array from the current physical layout into the specified + /// v2 physical layout. + pub fn rewrite_to_v2( + &self, + field: &Field, + logical_settings: &JsonSettings, + target_layout: &JsonSettings, + ) -> Result { + let is_v2 = json2_remainder_field(field)?.is_some(); + if is_v2 && self.inner.data_type() == &json2_physical_data_type(target_layout) { + return Ok(self.inner.clone()); + } + + let values = if is_v2 { + self.json2_values()? + } else { + (0..self.inner.len()) + .map(|i| self.try_get_value(i)) + .collect::>>()? + }; + let mut builder = JsonVectorBuilder::with_settings(target_layout, values.len()); + for value in values { + if value.is_null() { + builder.push_null(); + } else { + let value = logical_settings.encode(value)?; + builder.try_push_value_ref(&value.as_value_ref())?; + } + } + Ok(builder.to_vector().to_arrow_array()) + } + fn json2_values(&self) -> Result> { let structs = self.inner.as_struct_opt().context(AlignJsonArraySnafu { reason: "JSON2 layout v2 root array must be a struct", @@ -149,7 +184,14 @@ impl JsonArray<'_> { if child.name() == JSON2_REMAINDER_FIELD_NAME { continue; } - let value = JsonArray::from(column).try_get_value(i)?; + let mut value = JsonArray::from(column).try_get_value(i)?; + // Arrow child nulls cannot distinguish a missing path from an explicit JSON + // null. Builders preserve explicit null presence in the remainder, so nulls + // from the explicit branch must be discarded before merging both branches. + remove_null_object_fields(&mut value); + if value.is_null() { + continue; + } merge_explicit_value(&mut object, child.name().clone(), value, &mut path)?; } values.push(Value::Object(object)); @@ -420,6 +462,16 @@ fn merge_explicit_value( Ok(()) } +fn remove_null_object_fields(value: &mut Value) { + let Value::Object(object) = value else { + return; + }; + object.retain(|_, value| { + remove_null_object_fields(value); + !value.is_null() + }); +} + /// Returns whether Arrow can cast between the types without JSON-aware projection. /// Binary and nested types require JSONB decoding or recursive projection. fn can_fast_cast_types(from_type: &DataType, to_type: &DataType) -> bool { @@ -545,6 +597,8 @@ impl<'a> From<&'a ArrayRef> for JsonArray<'a> { #[cfg(test)] mod test { + use std::sync::Arc; + use arrow_array::types::Int64Type; use arrow_array::{ BinaryArray, BooleanArray, Float32Array, Float64Array, Int8Array, Int16Array, Int32Array, @@ -555,7 +609,7 @@ mod test { use super::*; use crate::extension::json::{Json2ExtensionType, JsonMetadata}; - use crate::json::JsonSettings; + use crate::json::{JsonSettings, JsonTypeHint}; use crate::vectors::json::variant::{json_values_to_variant, variant_field}; #[test] @@ -1049,7 +1103,7 @@ mod test { None, )); let field = Field::new("data", DataType::Struct(fields), true).with_extension_type( - Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(JsonSettings::default()))), + Json2ExtensionType::new(Arc::new(JsonMetadata::new(JsonSettings::default()))), ); assert_eq!( @@ -1063,12 +1117,10 @@ mod test { assert_eq!( json!({ "!__remainder__!": "user value", - "count": null, - "nested": {"left": null} + "nested": {} }), JsonArray::from(&array).json2_values()?[1] ); - let target = DataType::Struct( vec![ Arc::new(Field::new("cold", DataType::UInt64, true)), @@ -1088,6 +1140,38 @@ mod test { Ok(()) } + #[test] + fn test_rewrite_to_v2_reuses_matching_layout() -> Result<()> { + let settings = JsonSettings::try_new( + vec![JsonTypeHint { + path: vec!["kind".to_string()], + data_type: ConcreteDataType::string_datatype(), + nullable: true, + default_constraint: None, + inverted_index: false, + }], + Some(0), + )?; + let value = settings.encode(json!({"kind": "access", "cold": 1}))?; + let mut builder = JsonVectorBuilder::with_settings(&settings, 1); + builder.try_push_value_ref(&value.as_value_ref())?; + let array = builder.to_vector().to_arrow_array(); + let structs = array.as_struct(); + assert!(structs.column_by_name("kind").is_some()); + assert_eq!( + vec![Some(json!({"cold": 1}))], + variant_to_json_values(structs.column_by_name(JSON2_REMAINDER_FIELD_NAME).unwrap())? + ); + let field = Field::new("data", array.data_type().clone(), true).with_extension_type( + Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings.clone()))), + ); + + let rewritten = JsonArray::from(&array).rewrite_to_v2(&field, &settings, &settings)?; + + assert!(Arc::ptr_eq(&array, &rewritten)); + Ok(()) + } + #[test] fn test_project_partial_json2_v2_without_remainder() -> Result<()> { let fields = Fields::from(vec![Arc::new(Field::new("hot", DataType::Int64, true))]); @@ -1097,7 +1181,7 @@ mod test { None, )); let field = Field::new("data", DataType::Struct(fields), true).with_extension_type( - Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(JsonSettings::default()))), + Json2ExtensionType::new(Arc::new(JsonMetadata::new(JsonSettings::default()))), ); let projected = JsonArray::from(&array).project_to_v2(&field, field.data_type())?; @@ -1124,6 +1208,18 @@ mod test { ) ); + let Value::Object(mut remainder) = json!({"count": 1}) else { + unreachable!(); + }; + let error = merge_explicit_value( + &mut remainder, + "count".to_string(), + json!(1), + &mut Vec::new(), + ) + .unwrap_err(); + assert!(error.to_string().contains("cannot merge 'count'")); + let Value::Object(mut remainder) = json!({"nested": {"count": 1}}) else { unreachable!(); }; diff --git a/src/datatypes/src/vectors/json/builder.rs b/src/datatypes/src/vectors/json/builder.rs index ab3fd6d043..2e97c82730 100644 --- a/src/datatypes/src/vectors/json/builder.rs +++ b/src/datatypes/src/vectors/json/builder.rs @@ -29,11 +29,13 @@ use crate::error::{ }; use crate::extension::json::JSON2_REMAINDER_FIELD_NAME; use crate::json::value::{JsonNumber, JsonVariant, JsonVariantRef, encode_json_variant}; -use crate::json::{JSON2_MAX_STRUCTURED_DEPTH, JsonSettings}; +use crate::json::{ + JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS, JSON2_MAX_STRUCTURED_DEPTH, JsonSettings, +}; use crate::prelude::{ValueRef, Vector, VectorRef}; use crate::types::StructType; use crate::types::json_type::{JsonNativeType, is_include}; -use crate::value::{ListValue, StructValue, StructValueRef, Value}; +use crate::value::{ListValue, ListValueRef, StructValue, StructValueRef, Value}; use crate::vectors::json::variant::{append_json_variant, append_json_variant_ref, variant_field}; use crate::vectors::{Helper, MutableVector, NullVector, StructVectorBuilder}; @@ -216,6 +218,31 @@ impl JsonVectorBuilderState { } } +/// Returns the fixed v2 Arrow physical type produced from `settings`. +pub fn json2_physical_data_type(settings: &JsonSettings) -> DataType { + let DataType::Struct(fields) = explicit_type(settings).as_arrow_type() else { + unreachable!("JSON2 explicit type must map to Arrow Struct") + }; + let mut fields = fields + .iter() + .cloned() + .chain(std::iter::once(Arc::new(variant_field( + JSON2_REMAINDER_FIELD_NAME, + true, + )))) + .collect::>(); + fields.sort_unstable_by(|x, y| x.name().cmp(y.name())); + DataType::Struct(fields.into()) +} + +fn explicit_type(settings: &JsonSettings) -> JsonNativeType { + let mut explicit_type = JsonNativeType::Object(Default::default()); + for hint in settings.type_hints() { + insert_dynamic_type(&mut explicit_type, &hint.path, (&hint.data_type).into()); + } + explicit_type +} + impl JsonVectorBuilder { /// Creates a builder that merges all observed paths into the explicit schema. pub(crate) fn new(initial_native_type: JsonNativeType, capacity: usize) -> Self { @@ -232,12 +259,8 @@ impl JsonVectorBuilder { } /// Creates a builder bounded by the JSON settings and their type hints. - #[cfg_attr(not(test), expect(dead_code))] pub(crate) fn with_settings(settings: &JsonSettings, capacity: usize) -> Self { - let mut explicit_type = JsonNativeType::Object(Default::default()); - for hint in settings.type_hints() { - insert_dynamic_type(&mut explicit_type, &hint.path, (&hint.data_type).into()); - } + let explicit_type = explicit_type(settings); let state = if settings.max_auto_expanded_paths() == Some(0) { let DataType::Struct(fields) = explicit_type.as_arrow_type() else { unreachable!("JSON2 explicit type must map to Arrow Struct") @@ -255,7 +278,9 @@ impl JsonVectorBuilder { } else { JsonVectorBuilderState::AutoExpanding { explicit_type, - max_auto_expanded_paths: settings.max_auto_expanded_paths().unwrap_or(u32::MAX), + max_auto_expanded_paths: settings + .max_auto_expanded_paths() + .unwrap_or(JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS), values: Vec::with_capacity(capacity), } }; @@ -618,6 +643,17 @@ fn json_variant_ref_into_value_ref<'a>( ValueRef::Float64(*x) } (JsonVariantRef::String(x), ConcreteDataType::String(_)) => ValueRef::String(x), + (JsonVariantRef::Array(array), ConcreteDataType::List(list_type)) => { + let item_type = list_type.item_type().clone(); + let values = array + .iter() + .map(|x| json_variant_ref_into_value_ref(x, &item_type)) + .collect::>>()?; + ValueRef::List(ListValueRef::RefList { + val: values, + item_datatype: Arc::new(item_type), + }) + } (value, expected_type) => { return TryFromValueSnafu { reason: format!("unable to convert json value {value:?} to {expected_type}"), @@ -640,9 +676,15 @@ fn remainder_ref<'a>( }; let mut remainder = BTreeMap::new(); for (&name, value) in object { + // Preserve explicit JSON nulls in the remainder because Arrow child nulls cannot + // distinguish a present JSON null from a missing path. + if *value == JsonVariantRef::Null { + remainder.insert(name, JsonVariantRef::Null); + continue; + } + match fields.get(name) { Some(data_type @ JsonNativeType::Object(_)) => match value { - JsonVariantRef::Null => {} JsonVariantRef::Object(object) => { let child = remainder_ref(object, data_type)?; if !child.is_empty() { @@ -689,11 +731,14 @@ fn split_to_explicit( let Some(value) = remainder.remove(name) else { continue; }; + if value == JsonVariant::Null { + explicit.insert(name.clone(), JsonVariant::Null); + // Preserve explicit JSON nulls in the remainder because Arrow child nulls cannot + // distinguish a present JSON null from a missing path. + remainder.insert(name.clone(), JsonVariant::Null); + continue; + } if matches!(data_type, JsonNativeType::Object(_)) { - if matches!(value, JsonVariant::Null) { - explicit.insert(name.clone(), value); - continue; - } let (child_explicit, child_remainder) = split_to_explicit(value, data_type)?; explicit.insert(name.clone(), JsonVariant::Object(child_explicit)); if !child_remainder.is_empty() { @@ -1042,6 +1087,7 @@ mod tests { builder.try_push_value_ref(&value.as_value_ref())?; } let array = builder.to_vector().to_arrow_array(); + assert_eq!(&json2_physical_data_type(&settings), array.data_type()); assert_eq!(0, builder.len()); let structs = array.as_struct(); assert_eq!( @@ -1055,13 +1101,13 @@ mod tests { assert_eq!( vec![ Some(json!({"commit": {"collection": "post"}, "extra": 1})), - Some(json!({"dynamic": true})), + Some(json!({"commit": {"operation": null}, "dynamic": true})), ], variant_to_json_values(structs.column_by_name(JSON2_REMAINDER_FIELD_NAME).unwrap())? ); let field = Field::new("data", array.data_type().clone(), true).with_extension_type( - Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(settings))), + Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings))), ); let reconstructed = JsonArray::from(&array).project_to_v2(&field, &DataType::Binary)?; let reconstructed = reconstructed.as_binary::(); @@ -1099,6 +1145,15 @@ mod tests { #[test] fn test_finite_budget_selects_dynamic_paths() -> Result<()> { + let builder = JsonVectorBuilder::with_settings(&JsonSettings::default(), 0); + assert!(matches!( + builder.state, + JsonVectorBuilderState::AutoExpanding { + max_auto_expanded_paths: JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS, + .. + } + )); + let settings = JsonSettings::try_new( vec![JsonTypeHint { path: vec!["hint".to_string()], @@ -1154,7 +1209,7 @@ mod tests { ); let field = Field::new("data", array.data_type().clone(), true).with_extension_type( - Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(settings))), + Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings))), ); let reconstructed = JsonArray::from(&array).project_to_v2(&field, &DataType::Binary)?; let reconstructed = reconstructed.as_binary::(); @@ -1169,9 +1224,7 @@ mod tests { assert_eq!( json!({ "hint": "third", - "popular": "scalar", - "tie_a": null, - "tie_b": null + "popular": "scalar" }), decode_json_variant(reconstructed.value(2)).unwrap() ); @@ -1179,6 +1232,46 @@ mod tests { Ok(()) } + #[test] + fn test_v2_builder_preserves_explicit_null_presence() -> Result<()> { + let settings = JsonSettings::try_new(vec![], Some(1))?; + let values = [json!({"value": 1}), json!({"value": null}), json!({})]; + let mut builder = JsonVectorBuilder::with_settings(&settings, values.len()); + for value in values.clone() { + let value = settings.encode(value)?; + builder.try_push_value_ref(&value.as_value_ref())?; + } + let array = builder.to_vector().to_arrow_array(); + let structs = array.as_struct(); + assert_eq!( + vec![ + Some(json!({})), + Some(json!({"value": null})), + Some(json!({})) + ], + variant_to_json_values(structs.column_by_name(JSON2_REMAINDER_FIELD_NAME).unwrap())? + ); + let field = Field::new("data", array.data_type().clone(), true).with_extension_type( + Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings))), + ); + let reconstructed = JsonArray::from(&array).project_to_v2(&field, &DataType::Binary)?; + let reconstructed = reconstructed.as_binary::(); + + assert_eq!( + values[0], + decode_json_variant(reconstructed.value(0)).unwrap() + ); + assert_eq!( + values[1], + decode_json_variant(reconstructed.value(1)).unwrap() + ); + assert_eq!( + values[2], + decode_json_variant(reconstructed.value(2)).unwrap() + ); + Ok(()) + } + #[test] fn test_reconstruct_nested_remainder_only_value() -> Result<()> { let settings = JsonSettings::try_new(vec![], Some(1))?; @@ -1195,16 +1288,12 @@ mod tests { let array = builder.to_vector().to_arrow_array(); let field = Field::new("data", array.data_type().clone(), true).with_extension_type( - Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(settings))), + Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings))), ); let reconstructed = JsonArray::from(&array).project_to_v2(&field, &DataType::Binary)?; let reconstructed = reconstructed.as_binary::(); assert_eq!( - vec![ - values[0].clone(), - values[1].clone(), - json!({"a": {"cold": 3, "hot": null}}), - ], + vec![values[0].clone(), values[1].clone(), values[2].clone()], (0..reconstructed.len()) .map(|i| decode_json_variant(reconstructed.value(i)).unwrap()) .collect::>() diff --git a/src/datatypes/src/vectors/json/variant.rs b/src/datatypes/src/vectors/json/variant.rs index 9f90f0aaac..79b3b56165 100644 --- a/src/datatypes/src/vectors/json/variant.rs +++ b/src/datatypes/src/vectors/json/variant.rs @@ -39,6 +39,10 @@ pub fn variant_field(name: impl Into, nullable: bool) -> Field { .with_extension_type(VariantType) } +/// Encodes JSON values as an unshredded Parquet Variant array. +/// +/// `None` represents an Arrow null while `Some(Value::Null)` represents a JSON +/// null, preserving the distinction required by JSON2. #[cfg(test)] pub(crate) fn json_values_to_variant(values: &[Option]) -> Result { let mut builder = VariantArrayBuilder::new(values.len()); @@ -53,7 +57,7 @@ pub(crate) fn json_values_to_variant(values: &[Option]) -> Re /// Encodes JSON variants as an unshredded Parquet Variant array. #[cfg(test)] -fn json_variants_to_variant(values: &[Option]) -> Result { +pub(crate) fn json_variants_to_variant(values: &[Option]) -> Result { let mut builder = VariantArrayBuilder::new(values.len()); for value in values { match value { @@ -201,6 +205,7 @@ fn append_large_u64( builder: &mut impl VariantBuilderExt, value: u64, ) -> std::result::Result<(), ArrowError> { + // Parquet Variant has no unsigned integer primitive. Decimal16 preserves the full u64 range. let value = VariantDecimal16::try_new(value as i128, 0).map_err(|e| { ArrowError::InvalidArgumentError(format!( "Failed to encode JSON large integer as Variant Decimal16: {e}" diff --git a/src/mito2/Cargo.toml b/src/mito2/Cargo.toml index e48a5758ea..bf8e1cdaec 100644 --- a/src/mito2/Cargo.toml +++ b/src/mito2/Cargo.toml @@ -65,7 +65,7 @@ log-store = { workspace = true } mito-codec.workspace = true moka = { workspace = true, features = ["sync", "future"] } object-store = { workspace = true, features = ["testing"] } -parquet = { workspace = true, features = ["async"] } +parquet = { workspace = true, features = ["async", "variant_experimental"] } paste.workspace = true pin-project.workspace = true prometheus.workspace = true diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index 642e7b1ab8..19f104ce1a 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -34,6 +34,7 @@ use common_datasource::compression::CompressionType; use common_telemetry::warn; use datatypes::arrow::buffer::BooleanBuffer; use datatypes::arrow::record_batch::RecordBatch; +use datatypes::types::json_type::JsonNativeType; use datatypes::value::Value; use datatypes::vectors::VectorRef; use index::bloom_filter_index::{BloomFilterIndexCache, BloomFilterIndexCacheRef}; @@ -2066,8 +2067,7 @@ impl SelectorResultValue { SelectorResult::Flat(batches) => batches.iter().map(record_batch_estimated_size).sum(), }; result_size - + self.json_target_types.len() - * (mem::size_of::() + mem::size_of::()) + + self.json_target_types.len() * (size_of::() + size_of::()) } } diff --git a/src/mito2/src/compaction.rs b/src/mito2/src/compaction.rs index 58ddb29dad..cc98c3083e 100644 --- a/src/mito2/src/compaction.rs +++ b/src/mito2/src/compaction.rs @@ -14,6 +14,7 @@ mod buckets; pub mod compactor; +mod json2; pub mod memory_manager; pub mod picker; mod reader; @@ -31,6 +32,9 @@ use common_meta::key::SchemaMetadataManagerRef; use common_telemetry::{debug, error}; use common_time::TimeToLive; use common_time::range::TimestampRange; +pub(crate) use json2::{ + Json2RewritePlans, collect_json2_rewrite_plans, rewrite_json2_batch, rewrite_json2_schema, +}; pub use scheduler::CompactionRequest; pub(crate) use scheduler::{ CompactionExecution, CompactionPickFinished, CompactionScheduler, CompactionTransition, diff --git a/src/mito2/src/compaction/json2.rs b/src/mito2/src/compaction/json2.rs new file mode 100644 index 0000000000..02782b6bf4 --- /dev/null +++ b/src/mito2/src/compaction/json2.rs @@ -0,0 +1,658 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use arrow_schema::extension::ExtensionType; +use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field, Schema, SchemaRef}; +use datatypes::arrow::record_batch::RecordBatch; +use datatypes::extension::json::{JSON2_REMAINDER_FIELD_NAME, Json2ExtensionType, JsonMetadata}; +use datatypes::json::{JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS, JsonSettings, JsonTypeHint}; +use datatypes::prelude::ConcreteDataType; +use datatypes::types::json_type::JsonNativeType; +use datatypes::vectors::json::array::JsonArray; +use datatypes::vectors::json::json2_physical_data_type; +use parquet::arrow::parquet_to_arrow_schema; +use parquet::file::metadata::ParquetMetaData; +use snafu::{OptionExt, ResultExt, ensure}; +use store_api::metadata::RegionMetadataRef; + +use crate::error::{ + ConvertValueSnafu, DataTypeMismatchSnafu, InvalidRecordBatchSnafu, NewRecordBatchSnafu, Result, +}; + +/// Plan for rewriting one JSON2 column to a fixed compaction layout. +/// +/// A compaction input may contain v1 and v2 SSTs with different physical schemas. This plan is +/// derived only from the current region metadata and remains fixed while all input batches are +/// decoded and rewritten. It therefore prevents source-only paths from expanding the output +/// schema without a bound. +pub(crate) struct Json2RewritePlan { + /// User-defined settings used to encode logical values. + logical_settings: JsonSettings, + /// Fixed settings used to build the target physical layout. + pub(super) target_layout: JsonSettings, +} + +/// JSON2 rewrite plans keyed by logical column name. +pub(crate) type Json2RewritePlans = HashMap; + +#[derive(Clone)] +struct Json2LeafPathStats { + rows: u64, + data_type: JsonNativeType, + is_type_conflicted: bool, +} + +/// Builds the JSON2 rewrite plans for a compaction. +/// +/// Type hints from current region metadata are always retained. Existing explicit dynamic paths +/// from all input SST schemas are ranked once to produce a fixed layout; paths found only in a v2 +/// remainder are deliberately not promoted. [`rewrite_json2_batch`] decodes inputs and rewrites +/// them according to these plans. Non-JSON2 columns are omitted from the returned map. +/// +/// Returns an error when a JSON2 column has invalid or missing extension metadata, or when its +/// output layout is not v2. Legacy region metadata is upgraded in memory by the region opener +/// before compaction reaches this function. +pub(super) fn collect_json2_rewrite_plans_from_parquet( + metadata: &RegionMetadataRef, + parquet_metadata: &[Arc], +) -> Result { + let schemas = parquet_metadata + .iter() + .map(|metadata| { + let file = metadata.file_metadata(); + let schema = parquet_to_arrow_schema(file.schema_descr(), file.key_value_metadata()) + .map_err(|error| { + InvalidRecordBatchSnafu { + reason: format!("Failed to read compaction input Arrow schema: {error}"), + } + .build() + })?; + let rows = metadata + .row_groups() + .iter() + .map(|x| x.num_rows()) + .sum::() as u64; + Ok((Arc::new(schema), rows)) + }) + .collect::>>()?; + + collect_json2_rewrite_plans(metadata, &schemas) +} + +/// Builds JSON2 rewrite plans from existing physical schemas. +/// +/// Each source row count weights all of its existing explicit leaves. Ordinary paths are never +/// discovered from the v2 remainder, so this operation cannot unexpectedly promote opaque data. +pub(crate) fn collect_json2_rewrite_plans( + metadata: &RegionMetadataRef, + schemas: &[(SchemaRef, u64)], +) -> Result { + let json2_columns = metadata + .column_metadatas + .iter() + .filter_map(|x| { + x.column_schema + .data_type + .is_json2() + .then_some(&x.column_schema) + }) + .collect::>(); + + let mut plans = HashMap::with_capacity(json2_columns.len()); + for column in json2_columns { + let extension = column + .extension_type::() + .context(DataTypeMismatchSnafu)? + .with_context(|| InvalidRecordBatchSnafu { + reason: format!("JSON2 column '{}' has no extension metadata", column.name), + })?; + // Source SSTs may use v1, but current region metadata is copied to the rewritten output. + // Since the target physical layout is always v2, v1 metadata would produce an + // inconsistent persisted field. The region opener normally upgraded this metadata already. + ensure!( + extension.metadata().is_version_2(), + InvalidRecordBatchSnafu { + reason: format!("JSON2 column '{}' is not layout v2", column.name), + } + ); + + let settings = extension.metadata().json_settings(); + let hint_paths = settings + .type_hints() + .iter() + .map(|hint| hint.path.iter().map(String::as_str).collect::>()) + .collect::>(); + let mut stats = HashMap::new(); + for (schema, rows) in schemas { + let Some((_, field)) = schema.fields().find(&column.name) else { + continue; + }; + collect_json2_path_stats(field, *rows, &hint_paths, &mut stats)?; + } + + let mut hints = settings.type_hints().to_vec(); + hints.extend(select_dynamic_hints(settings, &hint_paths, &stats)); + let target_layout = JsonSettings::try_new(hints, Some(0)).context(DataTypeMismatchSnafu)?; + plans.insert( + column.name.clone(), + Json2RewritePlan { + logical_settings: settings.clone(), + target_layout, + }, + ); + } + Ok(plans) +} + +fn collect_json2_path_stats<'a>( + field: &'a Field, + rows: u64, + hint_paths: &HashSet>, + stats: &mut HashMap, Json2LeafPathStats>, +) -> Result<()> { + let ArrowDataType::Struct(fields) = field.data_type() else { + return InvalidRecordBatchSnafu { + reason: format!("JSON2 column '{}' is not a struct", field.name()), + } + .fail(); + }; + let mut paths = Vec::new(); + for field in fields { + if field.name() == JSON2_REMAINDER_FIELD_NAME { + continue; + } + collect_leaf_path_types(field, &mut Vec::new(), &mut paths)?; + } + + for (path, data_type) in paths { + if hint_paths.contains(path.as_slice()) { + continue; + } + let Some(stat) = stats.get_mut(&path) else { + stats.insert( + path, + Json2LeafPathStats { + rows, + data_type, + is_type_conflicted: false, + }, + ); + continue; + }; + if stat.data_type != data_type { + stat.is_type_conflicted = true; + } else { + stat.rows += rows; + } + } + Ok(()) +} + +fn collect_leaf_path_types<'a>( + field: &'a Field, + path: &mut Vec<&'a str>, + paths: &mut Vec<(Vec<&'a str>, JsonNativeType)>, +) -> Result<()> { + path.push(field.name()); + if let ArrowDataType::Struct(fields) = field.data_type() + && !fields.is_empty() + { + for field in fields { + collect_leaf_path_types(field, path, paths)?; + } + } else { + let json_type = + JsonNativeType::try_from(field.data_type()).context(DataTypeMismatchSnafu)?; + paths.push((path.clone(), json_type)); + } + path.pop(); + Ok(()) +} + +fn select_dynamic_hints( + settings: &JsonSettings, + hint_paths: &HashSet>, + stats: &HashMap, Json2LeafPathStats>, +) -> Vec { + let all_paths = stats + .keys() + .map(Vec::as_slice) + .chain(hint_paths.iter().map(Vec::as_slice)) + .collect::>(); + let has_ancestor_path = + |path: &[&str]| (1..path.len()).any(|len| all_paths.contains(&path[..len])); + + let prefixes = all_paths + .iter() + .copied() + .flat_map(|path| (1..path.len()).map(|len| &path[..len])) + .collect::>(); + let has_descendant_path = |path: &[&str]| prefixes.contains(path); + + let mut candidates = stats + .iter() + .filter(|(path, stat)| { + !stat.is_type_conflicted + && stat.data_type.is_primitive() + && !has_ancestor_path(path) + && !has_descendant_path(path) + }) + .collect::>(); + candidates.sort_unstable_by(|(x_path, x), (y_path, y)| { + y.rows.cmp(&x.rows).then_with(|| x_path.cmp(y_path)) + }); + candidates + .into_iter() + .take( + settings + .max_auto_expanded_paths() + .unwrap_or(JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS) as usize, + ) + .map(|(path, stat)| JsonTypeHint { + path: path.iter().map(|x| (*x).to_owned()).collect(), + data_type: ConcreteDataType::from_arrow_type(&stat.data_type.as_arrow_type()), + nullable: true, + default_constraint: None, + inverted_index: false, + }) + .collect() +} + +/// Replaces JSON2 physical field types according to the computed rewrite plans. +pub(crate) fn rewrite_json2_schema(schema: &SchemaRef, plans: &Json2RewritePlans) -> SchemaRef { + if plans.is_empty() { + return schema.clone(); + } + let fields = schema + .fields() + .iter() + .map(|field| { + let Some(plan) = plans.get(field.name()) else { + return field.clone(); + }; + let mut field = Field::clone(field); + field.set_data_type(json2_physical_data_type(&plan.target_layout)); + field = field.with_extension_type(Json2ExtensionType::new(Arc::new( + JsonMetadata::new(plan.logical_settings.clone()), + ))); + Arc::new(field) + }) + .collect::>(); + Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())) +} + +/// Rewrites JSON2 columns in `batch` according to the computed plans. +pub(crate) fn rewrite_json2_batch( + batch: RecordBatch, + plans: &Json2RewritePlans, +) -> Result { + if plans.is_empty() { + return Ok(batch); + } + let mut fields = Vec::with_capacity(batch.num_columns()); + let mut columns = Vec::with_capacity(batch.num_columns()); + + for (field, array) in batch.schema_ref().fields().iter().zip(batch.columns()) { + let Some(plan) = plans.get(field.name()) else { + fields.push(field.clone()); + columns.push(array.clone()); + continue; + }; + + let array = JsonArray::from(array) + .rewrite_to_v2(field, &plan.logical_settings, &plan.target_layout) + .context(ConvertValueSnafu)?; + debug_assert_eq!( + &json2_physical_data_type(&plan.target_layout), + array.data_type() + ); + + let mut field = Field::clone(field); + field.set_data_type(array.data_type().clone()); + field = field.with_extension_type(Json2ExtensionType::new(Arc::new(JsonMetadata::new( + plan.logical_settings.clone(), + )))); + fields.push(Arc::new(field)); + columns.push(array); + } + + let schema = Arc::new(Schema::new_with_metadata( + fields, + batch.schema_ref().metadata().clone(), + )); + RecordBatch::try_new(schema, columns).context(NewRecordBatchSnafu) +} + +#[cfg(test)] +mod tests { + use datatypes::extension::json::{ + JSON2_REMAINDER_FIELD_NAME, Json2PhysicalLayout, JsonMetadata, + }; + use datatypes::json::JsonTypeHint; + use datatypes::prelude::{ConcreteDataType, DataType}; + use datatypes::schema::ColumnSchema; + use datatypes::types::json_type::{JsonNativeType, JsonObjectType}; + use serde_json::json; + + use super::*; + + #[test] + fn test_select_dynamic_hints_rejects_type_and_prefix_conflicts() + -> Result<(), Box> { + let settings = JsonSettings::try_new( + vec![JsonTypeHint { + path: vec!["hint".to_string()], + data_type: ConcreteDataType::string_datatype(), + nullable: true, + default_constraint: None, + inverted_index: false, + }], + Some(2), + )?; + let stat = |rows, data_type, is_type_conflicted| Json2LeafPathStats { + rows, + data_type, + is_type_conflicted, + }; + let stats = HashMap::from([ + ( + vec!["hint", "nested"], + stat(10, JsonNativeType::String, false), + ), + (vec!["popular"], stat(9, JsonNativeType::String, false)), + ( + vec!["popular", "nested"], + stat(8, JsonNativeType::u64(), false), + ), + ( + vec!["type_conflicted"], + stat(7, JsonNativeType::String, true), + ), + ( + vec!["array"], + stat( + 6, + JsonNativeType::Array(Box::new(JsonNativeType::String)), + false, + ), + ), + (vec!["variant"], stat(5, JsonNativeType::Variant, false)), + (vec!["tie_a"], stat(2, JsonNativeType::String, false)), + (vec!["tie_b"], stat(2, JsonNativeType::Bool, false)), + ]); + + let hint_paths = settings + .type_hints() + .iter() + .map(|hint| hint.path.iter().map(String::as_str).collect::>()) + .collect::>(); + let hints = select_dynamic_hints(&settings, &hint_paths, &stats); + assert_eq!( + vec![vec!["tie_a".to_string()], vec!["tie_b".to_string()]], + hints.into_iter().map(|x| x.path).collect::>() + ); + Ok(()) + } + + #[test] + fn test_rewrite_json2_v1_batch_to_target_layout() -> Result<(), Box> { + let settings = JsonSettings::try_new( + vec![JsonTypeHint { + path: vec!["kind".to_string()], + data_type: ConcreteDataType::string_datatype(), + nullable: true, + default_constraint: None, + inverted_index: false, + }], + Some(0), + )?; + let target = json2_physical_data_type(&settings); + let plans = HashMap::from([( + "j".to_string(), + Json2RewritePlan { + logical_settings: settings.clone(), + target_layout: settings, + }, + )]); + let values = [ + json!({"kind": "a", "extra": {"x": 1}}), + json!({"kind": "b", "extra": {"x": 2}}), + ]; + let source_settings = JsonSettings::default(); + let source_extension = + Json2ExtensionType::new(Arc::new(JsonMetadata::new_v1(source_settings.clone()))); + let mut source_column = ColumnSchema::new( + "j", + ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())), + true, + ); + source_column.with_extension_type(&source_extension); + let mut source_builder = source_column.data_type.create_mutable_vector(values.len()); + for value in &values { + let value = source_settings.encode(value.clone())?; + source_builder.try_push_value_ref(&value.as_value_ref())?; + } + let source = source_builder.to_vector().to_arrow_array(); + let field = + Field::new("j", source.data_type().clone(), true).with_extension_type(source_extension); + let batch = RecordBatch::try_new(Arc::new(Schema::new(vec![field])), vec![source])?; + + let batch = rewrite_json2_batch(batch, &plans)?; + let field = batch.schema_ref().field(0); + assert!(Json2PhysicalLayout::try_from_root(field)?.is_version_2()); + assert_eq!(&target, field.data_type()); + let projected = + JsonArray::from(batch.column(0)).project_to_v2(field, &ArrowDataType::Binary)?; + let projected = JsonArray::from(&projected); + for (i, expected) in values.into_iter().enumerate() { + assert_eq!(expected, projected.try_get_value(i)?); + } + Ok(()) + } + + #[test] + fn test_rewrite_json2_v2_source_to_narrower_target() -> Result<(), Box> { + let target_settings = JsonSettings::try_new( + vec![JsonTypeHint { + path: vec!["kind".to_string()], + data_type: ConcreteDataType::string_datatype(), + nullable: true, + default_constraint: None, + inverted_index: false, + }], + Some(0), + )?; + let target_extension = + Json2ExtensionType::new(Arc::new(JsonMetadata::new(target_settings.clone()))); + let target_type = json2_physical_data_type(&target_settings); + let plans = HashMap::from([( + "j".to_string(), + Json2RewritePlan { + logical_settings: target_settings.clone(), + target_layout: target_settings, + }, + )]); + + let source_settings = JsonSettings::try_new( + vec![ + JsonTypeHint { + path: vec!["kind".to_string()], + data_type: ConcreteDataType::string_datatype(), + nullable: true, + default_constraint: None, + inverted_index: false, + }, + JsonTypeHint { + path: vec!["source_only".to_string()], + data_type: ConcreteDataType::int64_datatype(), + nullable: true, + default_constraint: None, + inverted_index: false, + }, + ], + None, + )?; + let source_extension = + Json2ExtensionType::new(Arc::new(JsonMetadata::new(source_settings.clone()))); + let mut source_column = ColumnSchema::new( + "j", + ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())), + true, + ); + source_column.with_extension_type(&source_extension); + let expected = json!({ + "kind": "a", + "source_only": 7, + "dynamic": {"nested": true} + }); + let mut builder = source_column.create_mutable_vector(1); + let value = source_settings.encode(expected.clone())?; + builder.try_push_value_ref(&value.as_value_ref())?; + let source = builder.to_vector().to_arrow_array(); + let source_field = + Field::new("j", source.data_type().clone(), true).with_extension_type(source_extension); + let source = + JsonArray::from(&source).project_to_v2(&source_field, &ArrowDataType::Binary)?; + let field = Field::new("j", ArrowDataType::Binary, true) + .with_extension_type(target_extension.clone()); + let batch = RecordBatch::try_new(Arc::new(Schema::new(vec![field])), vec![source])?; + + let batch = rewrite_json2_batch(batch, &plans)?; + let field = batch.schema_ref().field(0); + assert!(Json2PhysicalLayout::try_from_root(field)?.is_version_2()); + assert_eq!(&target_type, field.data_type()); + let ArrowDataType::Struct(fields) = field.data_type() else { + unreachable!() + }; + assert_eq!( + vec![JSON2_REMAINDER_FIELD_NAME, "kind"], + fields.iter().map(|x| x.name().as_str()).collect::>() + ); + + let projected = + JsonArray::from(batch.column(0)).project_to_v2(field, &ArrowDataType::Binary)?; + assert_eq!(expected, JsonArray::from(&projected).try_get_value(0)?); + + let first_schema = batch.schema(); + let field = + Field::new("j", ArrowDataType::Binary, true).with_extension_type(target_extension); + let batch = RecordBatch::try_new(Arc::new(Schema::new(vec![field])), vec![projected])?; + let batch = rewrite_json2_batch(batch, &plans)?; + assert_eq!(first_schema, batch.schema()); + + let field = batch.schema_ref().field(0); + let projected = + JsonArray::from(batch.column(0)).project_to_v2(field, &ArrowDataType::Binary)?; + assert_eq!(expected, JsonArray::from(&projected).try_get_value(0)?); + Ok(()) + } + + #[test] + fn test_rewrite_json2_v2_source_to_wider_target() -> Result<(), Box> { + let logical_settings = JsonSettings::try_new( + vec![JsonTypeHint { + path: vec!["kind".to_string()], + data_type: ConcreteDataType::string_datatype(), + nullable: true, + default_constraint: None, + inverted_index: false, + }], + Some(0), + )?; + let target_layout = JsonSettings::try_new( + vec![ + JsonTypeHint { + path: vec!["kind".to_string()], + data_type: ConcreteDataType::string_datatype(), + nullable: true, + default_constraint: None, + inverted_index: false, + }, + JsonTypeHint { + path: vec!["promoted".to_string()], + data_type: ConcreteDataType::int64_datatype(), + nullable: true, + default_constraint: None, + inverted_index: false, + }, + ], + Some(0), + )?; + let target_type = json2_physical_data_type(&target_layout); + let plans = HashMap::from([( + "j".to_string(), + Json2RewritePlan { + logical_settings: logical_settings.clone(), + target_layout, + }, + )]); + + let extension = + Json2ExtensionType::new(Arc::new(JsonMetadata::new(logical_settings.clone()))); + let mut column = ColumnSchema::new( + "j", + ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())), + true, + ); + column.with_extension_type(&extension); + let expected = json!({ + "kind": "a", + "promoted": 7, + "dynamic": {"nested": true} + }); + let mut builder = column.create_mutable_vector(1); + let value = logical_settings.encode(expected.clone())?; + builder.try_push_value_ref(&value.as_value_ref())?; + let source = builder.to_vector().to_arrow_array(); + assert_eq!( + &json2_physical_data_type(&logical_settings), + source.data_type() + ); + let field = + Field::new("j", source.data_type().clone(), true).with_extension_type(extension); + let batch = RecordBatch::try_new(Arc::new(Schema::new(vec![field])), vec![source])?; + + let batch = rewrite_json2_batch(batch, &plans)?; + let field = batch.schema_ref().field(0); + assert_eq!(&target_type, field.data_type()); + let ArrowDataType::Struct(fields) = field.data_type() else { + unreachable!() + }; + assert_eq!( + vec![JSON2_REMAINDER_FIELD_NAME, "kind", "promoted"], + fields.iter().map(|x| x.name().as_str()).collect::>() + ); + + let array = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let promoted = array + .column_by_name("promoted") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(7, promoted.value(0)); + + let projected = + JsonArray::from(batch.column(0)).project_to_v2(field, &ArrowDataType::Binary)?; + assert_eq!(expected, JsonArray::from(&projected).try_get_value(0)?); + Ok(()) + } +} diff --git a/src/mito2/src/compaction/reader.rs b/src/mito2/src/compaction/reader.rs index f862d11e1a..2f604799cf 100644 --- a/src/mito2/src/compaction/reader.rs +++ b/src/mito2/src/compaction/reader.rs @@ -12,26 +12,25 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::{BTreeMap, HashMap}; +use std::collections::BTreeMap; use std::sync::Arc; +use arrow_schema::extension::EXTENSION_TYPE_METADATA_KEY; use common_time::Timestamp; use common_time::range::TimestampRange; use common_time::timestamp::TimeUnit; use datafusion_common::ScalarValue; use datafusion_expr::Expr; -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}; -use snafu::{OptionExt, ResultExt}; +use snafu::OptionExt; use store_api::metadata::RegionMetadataRef; use crate::access_layer::AccessLayerRef; use crate::cache::{CacheManagerRef, CacheStrategy}; -use crate::error::{ - DataTypeMismatchSnafu, ParquetToArrowSchemaSnafu, Result, TimeRangePredicateOverflowSnafu, +use crate::compaction::json2::{ + Json2RewritePlans, collect_json2_rewrite_plans_from_parquet, rewrite_json2_schema, }; +use crate::error::{InvalidRecordBatchSnafu, Result, TimeRangePredicateOverflowSnafu}; use crate::read::FlatSource; use crate::read::flat_projection::FlatProjectionMapper; use crate::read::read_columns::ReadColumns; @@ -40,6 +39,7 @@ use crate::read::seq_scan::SeqScan; use crate::region::options::MergeMode; use crate::sst::file::FileHandle; use crate::sst::parquet::reader::MetadataCacheMetrics; +use crate::sst::parquet::{Json2RewriteTargets, Json2TargetLayout}; /// Builders to create [BoxedRecordBatchStream] for compaction. pub(crate) struct CompactionSstReaderBuilder<'a> { @@ -57,20 +57,24 @@ impl CompactionSstReaderBuilder<'_> { /// Build a [FlatSource] that yields Arrow `RecordBatch`s from reading all the input SST files, /// for compaction. The schema of the [FlatSource] is unified. pub(crate) async fn build_flat_sst_reader(self) -> Result { - let scan_input = self.build_scan_input().await?; + let parquet_metadata = self.collect_parquet_metadata().await?; + let plans = collect_json2_rewrite_plans_from_parquet(&self.metadata, &parquet_metadata)?; + let scan_input = self.build_scan_input(&parquet_metadata, &plans)?; let schema = scan_input.mapper.output_schema(); - let schema = schema.arrow_schema(); + let schema = rewrite_json2_schema(schema.arrow_schema(), &plans); let stream = SeqScan::new(scan_input) .build_flat_reader_for_compaction() .await?; - Ok(FlatSource::new_stream(schema.clone(), stream)) + Ok(FlatSource::new_stream(schema, stream)) } - async fn build_scan_input(self) -> Result { - let schema = self.metadata.schema.arrow_schema(); - let parquet_metadata = self.collect_parquet_metadata().await?; + fn build_scan_input( + self, + parquet_metadata: &[Arc], + plans: &Json2RewritePlans, + ) -> Result { let batch_size = crate::batch_size::estimate_batch_size( parquet_metadata .iter() @@ -84,38 +88,6 @@ impl CompactionSstReaderBuilder<'_> { (row_group.num_rows() as u64, uncompressed_bytes) }), ); - let json_type_hint = if schema.fields().iter().any(is_json2_extension_type) { - let mut json_type_hint = schema - .fields() - .iter() - .filter(|&field| is_json2_extension_type(field)) - .map(|field| (field.name().clone(), JsonNativeType::Null)) - .collect::>(); - - for metadata in &parquet_metadata { - let file_metadata = metadata.file_metadata(); - let schema = parquet_to_arrow_schema( - file_metadata.schema_descr(), - file_metadata.key_value_metadata(), - ) - .context(ParquetToArrowSchemaSnafu { - file: "compaction input", - })?; - for field in schema.fields() { - let Some(merged) = json_type_hint.get_mut(field.name()) else { - continue; - }; - - let json_type = JsonNativeType::try_from(field.data_type()) - .context(DataTypeMismatchSnafu)?; - merged.merge(&json_type); - } - } - - Some(json_type_hint) - } else { - None - }; let projection = (0..self.metadata.column_metadatas.len()).collect(); let read_column_ids = self @@ -124,24 +96,39 @@ impl CompactionSstReaderBuilder<'_> { .iter() .map(|x| x.column_id) .collect::>(); - let json_target_types = json_type_hint - .as_ref() - .map(|hint| { - hint.iter() - .filter_map(|(col_name, json_type)| { - self.metadata - .column_by_name(col_name) - .map(|col| (col.column_id, json_type.clone())) - }) - .collect::>() - }) - .unwrap_or_default(); - let read_columns = - ReadColumns::new(read_column_ids).with_json_target_types(json_target_types); - let mapper = - FlatProjectionMapper::new_with_read_columns(&self.metadata, projection, read_columns)?; + + let mut json2_target_layouts = BTreeMap::new(); + for (name, plan) in plans { + let Some(column) = self.metadata.column_by_name(name) else { + continue; + }; + let extension_metadata = column + .column_schema + .metadata() + .get(EXTENSION_TYPE_METADATA_KEY) + .cloned() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!("JSON2 target column '{name}' has no extension metadata"), + })?; + json2_target_layouts.insert( + column.column_id, + Json2TargetLayout { + extension_metadata, + target_layout: plan.target_layout.clone(), + }, + ); + } + let read_columns = ReadColumns::new(read_column_ids); + let targets: Json2RewriteTargets = Arc::new(json2_target_layouts); + let mapper = FlatProjectionMapper::new_with_json2_rewrite_targets( + &self.metadata, + projection, + read_columns, + &targets, + )?; let mut scan_input = ScanInput::new(self.sst_layer, mapper) + .with_json2_rewrite_targets(targets) .with_files(self.inputs.to_vec()) .with_compaction(true) .with_batch_size(batch_size) diff --git a/src/mito2/src/engine/create_test.rs b/src/mito2/src/engine/create_test.rs index 25b6995190..2c04e2e24f 100644 --- a/src/mito2/src/engine/create_test.rs +++ b/src/mito2/src/engine/create_test.rs @@ -15,9 +15,13 @@ use std::time::Duration; use api::v1::Rows; +use common_error::ext::{ErrorExt, WhateverResult}; +use common_error::status_code::StatusCode; use common_recordbatch::RecordBatches; +use datatypes::prelude::ConcreteDataType; +use datatypes::types::json_type::{JsonNativeType, JsonObjectType}; use store_api::region_engine::RegionEngine; -use store_api::region_request::{RegionCloseRequest, RegionRequest}; +use store_api::region_request::{PathType, RegionCloseRequest, RegionOpenRequest, RegionRequest}; use store_api::storage::{RegionId, ScanRequest}; use crate::config::MitoConfig; @@ -32,6 +36,81 @@ async fn test_engine_create_new_region() { test_engine_create_new_region_with_format(true).await; } +#[tokio::test] +async fn test_engine_rejects_json2_with_time_series_memtable_on_create_and_open() +-> WhateverResult<()> { + let mut env = TestEnv::with_prefix("json2-rejects-time-series-memtable").await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: false, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new() + .field_datatype(ConcreteDataType::json2(JsonNativeType::Object( + JsonObjectType::new(), + ))) + .insert_option("append_mode", "true") + .insert_option("memtable.type", "time_series") + .build(); + + let err = engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap_err(); + assert_eq!(StatusCode::InvalidArguments, err.status_code()); + assert!( + err.to_string() + .contains("JSON2 columns only support BulkMemtable") + ); + + let request = CreateRequestBuilder::new() + .field_datatype(ConcreteDataType::json2(JsonNativeType::Object( + JsonObjectType::new(), + ))) + .insert_option("append_mode", "true") + .insert_option("memtable.type", "bulk") + .build(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await?; + engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await?; + + let options = [ + ("append_mode".to_string(), "true".to_string()), + ("memtable.type".to_string(), "time_series".to_string()), + ] + .into_iter() + .collect(); + let err = engine + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir: "test".to_string(), + path_type: PathType::Bare, + options, + skip_wal_replay: false, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap_err(); + assert_eq!(StatusCode::InvalidArguments, err.status_code()); + assert!( + err.output_msg() + .contains("JSON2 columns only support BulkMemtable") + ); + Ok(()) +} + async fn test_engine_create_new_region_with_format(flat_format: bool) { let mut env = TestEnv::with_prefix("new-region").await; let engine = env diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index a53b77c825..3371d7a21a 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -19,6 +19,7 @@ use api::helper::encode_json_value; use api::v1::helper::row; use api::v1::value::ValueData; use api::v1::{ColumnDataType, Rows, SemanticType, WriteHint}; +use arrow_schema::extension::ExtensionType; use common_base::readable_size::ReadableSize; use common_error::ext::{ErrorExt, WhateverResult}; use common_error::status_code::StatusCode; @@ -28,16 +29,19 @@ use datafusion_common::ScalarValue; use datafusion_expr::{col, lit}; use datatypes::arrow::array::AsArray; use datatypes::arrow::datatypes::{Float64Type, TimestampMillisecondType, UInt64Type}; +use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; +use datatypes::json::JsonSettings; use datatypes::json::value::JsonValue; use datatypes::prelude::ConcreteDataType; use datatypes::types::json_type::{JsonNativeType, JsonObjectType}; +use datatypes::vectors::json::array::JsonArray; use futures::TryStreamExt; use futures::future::try_join_all; use serde_json::json; use store_api::codec::PrimaryKeyEncoding; use store_api::metric_engine_consts::PRIMARY_KEY_ENCODING; use store_api::region_engine::{PrepareRequest, RegionEngine, RegionScanner}; -use store_api::region_request::{RegionPutRequest, RegionRequest}; +use store_api::region_request::{RegionCompactRequest, RegionPutRequest, RegionRequest}; use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; use store_api::storage::{RegionId, ScanRequest, TimeSeriesDistribution}; @@ -47,7 +51,7 @@ use crate::read::read_columns::ReadColumns; use crate::read::scan_region::Scanner; use crate::test_util; use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_with_encoding}; -use crate::test_util::{CreateRequestBuilder, TestEnv}; +use crate::test_util::{CreateRequestBuilder, TestEnv, reopen_region}; #[tokio::test] async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResult<()> { @@ -79,28 +83,34 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul // Write full JSON objects, then flush them so the scanner has an Parquet file where nested // projection can be pushed down. - let rows = Rows { - schema, - rows: vec![ - row(vec![ - ValueData::StringValue("tag-1".to_string()), - ValueData::JsonValue(encode_json_value(JsonValue::from(json!({ - "a": { "x": 10, "y": "ignored-a" }, - "b": "ignored-b" - })))), - ValueData::TimestampMillisecondValue(1000), - ]), - row(vec![ - ValueData::StringValue("tag-2".to_string()), - ValueData::JsonValue(encode_json_value(JsonValue::from(json!({ - "a": { "x": 20, "y": "ignored-c" }, - "b": "ignored-d" - })))), - ValueData::TimestampMillisecondValue(2000), - ]), + for values in [ + vec![ + ValueData::StringValue("tag-1".to_string()), + ValueData::JsonValue(encode_json_value(JsonValue::from(json!({ + "a": { "x": 10, "y": "ignored-a" }, + "b": "ignored-b" + })))), + ValueData::TimestampMillisecondValue(1000), ], - }; - test_util::put_rows(&engine, region_id, rows).await; + vec![ + ValueData::StringValue("tag-2".to_string()), + ValueData::JsonValue(encode_json_value(JsonValue::from(json!({ + "a": { "x": 20, "y": "ignored-c" }, + "b": "ignored-d" + })))), + ValueData::TimestampMillisecondValue(2000), + ], + ] { + test_util::put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: vec![row(values)], + }, + ) + .await; + } test_util::flush_region(&engine, region_id, None).await; // Without a type hint, the scanner reads the whole JSON2 root column. @@ -214,6 +224,231 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul Ok(()) } +#[tokio::test] +async fn test_json2_v1_region_reopen_and_compaction() -> WhateverResult<()> { + let mut request = CreateRequestBuilder::new() + .field_datatype(ConcreteDataType::json2(JsonNativeType::Object( + JsonObjectType::new(), + ))) + .insert_option("memtable.type", "bulk") + .build(); + let settings = JsonSettings::default(); + request.column_metadatas[1] + .column_schema + .with_extension_type(&Json2ExtensionType::new(Arc::new(JsonMetadata::new_v1( + settings, + )))); + let table_dir = request.table_dir.clone(); + let schema = test_util::rows_schema(&request); + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1024, 0); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await?; + + let values = [ + json!({"route": 1}), + json!({"b": {"c": "x"}}), + json!({"route": 3, "written_after_reopen": true}), + ]; + for (i, value) in values[..2].iter().enumerate() { + test_util::put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: vec![row(vec![ + ValueData::StringValue("tag".to_string()), + ValueData::JsonValue(encode_json_value(JsonValue::from(value.clone()))), + ValueData::TimestampMillisecondValue((i as i64 + 1) * 1000), + ])], + }, + ) + .await; + test_util::flush_region(&engine, region_id, None).await; + } + + reopen_region( + &engine, + region_id, + table_dir, + true, + HashMap::from([("memtable.type".to_string(), "bulk".to_string())]), + ) + .await; + let region = engine.get_region(region_id).unwrap(); + let version = region.version(); + let column = &version + .metadata + .column_by_name("field_0") + .unwrap() + .column_schema; + let extension = column.extension_type::()?.unwrap(); + assert!(extension.metadata().is_version_2()); + + test_util::put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: vec![row(vec![ + ValueData::StringValue("tag".to_string()), + ValueData::JsonValue(encode_json_value(JsonValue::from(values[2].clone()))), + ValueData::TimestampMillisecondValue(3000), + ])], + }, + ) + .await; + test_util::flush_region(&engine, region_id, None).await; + + let region = engine.get_region(region_id).unwrap(); + let input_files = region + .version() + .ssts + .levels() + .iter() + .map(|level| level.files.len()) + .sum::(); + assert_eq!(3, input_files); + + // The same requested path is stored as a v1 explicit leaf in the first SST, is absent from + // the second SST, and lives in the v2 remainder in the third SST. A per-file route preserves + // those differences while exposing one logical query type to the merge reader. + let scanner = engine + .scanner( + region_id, + ScanRequest { + projection: Some(vec![1]), + json_type_hint: HashMap::from([( + "field_0".to_string(), + JsonNativeType::Object(JsonObjectType::from([( + "route".to_string(), + JsonNativeType::i64(), + )])), + )]), + ..Default::default() + }, + ) + .await?; + let batches = RecordBatches::try_collect(scanner.scan().await?).await?; + let mut routed = Vec::new(); + for batch in batches.iter() { + let array = batch.column_by_name("field_0").unwrap().clone(); + let json = JsonArray::from(&array); + for i in 0..array.len() { + routed.push(json.try_get_value(i)?); + } + } + assert_eq!( + [json!({"route": 1}), json!(null), json!({"route": 3})], + routed.as_slice() + ); + + engine + .handle_request( + region_id, + RegionRequest::Compact(RegionCompactRequest::default()), + ) + .await?; + + let scanner = engine + .scanner( + region_id, + ScanRequest { + projection: Some(vec![1]), + ..Default::default() + }, + ) + .await?; + let batches = RecordBatches::try_collect(scanner.scan().await?).await?; + let mut actual = Vec::new(); + for batch in batches.iter() { + let array = batch.column_by_name("field_0").unwrap().clone(); + let json = JsonArray::from(&array); + for i in 0..array.len() { + actual.push(json.try_get_value(i)?); + } + } + assert_eq!(values, actual.as_slice()); + Ok(()) +} + +#[tokio::test] +async fn test_flush_aligns_different_json2_layouts() -> WhateverResult<()> { + let mut request = CreateRequestBuilder::new() + .field_datatype(ConcreteDataType::json2(JsonNativeType::Object( + JsonObjectType::new(), + ))) + .insert_option("append_mode", "true") + .insert_option("memtable.type", "bulk") + .build(); + let settings = JsonSettings::try_new(vec![], Some(1))?; + request.column_metadatas[1] + .column_schema + .with_extension_type(&Json2ExtensionType::new(Arc::new(JsonMetadata::new( + settings, + )))); + let schema = test_util::rows_schema(&request); + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1025, 0); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await?; + + for (offset, name) in [(0, "a"), (1024, "b")] { + let rows = (0..1024) + .map(|i| { + row(vec![ + ValueData::StringValue("tag".to_string()), + ValueData::JsonValue(encode_json_value(JsonValue::from(json!({(name): i})))), + ValueData::TimestampMillisecondValue(offset + i), + ]) + }) + .collect(); + test_util::put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows, + }, + ) + .await; + } + + test_util::flush_region(&engine, region_id, None).await; + + let scanner = engine + .scanner( + region_id, + ScanRequest { + projection: Some(vec![1]), + ..Default::default() + }, + ) + .await?; + let batches = RecordBatches::try_collect(scanner.scan().await?).await?; + let mut counts = HashMap::new(); + for batch in batches.iter() { + let array = batch.column_by_name("field_0").unwrap().clone(); + let json = JsonArray::from(&array); + for i in 0..array.len() { + let value = json.try_get_value(i)?; + let name = match (value.get("a"), value.get("b")) { + (Some(_), None) => "a", + (None, Some(_)) => "b", + _ => panic!("expected exactly one dynamic JSON2 field, got {value}"), + }; + *counts.entry(name).or_insert(0) += 1; + } + } + assert_eq!(Some(&1024), counts.get("a")); + assert_eq!(Some(&1024), counts.get("b")); + Ok(()) +} + #[tokio::test] async fn test_incremental_query_stale_error() { let mut env = TestEnv::with_prefix("test_incremental_query_stale_error").await; diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index a4a42970e4..7c088dea58 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -24,10 +24,10 @@ use bytes::Bytes; use common_base::cancellation::CancellableFuture; use common_telemetry::{debug, error, info}; use datatypes::arrow::datatypes::SchemaRef; -use datatypes::extension::json::is_json2_extension_type; use partition::expr::PartitionExpr; use smallvec::{SmallVec, smallvec}; -use snafu::ResultExt; +use snafu::{ResultExt, ensure}; +use store_api::metadata::RegionMetadataRef; use store_api::region_request::RegionFlushReason; use store_api::storage::{RegionId, SequenceNumber}; use strum::IntoStaticStr; @@ -37,15 +37,15 @@ use crate::access_layer::{ AccessLayerRef, Metrics, OperationType, SstInfoArray, SstWriteRequest, WriteType, }; use crate::cache::CacheManagerRef; +use crate::compaction::{collect_json2_rewrite_plans, rewrite_json2_batch, rewrite_json2_schema}; use crate::config::MitoConfig; use crate::engine::region_hook::SstFileInfo; use crate::error::{ Error, FlushCancelledSnafu, FlushRegionSnafu, JoinSnafu, RegionBusySnafu, RegionClosedSnafu, - RegionDroppedSnafu, RegionTruncatedSnafu, Result, + RegionDroppedSnafu, RegionTruncatedSnafu, Result, UnexpectedSnafu, }; use crate::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList}; use crate::memtable::bulk::ENCODE_ROW_THRESHOLD; -use crate::memtable::bulk::json_align::Json2Aligner; use crate::memtable::{BoxedRecordBatchIterator, EncodedRange, MemtableRanges, RangesOptions}; use crate::metrics::{ FLUSH_BYTES_TOTAL, FLUSH_ELAPSED, FLUSH_FAILURE_TOTAL, FLUSH_FILE_TOTAL, FLUSH_REQUESTS_TOTAL, @@ -708,6 +708,7 @@ impl RegionFlushTask { let flat_sources = memtable_flat_sources( batch_schema, mem_ranges, + &version.metadata, &version.options, field_column_start, )?; @@ -901,6 +902,7 @@ struct FlatSources { fn memtable_flat_sources( schema: SchemaRef, mem_ranges: MemtableRanges, + metadata: &RegionMetadataRef, options: &RegionOptions, field_column_start: usize, ) -> Result { @@ -918,6 +920,7 @@ fn memtable_flat_sources( if let Some(encoded) = only_range.encoded() { flat_sources.encoded.push((encoded, max_sequence)); } else { + let schema = only_range.record_batch_schema_hint().unwrap_or(schema); let iter = only_range.build_record_batch_iter(None, None)?; // Dedup according to append mode and merge mode. // Even single range may have duplicate rows. @@ -951,12 +954,23 @@ 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_json2_extension_type); - let mut json_align_schemas = if has_json2 { - Some(Vec::with_capacity(num_ranges)) - } else { - None - }; + let schemas = ranges + .values() + .filter(|range| range.encoded().is_none()) + .map(|range| { + ( + range + .record_batch_schema_hint() + .unwrap_or_else(|| schema.clone()), + range.num_rows() as u64, + ) + }) + .collect::>(); + let plans = Arc::new(collect_json2_rewrite_plans(metadata, &schemas)?); + let schema = rewrite_json2_schema( + schemas.first().map(|(schema, _)| schema).unwrap_or(&schema), + &plans, + ); for (_range_id, range) in ranges { if let Some(encoded) = range.encoded() { @@ -965,15 +979,26 @@ fn memtable_flat_sources( continue; } - // Collect schemas if has json2 field. - if let Some(schemas) = json_align_schemas.as_mut() { - let schema = range - .record_batch_schema_hint() - .unwrap_or_else(|| schema.clone()); - schemas.push(schema); + if let Some(actual) = range.record_batch_schema_hint() { + let actual = rewrite_json2_schema(&actual, &plans); + ensure!( + actual == schema, + UnexpectedSnafu { + reason: format!( + "Different schemas found in a MemtableRanges, expected: {}, actual: {}", + schema, actual, + ), + } + ) } let iter = range.build_record_batch_iter(None, None)?; + let iter: BoxedRecordBatchIterator = if plans.is_empty() { + iter + } else { + let plans = plans.clone(); + Box::new(iter.map(move |batch| rewrite_json2_batch(batch?, &plans))) + }; input_iters.push(iter); let range_rows = range.num_rows(); last_iter_rows += range_rows; @@ -1007,11 +1032,6 @@ fn memtable_flat_sources( let input_iters = std::mem::replace(&mut input_iters, Vec::with_capacity(num_ranges)); - let (schema, input_iters) = maybe_align_json2_iters( - schema.clone(), - json_align_schemas.take(), - input_iters, - )?; let maybe_dedup = merge_and_dedup_with_batch_size( &schema, @@ -1022,17 +1042,12 @@ fn memtable_flat_sources( batch_size, )?; - flat_sources - .sources - .push((FlatSource::new_iter(schema, maybe_dedup), max_sequence)); + flat_sources.sources.push(( + FlatSource::new_iter(schema.clone(), maybe_dedup), + max_sequence, + )); last_iter_rows = 0; current_ranges.clear(); - - json_align_schemas = if has_json2 { - Some(Vec::with_capacity(num_ranges)) - } else { - None - }; } } @@ -1046,9 +1061,6 @@ fn memtable_flat_sources( rows_remaining ); - let (schema, input_iters) = - maybe_align_json2_iters(schema, json_align_schemas, input_iters)?; - let max_sequence = current_ranges .iter() .map(|r| r.stats().max_sequence()) @@ -1078,24 +1090,6 @@ fn memtable_flat_sources( Ok(flat_sources) } -fn maybe_align_json2_iters( - schema: SchemaRef, - schemas: Option>, - input_iters: Vec, -) -> Result<(SchemaRef, Vec)> { - let Some(schemas) = schemas else { - return Ok((schema, input_iters)); - }; - - let aligner = Json2Aligner::try_new(schemas)?; - let input_iters = input_iters - .into_iter() - .map(|input_iter| aligner.wrap_iter(input_iter)) - .collect(); - - Ok((aligner.schema().clone(), input_iters)) -} - /// Merges multiple record batch iterators and applies deduplication based on the specified mode. /// /// This function is used during the flush process to combine data from multiple memtable ranges @@ -1646,6 +1640,8 @@ mod tests { use api::v1::{OpType, Rows}; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; + use datatypes::arrow::datatypes::Schema; + use datatypes::arrow::record_batch::RecordBatch; use mito_codec::row_converter::build_primary_key_codec; use tokio::sync::oneshot; @@ -1654,7 +1650,9 @@ mod tests { use crate::error::InvalidSchedulerStateSnafu; use crate::memtable::bulk::part::BulkPartConverter; use crate::memtable::time_series::TimeSeriesMemtableBuilder; - use crate::memtable::{Memtable, RangesOptions}; + use crate::memtable::{ + IterBuilder, Memtable, MemtableRange, MemtableRangeContext, MemtableStats, RangesOptions, + }; use crate::request::WriteRequest; use crate::schedule::scheduler::Scheduler; use crate::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; @@ -2243,6 +2241,7 @@ mod tests { let flat_sources = memtable_flat_sources( schema.clone(), mem_ranges, + &metadata, &options, metadata.primary_key.len(), ) @@ -2271,9 +2270,14 @@ mod tests { ..Default::default() }; - let flat_sources = - memtable_flat_sources(schema, mem_ranges, &options, metadata.primary_key.len()) - .unwrap(); + let flat_sources = memtable_flat_sources( + schema, + mem_ranges, + &metadata, + &options, + metadata.primary_key.len(), + ) + .unwrap(); assert!(flat_sources.encoded.is_empty()); assert_eq!(1, flat_sources.sources.len()); @@ -2288,6 +2292,118 @@ mod tests { } } + #[test] + fn test_memtable_flat_sources_uses_non_encoded_schema() -> Result<()> { + struct TestIterBuilder { + schema: SchemaRef, + batch: Option, + } + + impl IterBuilder for TestIterBuilder { + fn build( + &self, + _metrics: Option, + ) -> Result { + unimplemented!() + } + + fn is_record_batch(&self) -> bool { + true + } + + fn build_record_batch( + &self, + _time_range: Option<(common_time::Timestamp, common_time::Timestamp)>, + _metrics: Option, + ) -> Result { + let Some(batch) = self.batch.clone() else { + unimplemented!() + }; + Ok(Box::new(std::iter::once(Ok(batch)))) + } + + fn record_batch_schema_hint(&self) -> Option { + Some(self.schema.clone()) + } + + fn encoded_range(&self) -> Option { + self.batch.is_none().then(|| EncodedRange { + data: Bytes::new(), + sst_info: SstInfo::default(), + }) + } + } + + let metadata = metadata_for_test(); + let schema = to_flat_sst_arrow_schema( + &metadata, + &FlatSchemaOptions::from_encoding(metadata.primary_key_encoding), + ); + let pk_codec = build_primary_key_codec(&metadata); + let mut converter = BulkPartConverter::new(&metadata, schema.clone(), 1, pk_codec, true); + let kvs = build_key_values_with_ts_seq_values( + &metadata, + "key".to_string(), + 1, + std::iter::once(1000), + std::iter::once(Some(1.0)), + 1, + ); + converter.append_key_values(&kvs)?; + let batch = converter.convert()?.batch; + let encoded_schema = Arc::new(Schema::empty()); + + let new_range = |id, builder| { + MemtableRange::new( + Arc::new(MemtableRangeContext::new( + id, + Box::new(builder), + Default::default(), + )), + MemtableStats { + num_rows: 1, + ..Default::default() + }, + ) + }; + let mut ranges = std::collections::BTreeMap::new(); + ranges.insert( + 0, + new_range( + 0, + TestIterBuilder { + schema: encoded_schema.clone(), + batch: None, + }, + ), + ); + ranges.insert( + 1, + new_range( + 0, + TestIterBuilder { + schema: schema.clone(), + batch: Some(batch), + }, + ), + ); + + let sources = memtable_flat_sources( + encoded_schema, + MemtableRanges { ranges }, + &metadata, + &RegionOptions { + append_mode: true, + ..Default::default() + }, + metadata.primary_key.len(), + )?; + assert_eq!(1, sources.encoded.len()); + assert_eq!(1, sources.sources.len()); + assert_eq!(&schema, sources.sources[0].0.schema()); + Ok(()) + } + #[tokio::test] async fn test_schedule_pending_request_on_flush_success() { common_telemetry::init_default_ut_logging(); diff --git a/src/mito2/src/memtable.rs b/src/mito2/src/memtable.rs index 6cc6e327ae..8c107d15f7 100644 --- a/src/mito2/src/memtable.rs +++ b/src/mito2/src/memtable.rs @@ -30,11 +30,11 @@ pub use mito_codec::key_values::KeyValues; use mito_codec::row_converter::{PrimaryKeyCodec, build_primary_key_codec}; use snafu::ensure; use store_api::codec::PrimaryKeyEncoding; -use store_api::metadata::RegionMetadataRef; +use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::storage::{ColumnId, SequenceNumber, SequenceRange}; use crate::config::MitoConfig; -use crate::error::{Result, UnsupportedOperationSnafu}; +use crate::error::{InvalidRegionOptionsSnafu, Result, UnsupportedOperationSnafu}; use crate::flush::WriteBufferManagerRef; use crate::memtable::bulk::{BulkMemtableBuilder, CompactDispatcher}; use crate::memtable::time_series::TimeSeriesMemtableBuilder; @@ -401,6 +401,26 @@ pub(crate) struct MemtableBuilderProvider { compact_dispatcher: Arc, } +/// Ensures JSON2 columns are not used with [`TimeSeriesMemtable`]. +pub(crate) fn ensure_json2_not_use_time_series_memtable( + metadata: &RegionMetadata, + options: &RegionOptions, +) -> Result<()> { + if metadata + .column_metadatas + .iter() + .any(|x| x.column_schema.data_type.is_json2()) + { + ensure!( + !matches!(&options.memtable, Some(MemtableOptions::TimeSeries)), + InvalidRegionOptionsSnafu { + reason: "JSON2 columns only support BulkMemtable", + } + ); + } + Ok(()) +} + impl MemtableBuilderProvider { pub(crate) fn new( write_buffer_manager: Option, @@ -772,9 +792,15 @@ impl MemtableRange { mod tests { use std::sync::Arc; + use common_error::ext::WhateverResult; + use datatypes::prelude::ConcreteDataType; + use datatypes::types::json_type::{JsonNativeType, JsonObjectType}; + use store_api::metadata::RegionMetadataBuilder; + use super::*; use crate::flush::{WriteBufferManager, WriteBufferManagerImpl}; use crate::memtable::bulk::BulkMemtableConfig; + use crate::test_util::sst_util::sst_region_metadata; #[test] fn test_alloc_tracker_without_manager() { @@ -848,4 +874,27 @@ mod tests { assert_eq!(&config, builder.config()); } + + #[test] + fn test_json2_requires_bulk_memtable() -> WhateverResult<()> { + let mut metadata = sst_region_metadata(); + metadata.column_metadatas[2].column_schema.data_type = + ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())); + let metadata = RegionMetadataBuilder::from_existing(metadata).build()?; + let mut options = RegionOptions { + sst_format: Some(FormatType::PrimaryKey), + memtable: Some(MemtableOptions::TimeSeries), + ..Default::default() + }; + + let err = ensure_json2_not_use_time_series_memtable(&metadata, &options).unwrap_err(); + assert!( + err.to_string() + .contains("JSON2 columns only support BulkMemtable") + ); + + options.memtable = Some(MemtableOptions::Bulk(BulkMemtableConfig::default())); + ensure_json2_not_use_time_series_memtable(&metadata, &options)?; + Ok(()) + } } diff --git a/src/mito2/src/memtable/builder.rs b/src/mito2/src/memtable/builder.rs index 7e37077f3e..d072b85679 100644 --- a/src/mito2/src/memtable/builder.rs +++ b/src/mito2/src/memtable/builder.rs @@ -22,8 +22,8 @@ use datatypes::arrow::array::{ }; use datatypes::arrow::buffer::Buffer; use datatypes::arrow_array::StringArray; -use datatypes::data_type::DataType; use datatypes::prelude::{ConcreteDataType, MutableVector, VectorRef}; +use datatypes::schema::ColumnSchema; use datatypes::value::ValueRef; use datatypes::vectors::StringVector; @@ -35,11 +35,11 @@ pub(crate) enum FieldBuilder { impl FieldBuilder { /// Creates a [FieldBuilder] instance with given type and capacity. - pub fn create(data_type: &ConcreteDataType, init_cap: usize) -> Self { - if let ConcreteDataType::String(_) = data_type { + pub(crate) fn create(column_schema: &ColumnSchema, init_cap: usize) -> Self { + if let ConcreteDataType::String(_) = &column_schema.data_type { Self::String(StringBuilder::with_capacity(init_cap / 16, init_cap)) } else { - Self::Other(data_type.create_mutable_vector(init_cap)) + Self::Other(column_schema.create_mutable_vector(init_cap)) } } diff --git a/src/mito2/src/memtable/bulk.rs b/src/mito2/src/memtable/bulk.rs index 596db7aa34..1b67d487a5 100644 --- a/src/mito2/src/memtable/bulk.rs +++ b/src/mito2/src/memtable/bulk.rs @@ -16,7 +16,6 @@ pub(crate) mod chunk_reader; pub mod context; -pub(crate) mod json_align; pub mod part; pub mod part_reader; mod row_group_reader; @@ -44,10 +43,12 @@ use store_api::metadata::RegionMetadataRef; use store_api::storage::{ColumnId, FileId, RegionId, SequenceRange}; use tokio::sync::Semaphore; +use crate::compaction::{ + Json2RewritePlans, collect_json2_rewrite_plans, rewrite_json2_batch, rewrite_json2_schema, +}; use crate::error::{Result, UnsupportedOperationSnafu}; use crate::flush::WriteBufferManagerRef; use crate::memtable::bulk::context::BulkIterContext; -use crate::memtable::bulk::json_align::Json2Aligner; use crate::memtable::bulk::part::{ BulkPart, BulkPartEncodeMetrics, BulkPartEncoder, MultiBulkPart, UnorderedPart, should_prune_bulk_part, @@ -464,7 +465,8 @@ impl Memtable for BulkMemtable { // Compacts unordered_part if the row or byte threshold is exceeded. if bulk_parts.should_compact_unordered_part(self.config.encode_bytes_threshold) - && let Some(bulk_part) = bulk_parts.unordered_part.to_bulk_part()? + && let Some(bulk_part) = + bulk_parts.unordered_part.to_bulk_part(&self.metadata)? { bulk_parts.parts.push(BulkPartWrapper { part: PartToMerge::Bulk { @@ -525,7 +527,8 @@ impl Memtable for BulkMemtable { // Adds range for unordered part if not empty if !bulk_parts.unordered_part.is_empty() - && let Some(unordered_bulk_part) = bulk_parts.unordered_part.to_bulk_part()? + && let Some(unordered_bulk_part) = + bulk_parts.unordered_part.to_bulk_part(&self.metadata)? { let part_stats = unordered_bulk_part.to_memtable_stats(&self.metadata); let range = MemtableRange::new( @@ -1116,8 +1119,9 @@ impl PartToMerge { fn create_iterator( self, context: Arc, + plans: Arc, ) -> Result> { - match self { + let iter = match self { PartToMerge::Bulk { part, .. } => { let series_count = part.estimated_series_count(); let iter = BulkPartBatchIter::from_single( @@ -1127,10 +1131,18 @@ impl PartToMerge { series_count, None, // No metrics for merging ); - Ok(Some(Box::new(iter) as BoxedRecordBatchIterator)) + Some(Box::new(iter) as BoxedRecordBatchIterator) } - PartToMerge::Multi { part, .. } => part.read(context, None, None), - PartToMerge::Encoded { part, .. } => part.read(context, None, None), + PartToMerge::Multi { part, .. } => part.read(context, None, None)?, + PartToMerge::Encoded { part, .. } => part.read(context, None, None)?, + }; + if plans.is_empty() { + Ok(iter) + } else { + Ok(iter.map(|x| { + Box::new(x.map(move |batch| rewrite_json2_batch(batch?, &plans))) + as BoxedRecordBatchIterator + })) } } } @@ -1289,19 +1301,37 @@ impl MemtableCompactor { batch_size, )?); - let aligner = Json2Aligner::try_new(parts_to_merge.iter().map(PartToMerge::arrow_schema))?; + let schemas = parts_to_merge + .iter() + .map(|part| (part.arrow_schema(), part.num_rows() as u64)) + .collect::>(); + let plans = Arc::new(collect_json2_rewrite_plans(metadata, &schemas)?); + + debug_assert!(parts_to_merge.windows(2).all(|w| rewrite_json2_schema( + &w[0].arrow_schema(), + &plans + ) == rewrite_json2_schema( + &w[1].arrow_schema(), + &plans + ))); + // Parts in one merge group may differ only in their JSON2 physical layouts. So every source + // schema is therefore a valid template for producing the final target schema that has + // the union JSON2 types (rewritten). + let schema = rewrite_json2_schema(&parts_to_merge[0].arrow_schema(), &plans); let iterators: Vec = parts_to_merge .into_iter() - .filter_map(|part| part.create_iterator(context.clone()).ok().flatten()) - .map(|iter| aligner.wrap_iter(iter)) + .map(|part| part.create_iterator(context.clone(), plans.clone())) + .collect::>>()? + .into_iter() + .flatten() .collect(); if iterators.is_empty() { return Ok(None); } - let merged_iter = FlatMergeIterator::new(aligner.schema().clone(), iterators, batch_size)?; + let merged_iter = FlatMergeIterator::new(schema.clone(), iterators, batch_size)?; let boxed_iter: BoxedRecordBatchIterator = if dedup { match merge_mode { @@ -1310,8 +1340,7 @@ impl MemtableCompactor { Box::new(dedup_iter) } MergeMode::LastNonNull => { - let field_column_start = - field_column_start(metadata, aligner.schema().fields().len()); + let field_column_start = field_column_start(metadata, schema.fields().len()); let dedup_iter = FlatDedupIterator::new( merged_iter, @@ -1332,7 +1361,7 @@ impl MemtableCompactor { let mut metrics = BulkPartEncodeMetrics::default(); let encoded_part = encoder.encode_record_batch_iter( boxed_iter, - aligner.schema().clone(), + schema, min_timestamp, max_timestamp, max_sequence, @@ -1532,9 +1561,14 @@ mod tests { use api::helper::encode_json_value; use api::v1::value::ValueData; use api::v1::{Mutation, Row, Rows, SemanticType}; + use common_error::ext::WhateverResult; + use datatypes::arrow::datatypes::DataType as ArrowDataType; use datatypes::data_type::ConcreteDataType; - use datatypes::extension::json::Json2ExtensionType; + use datatypes::extension::json::{ + JSON2_REMAINDER_FIELD_NAME, Json2ExtensionType, Json2PhysicalLayout, JsonMetadata, + }; use datatypes::json::value::JsonValue; + use datatypes::json::{JsonSettings, JsonTypeHint}; use datatypes::schema::ColumnSchema; use datatypes::types::json_type::{JsonNativeType, JsonObjectType}; use mito_codec::row_converter::build_primary_key_codec; @@ -1678,7 +1712,7 @@ mod tests { #[test] fn test_bulk_memtable_compact_parts_with_json2() { - let metadata = mock_metadata_with_json2(); + let metadata = mock_metadata_with_json2(JsonSettings::default()); let config = BulkMemtableConfig { merge_threshold: 2, @@ -1717,7 +1751,126 @@ mod tests { assert_eq!(4, total_rows); } - fn mock_metadata_with_json2() -> RegionMetadataRef { + #[test] + fn test_bulk_memtable_merge_bounds_json2_paths() -> WhateverResult<()> { + let metadata = mock_metadata_with_json2(JsonSettings::try_new(vec![], Some(1))?); + let first = mock_bulk_part_with_json2_values( + &metadata, + vec![1000, 2000], + vec![json!({"a": 1}), json!({"a": 2})], + 100, + )?; + let second = mock_bulk_part_with_json2_values( + &metadata, + vec![3000, 4000], + vec![json!({"b": 3}), json!({"b": 4})], + 200, + )?; + let parts = vec![ + PartToMerge::Bulk { + part: first, + file_id: FileId::random(), + }, + PartToMerge::Bulk { + part: second, + file_id: FileId::random(), + }, + ]; + + let merged = MemtableCompactor::merge_parts_group( + parts, + &metadata, + false, + MergeMode::LastRow, + usize::MAX, + usize::MAX, + DEFAULT_ROW_GROUP_SIZE, + )? + .unwrap(); + let MergedPart::Multi(part) = merged else { + unreachable!() + }; + let schema = part.schemas().next().unwrap(); + let ArrowDataType::Struct(fields) = schema.field(0).data_type() else { + unreachable!() + }; + assert_eq!( + vec![JSON2_REMAINDER_FIELD_NAME, "a"], + fields.iter().map(|x| x.name().as_str()).collect::>() + ); + Ok(()) + } + + #[test] + fn test_unordered_parts_align_json2_layouts() -> WhateverResult<()> { + let metadata = mock_metadata_with_json2(JsonSettings::try_new(vec![], Some(2))?); + let memtable = BulkMemtable::new( + 42, + BulkMemtableConfig::default(), + metadata.clone(), + None, + None, + false, + MergeMode::LastRow, + ); + memtable.write_bulk(mock_bulk_part_with_json2_values( + &metadata, + vec![1000, 2000], + vec![json!({"a": 1}), json!({"a": 2})], + 100, + )?)?; + memtable.write_bulk(mock_bulk_part_with_json2_values( + &metadata, + vec![3000, 4000], + vec![json!({"b": 3}), json!({"b": 4})], + 200, + )?)?; + + let predicate = PredicateGroup::new(&metadata, &[])?; + let ranges = memtable.ranges(None, RangesOptions::default().with_predicate(predicate))?; + let range = ranges.ranges.values().next().unwrap(); + let batch = range.build_record_batch_iter(None, None)?.next().unwrap()?; + let schema = batch.schema(); + let ArrowDataType::Struct(fields) = schema.field(0).data_type() else { + unreachable!() + }; + assert_eq!( + vec![JSON2_REMAINDER_FIELD_NAME, "a", "b"], + fields.iter().map(|x| x.name().as_str()).collect::>() + ); + Ok(()) + } + + #[test] + fn test_bulk_part_converter_uses_json2_v2_layout() -> WhateverResult<()> { + let settings = JsonSettings::try_new( + vec![JsonTypeHint { + path: vec!["id".to_string()], + data_type: ConcreteDataType::int64_datatype(), + nullable: true, + default_constraint: None, + inverted_index: false, + }], + Some(0), + )?; + let metadata = mock_metadata_with_json2(settings); + let part = mock_bulk_part_with_json2(&metadata, vec![1000, 2000], 100)?; + let schema = part.batch.schema(); + let field = schema.field(0); + let layout = Json2PhysicalLayout::try_from_root(field)?; + + assert!(layout.is_version_2()); + let ArrowDataType::Struct(fields) = field.data_type() else { + unreachable!() + }; + assert_eq!( + vec![JSON2_REMAINDER_FIELD_NAME, "id"], + fields.iter().map(|x| x.name().as_str()).collect::>() + ); + Ok(()) + } + + fn mock_metadata_with_json2(settings: JsonSettings) -> RegionMetadataRef { let col_meta_1 = ColumnMetadata { column_schema: ColumnSchema::new( "ts", @@ -1730,7 +1883,9 @@ mod tests { let data_type = ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())); let mut col_schema = ColumnSchema::new("data", data_type, true); - col_schema.with_extension_type(&Json2ExtensionType::default()); + col_schema.with_extension_type(&Json2ExtensionType::new(Arc::new(JsonMetadata::new( + settings, + )))); let col_meta_2 = ColumnMetadata { column_schema: col_schema, @@ -1749,8 +1904,29 @@ mod tests { metadata: &RegionMetadataRef, timestamps: Vec, sequence: u64, + ) -> Result { + let values = timestamps + .iter() + .map(|ts| { + json!({ + "id": ts, + "payload": { + "message": format!("row-{ts}"), + }, + }) + }) + .collect(); + mock_bulk_part_with_json2_values(metadata, timestamps, values, sequence) + } + + fn mock_bulk_part_with_json2_values( + metadata: &RegionMetadataRef, + timestamps: Vec, + values: Vec, + sequence: u64, ) -> Result { let capacity = timestamps.len(); + debug_assert_eq!(capacity, values.len()); let primary_key_codec = build_primary_key_codec(metadata); let json_type = JsonNativeType::Object(JsonObjectType::from([ ("id".to_string(), JsonNativeType::i64()), @@ -1773,16 +1949,12 @@ mod tests { let rows = timestamps .into_iter() - .map(|ts| { + .zip(values) + .map(|(ts, value)| { let val1 = api::v1::Value { value_data: Some(ValueData::TimestampMillisecondValue(ts)), }; - let value_data = ValueData::JsonValue(encode_json_value(JsonValue::from(json!({ - "id": ts, - "payload": { - "message": format!("row-{ts}"), - }, - })))); + let value_data = ValueData::JsonValue(encode_json_value(JsonValue::from(value))); let val2 = api::v1::Value { value_data: Some(value_data), }; diff --git a/src/mito2/src/memtable/bulk/json_align.rs b/src/mito2/src/memtable/bulk/json_align.rs deleted file mode 100644 index 3a5c6e7fbc..0000000000 --- a/src/mito2/src/memtable/bulk/json_align.rs +++ /dev/null @@ -1,451 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::collections::HashMap; -use std::sync::Arc; - -use datatypes::arrow::datatypes::{DataType as ArrowDataType, Schema, SchemaRef}; -use datatypes::arrow::record_batch::RecordBatch; -use datatypes::extension::json::is_json2_extension_type; -use datatypes::types::json_type::JsonNativeType; -use datatypes::vectors::json::array::JsonArray; -use snafu::{OptionExt, ResultExt}; - -use crate::error::{ - ConvertValueSnafu, DataTypeMismatchSnafu, NewRecordBatchSnafu, Result, UnexpectedSnafu, -}; -use crate::memtable::BoxedRecordBatchIterator; - -/// Aligns concrete JSON2 Arrow types across record batches. -/// -/// JSON2 column concrete Arrow types are derived from data. Different memtable -/// parts may therefore have different concrete types for the same JSON2 column. -/// This helper merges those concrete types and aligns batches to the merged schema. -#[derive(Clone)] -pub(crate) struct Json2Aligner { - /// Schema after merging all JSON2 column concrete types. - schema: SchemaRef, - /// JSON2 columns that may need per-batch alignment. - json_columns: Vec<(usize, ArrowDataType)>, -} - -impl Json2Aligner { - /// Builds an aligner from input schemas. - /// - /// Note: except for JSON2 columns, all input schemas must be identical. - pub(crate) fn try_new(input_schemas: I) -> Result - where - I: IntoIterator, - { - let mut input_schemas = input_schemas.into_iter(); - - // Use first schema as base: it defines column order and non-JSON types. - let base_schema = input_schemas.next().context(UnexpectedSnafu { - reason: "Json2Aligner requires at least one input schema", - })?; - - // Init merged types from base schema. - let mut merged_types = base_schema - .fields() - .iter() - .enumerate() - .filter(|&(_idx, field)| is_json2_extension_type(field)) - .map(|(idx, field)| { - let json_type = - JsonNativeType::try_from(field.data_type()).context(DataTypeMismatchSnafu)?; - Ok((idx, json_type)) - }) - .collect::>>()?; - - // No JSON2 columns, no alignment needed. - if merged_types.is_empty() { - return Ok(Self { - schema: base_schema, - json_columns: Vec::new(), - }); - } - - // Merge JSON2 types from remaining schemas. - for schema in input_schemas { - // Input schemas should only differ in JSON2 concrete types. - #[cfg(debug_assertions)] - assert_columns_match_except_json2(&base_schema, &schema); - - for (idx, merged) in &mut merged_types { - if *idx >= schema.fields().len() { - continue; - } - let json_type = JsonNativeType::try_from(schema.field(*idx).data_type()) - .context(DataTypeMismatchSnafu)?; - merged.merge(&json_type); - } - } - - // Build output schema with merged JSON2 types. - let mut json_columns = Vec::with_capacity(merged_types.len()); - let fields: Vec<_> = base_schema - .fields() - .iter() - .enumerate() - .map(|(idx, field)| { - if let Some(merged) = merged_types.get(&idx) { - let data_type = merged.as_arrow_type(); - json_columns.push((idx, data_type.clone())); - let mut field = (**field).clone(); - field.set_data_type(data_type); - Arc::new(field) - } else { - field.clone() - } - }) - .collect(); - - let schema = Arc::new(Schema::new_with_metadata( - fields, - base_schema.metadata().clone(), - )); - - Ok(Self { - schema, - json_columns, - }) - } - - /// Returns the aligned output schema. - pub(crate) fn schema(&self) -> &SchemaRef { - &self.schema - } - - /// Aligns a [`RecordBatch`] to [`Self::schema`]. - pub(crate) fn align_batch(&self, batch: RecordBatch) -> Result { - if self.json_columns.is_empty() { - return Ok(batch); - } - let mut cols = batch.columns().to_vec(); - for (idx, expected_type) in &self.json_columns { - if batch.schema_ref().field(*idx).data_type() != expected_type { - cols[*idx] = JsonArray::from(batch.column(*idx)) - .widen_to(expected_type) - .context(ConvertValueSnafu)?; - } - } - RecordBatch::try_new(self.schema.clone(), cols).context(NewRecordBatchSnafu) - } - - /// Aligns [`RecordBatch`]s to [`Self::schema`]. - pub(crate) fn align_batches(&self, batches: I) -> Result> - where - I: IntoIterator, - { - batches - .into_iter() - .map(|batch| self.align_batch(batch)) - .collect() - } - - /// Wraps an iterator so each yielded [`RecordBatch`] is lazily aligned. - pub(crate) fn wrap_iter(&self, iter: BoxedRecordBatchIterator) -> BoxedRecordBatchIterator { - let aligner = self.clone(); - Box::new(iter.map(move |batch| aligner.align_batch(batch?))) - } -} - -#[cfg(debug_assertions)] -fn assert_columns_match_except_json2(base_schema: &Schema, schema: &Schema) { - debug_assert_eq!( - base_schema.fields().len(), - schema.fields().len(), - "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_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" - ); - if !base_is_json2 && !is_json2 { - debug_assert_eq!( - base_field, field, - "non-JSON2 column {idx} must be identical across input schemas" - ); - } - } -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use datatypes::arrow::array::{ - Array, ArrayRef, AsArray, Int64Array, StringViewArray, StructArray, UInt64Array, - }; - use datatypes::arrow::datatypes::{DataType, Field, Fields, Schema}; - use datatypes::extension::json::{Json2ExtensionType, JsonExtensionType}; - use serde_json::json; - - use super::*; - - #[test] - fn test_try_new_rejects_empty_input() { - let err = match Json2Aligner::try_new([]) { - Ok(_) => panic!("expected empty input to fail"), - Err(err) => err, - }; - assert!( - err.to_string() - .contains("Json2Aligner requires at least one input schema") - ); - } - - #[test] - fn test_try_new_keeps_non_json_schema_unchanged() { - let schema = Arc::new(Schema::new(vec![ - Arc::new(Field::new("ts", DataType::Int64, false)), - Arc::new(Field::new("value", DataType::UInt64, true)), - ])); - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int64Array::from_iter_values([1, 2])) as ArrayRef, - Arc::new(UInt64Array::from(vec![Some(10), None])) as ArrayRef, - ], - ) - .unwrap(); - - let aligner = Json2Aligner::try_new([schema.clone()]).unwrap(); - assert!(Arc::ptr_eq(aligner.schema(), &schema)); - - let aligned = aligner.align_batch(batch).unwrap(); - assert!(Arc::ptr_eq(aligned.schema_ref(), &schema)); - } - - #[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), - ); - let schema = Arc::new(Schema::new(vec![ - Arc::new(Field::new("ts", DataType::Int64, false)), - legacy_jsonb_field, - ])); - - let aligner = Json2Aligner::try_new([schema.clone()]).unwrap(); - - assert!(Arc::ptr_eq(aligner.schema(), &schema)); - assert!(aligner.json_columns.is_empty()); - } - - #[test] - fn test_try_new_merges_json2_object_fields() { - let id_fields = Fields::from(vec![id_field()]); - let name_fields = Fields::from(vec![name_field()]); - let schema_with_id = schema_with_json_field(json_field("data", id_fields)); - let schema_with_name = schema_with_json_field(json_field("data", name_fields)); - - let aligner = Json2Aligner::try_new([schema_with_id, schema_with_name]).unwrap(); - let data_field = aligner.schema().field(1); - let DataType::Struct(fields) = data_field.data_type() else { - panic!("expected JSON2 field to be a struct"); - }; - - assert_eq!(2, fields.len()); - assert_eq!("id", fields[0].name()); - assert_eq!(&DataType::Int64, fields[0].data_type()); - assert_eq!("name", fields[1].name()); - assert_eq!(&DataType::Utf8View, fields[1].data_type()); - assert!(is_json2_extension_type(&aligner.schema().fields()[1])); - } - - #[test] - fn test_align_batch_fills_missing_json2_fields() { - let id_fields = Fields::from(vec![id_field()]); - let name_fields = Fields::from(vec![name_field()]); - let schema_with_id = schema_with_json_field(json_field("data", id_fields.clone())); - let schema_with_name = schema_with_json_field(json_field("data", name_fields.clone())); - - let batch_with_id = RecordBatch::try_new( - schema_with_id.clone(), - vec![ - Arc::new(Int64Array::from_iter_values([1, 2])) as ArrayRef, - struct_array( - id_fields, - vec![Arc::new(Int64Array::from_iter_values([10, 20])) as ArrayRef], - ), - ], - ) - .unwrap(); - let batch_with_name = RecordBatch::try_new( - schema_with_name.clone(), - vec![ - Arc::new(Int64Array::from_iter_values([3, 4])) as ArrayRef, - struct_array( - name_fields, - vec![ - Arc::new(StringViewArray::from(vec![Some("alice"), Some("bob")])) - as ArrayRef, - ], - ), - ], - ) - .unwrap(); - - let aligner = Json2Aligner::try_new([schema_with_id, schema_with_name]).unwrap(); - let aligned_with_id = aligner.align_batch(batch_with_id).unwrap(); - let aligned_with_name = aligner.align_batch(batch_with_name).unwrap(); - - let data_with_id = aligned_with_id - .column(1) - .as_any() - .downcast_ref::() - .unwrap(); - let id_values = data_with_id - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - let missing_names = data_with_id.column(1); - assert_eq!(10, id_values.value(0)); - assert_eq!(20, id_values.value(1)); - assert!(missing_names.is_null(0)); - assert!(missing_names.is_null(1)); - - let data_with_name = aligned_with_name - .column(1) - .as_any() - .downcast_ref::() - .unwrap(); - let missing_ids = data_with_name.column(0); - let name_values = data_with_name - .column(1) - .as_any() - .downcast_ref::() - .unwrap(); - assert!(missing_ids.is_null(0)); - assert!(missing_ids.is_null(1)); - assert_eq!("alice", name_values.value(0)); - assert_eq!("bob", name_values.value(1)); - } - - #[test] - fn test_align_conflicting_number_types_as_variant() { - let u64_fields = Fields::from(vec![Arc::new(Field::new("value", DataType::UInt64, true))]); - let i64_fields = Fields::from(vec![Arc::new(Field::new("value", DataType::Int64, true))]); - let u64_schema = schema_with_json_field(json_field("data", u64_fields.clone())); - let i64_schema = schema_with_json_field(json_field("data", i64_fields.clone())); - let u64_batch = RecordBatch::try_new( - u64_schema.clone(), - vec![ - Arc::new(Int64Array::from_iter_values([1])) as ArrayRef, - struct_array( - u64_fields, - vec![Arc::new(UInt64Array::from_iter_values([u64::MAX])) as ArrayRef], - ), - ], - ) - .unwrap(); - let i64_batch = RecordBatch::try_new( - i64_schema.clone(), - vec![ - Arc::new(Int64Array::from_iter_values([2])) as ArrayRef, - struct_array( - i64_fields, - vec![Arc::new(Int64Array::from_iter_values([i64::MIN])) as ArrayRef], - ), - ], - ) - .unwrap(); - - let aligner = Json2Aligner::try_new([u64_schema, i64_schema]).unwrap(); - let DataType::Struct(fields) = aligner.schema().field(1).data_type() else { - panic!("expected JSON2 field to be a struct"); - }; - assert_eq!(&DataType::Binary, fields[0].data_type()); - - for (batch, expected) in [(u64_batch, json!(u64::MAX)), (i64_batch, json!(i64::MIN))] { - let aligned = aligner.align_batch(batch).unwrap(); - let data = aligned.column(1).as_struct(); - assert_eq!( - expected, - JsonArray::from(data.column(0)).try_get_value(0).unwrap() - ); - } - } - - #[test] - fn test_wrap_iter_aligns_each_batch() { - let id_fields = Fields::from(vec![id_field()]); - let name_fields = Fields::from(vec![name_field()]); - let schema_with_id = schema_with_json_field(json_field("data", id_fields.clone())); - let schema_with_name = schema_with_json_field(json_field("data", name_fields.clone())); - - let batch_with_id = RecordBatch::try_new( - schema_with_id.clone(), - vec![ - Arc::new(Int64Array::from_iter_values([1])) as ArrayRef, - struct_array( - id_fields, - vec![Arc::new(Int64Array::from_iter_values([10])) as ArrayRef], - ), - ], - ) - .unwrap(); - let batch_with_name = RecordBatch::try_new( - schema_with_name.clone(), - vec![ - Arc::new(Int64Array::from_iter_values([2])) as ArrayRef, - struct_array( - name_fields, - vec![Arc::new(StringViewArray::from(vec![Some("alice")])) as ArrayRef], - ), - ], - ) - .unwrap(); - - let aligner = Json2Aligner::try_new([schema_with_id, schema_with_name]).unwrap(); - let iter: BoxedRecordBatchIterator = - Box::new(vec![Ok(batch_with_id), Ok(batch_with_name)].into_iter()); - let aligned = aligner.wrap_iter(iter).collect::>>().unwrap(); - - assert_eq!(2, aligned.len()); - assert!(Arc::ptr_eq(aligned[0].schema_ref(), aligner.schema())); - assert!(Arc::ptr_eq(aligned[1].schema_ref(), aligner.schema())); - } - - fn json_field(name: &str, fields: Fields) -> Arc { - Arc::new( - Field::new(name, DataType::Struct(fields), true) - .with_extension_type(Json2ExtensionType::default()), - ) - } - - fn schema_with_json_field(json_field: Arc) -> SchemaRef { - Arc::new(Schema::new(vec![ - Arc::new(Field::new("ts", DataType::Int64, false)), - json_field, - ])) - } - - fn id_field() -> Arc { - Arc::new(Field::new("id", DataType::Int64, true)) - } - - fn name_field() -> Arc { - Arc::new(Field::new("name", DataType::Utf8View, true)) - } - - fn struct_array(fields: Fields, columns: Vec) -> ArrayRef { - Arc::new(StructArray::new(fields, columns, None)) - } -} diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index cc2c37a785..2a0c42b703 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -56,13 +56,13 @@ use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; use store_api::storage::{ColumnId, FileId, SequenceNumber, SequenceRange}; +use crate::compaction::{collect_json2_rewrite_plans, rewrite_json2_batch, rewrite_json2_schema}; use crate::error::{ self, ColumnNotFoundSnafu, ComputeArrowSnafu, CreateDefaultSnafu, DataTypeMismatchSnafu, EncodeMemtableSnafu, EncodeSnafu, InvalidMetadataSnafu, InvalidRequestSnafu, NewRecordBatchSnafu, Result, }; use crate::memtable::bulk::context::{BulkIterContext, BulkIterContextRef}; -use crate::memtable::bulk::json_align::Json2Aligner; use crate::memtable::bulk::part_reader::EncodedBulkPartIter; use crate::memtable::time_series::{ValueBuilder, Values}; use crate::memtable::{BoxedRecordBatchIterator, MemScanMetrics, MemtableStats}; @@ -442,7 +442,7 @@ impl UnorderedPart { /// Concatenates and sorts all parts into a single RecordBatch. /// Returns None if the collection is empty. - pub fn concat_and_sort(&self) -> Result> { + pub fn concat_and_sort(&self, metadata: &RegionMetadataRef) -> Result> { if self.parts.is_empty() { return Ok(None); } @@ -452,17 +452,28 @@ impl UnorderedPart { return Ok(Some(self.parts[0].batch.clone())); } - // Get the schema from the first part - let schema = self.parts[0].batch.schema(); - 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()))?; - concat_batches(aligner.schema(), &aligned_batches).context(ComputeArrowSnafu)? - } else { - concat_batches(&schema, self.parts.iter().map(|x| &x.batch)) - .context(ComputeArrowSnafu)? - }; + let schemas = self + .parts + .iter() + .map(|x| (x.batch.schema(), x.num_rows() as u64)) + .collect::>(); + let plans = collect_json2_rewrite_plans(metadata, &schemas)?; + + debug_assert!(self.parts.windows(2).all(|w| rewrite_json2_schema( + &w[0].batch.schema(), + &plans + ) == rewrite_json2_schema( + &w[1].batch.schema(), + &plans + ))); + let schema = rewrite_json2_schema(&self.parts[0].batch.schema(), &plans); + + let batches = self + .parts + .iter() + .map(|x| rewrite_json2_batch(x.batch.clone(), &plans)) + .collect::>>()?; + let concatenated = concat_batches(&schema, &batches).context(ComputeArrowSnafu)?; // Sort the concatenated batch let sorted_batch = sort_primary_key_record_batch(&concatenated)?; @@ -472,8 +483,8 @@ impl UnorderedPart { /// Converts all parts into a single sorted BulkPart. /// Returns None if the collection is empty. - pub fn to_bulk_part(&self) -> Result> { - let Some(sorted_batch) = self.concat_and_sort()? else { + pub fn to_bulk_part(&self, metadata: &RegionMetadataRef) -> Result> { + let Some(sorted_batch) = self.concat_and_sort(metadata)? else { return Ok(None); }; diff --git a/src/mito2/src/memtable/time_series.rs b/src/mito2/src/memtable/time_series.rs index 7e89d569ee..fa9970137c 100644 --- a/src/mito2/src/memtable/time_series.rs +++ b/src/mito2/src/memtable/time_series.rs @@ -29,6 +29,7 @@ use datatypes::arrow::array::ArrayRef; use datatypes::arrow_array::StringArray; use datatypes::data_type::{ConcreteDataType, DataType}; use datatypes::prelude::{ScalarVector, Vector, VectorRef}; +use datatypes::schema::ColumnSchema; use datatypes::types::TimestampType; use datatypes::value::{Value, ValueRef}; use datatypes::vectors::{ @@ -46,7 +47,7 @@ use crate::error::{ self, ComputeArrowSnafu, ConvertVectorSnafu, EncodeSnafu, PrimaryKeyLengthMismatchSnafu, Result, }; use crate::flush::WriteBufferManagerRef; -use crate::memtable::builder::{FieldBuilder, StringBuilder}; +use crate::memtable::builder::FieldBuilder; use crate::memtable::bulk::part::BulkPart; use crate::memtable::simple_bulk_memtable::SimpleBulkMemtable; use crate::memtable::stats::WriteMetrics; @@ -918,7 +919,7 @@ pub(crate) struct ValueBuilder { sequence: Vec, op_type: Vec, fields: Vec>, - field_types: Vec, + field_schemas: Vec, } impl ValueBuilder { @@ -931,18 +932,18 @@ impl ValueBuilder { let sequence = Vec::with_capacity(capacity); let op_type = Vec::with_capacity(capacity); - let field_types = region_metadata + let field_schemas = region_metadata .field_columns() - .map(|c| c.column_schema.data_type.clone()) + .map(|c| c.column_schema.clone()) .collect::>(); - let fields = (0..field_types.len()).map(|_| None).collect(); + let fields = (0..field_schemas.len()).map(|_| None).collect(); Self { timestamp: Vec::with_capacity(capacity), timestamp_type, sequence, op_type, fields, - field_types, + field_schemas, } } @@ -984,15 +985,10 @@ impl ValueBuilder { .push(field_value) .unwrap_or_else(|e| panic!("Failed to push field value: {e:?}")); } else { - let mut mutable_vector = - if let ConcreteDataType::String(_) = &self.field_types[idx] { - FieldBuilder::String(StringBuilder::with_capacity(4, 8)) - } else { - FieldBuilder::Other( - self.field_types[idx] - .create_mutable_vector(num_rows.max(INITIAL_BUILDER_CAPACITY)), - ) - }; + let mut mutable_vector = FieldBuilder::create( + &self.field_schemas[idx], + num_rows.max(INITIAL_BUILDER_CAPACITY), + ); mutable_vector.push_nulls(num_rows - 1); mutable_vector .push(field_value) @@ -1017,7 +1013,12 @@ impl ValueBuilder { /// the Arrow string array offset limit and thus can never be accommodated, not even by an /// empty builder. pub(crate) fn can_accommodate(&self, fields: &[VectorRef]) -> Result { - scan_string_capacity(fields, &self.fields, &self.field_types, i32::MAX) + let data_types = self + .field_schemas + .iter() + .map(|x| x.data_type.clone()) + .collect::>(); + scan_string_capacity(fields, &self.fields, &data_types, i32::MAX) } pub(crate) fn extend( @@ -1082,7 +1083,7 @@ impl ValueBuilder { { let builder = field_dest.get_or_insert_with(|| { let mut field_builder = - FieldBuilder::create(&self.field_types[field_idx], INITIAL_BUILDER_CAPACITY); + FieldBuilder::create(&self.field_schemas[field_idx], INITIAL_BUILDER_CAPACITY); field_builder.push_nulls(num_rows_before); field_builder }); @@ -1140,9 +1141,9 @@ impl ValueBuilder { MEMTABLE_ACTIVE_FIELD_BUILDER_COUNT.dec(); v.finish_cloned() } else { - let mut single_null = self.field_types[i].create_mutable_vector(num_rows); - single_null.push_nulls(num_rows); - single_null.to_vector() + let mut builder = FieldBuilder::create(&self.field_schemas[i], num_rows); + builder.push_nulls(num_rows); + builder.finish() } }) .collect::>(); @@ -1282,9 +1283,9 @@ impl From for Values { MEMTABLE_ACTIVE_FIELD_BUILDER_COUNT.dec(); v.finish() } else { - let mut single_null = value.field_types[i].create_mutable_vector(num_rows); - single_null.push_nulls(num_rows); - single_null.to_vector() + let mut builder = FieldBuilder::create(&value.field_schemas[i], num_rows); + builder.push_nulls(num_rows); + builder.finish() } }) .collect::>(); @@ -1390,6 +1391,7 @@ mod tests { use store_api::storage::RegionId; use super::*; + use crate::memtable::builder::StringBuilder; use crate::test_util::column_metadata_to_column_schema; fn schema_for_test() -> RegionMetadataRef { diff --git a/src/mito2/src/read/compat.rs b/src/mito2/src/read/compat.rs index 3fa625694c..b0c4d52f04 100644 --- a/src/mito2/src/read/compat.rs +++ b/src/mito2/src/read/compat.rs @@ -95,16 +95,11 @@ impl FlatCompatBatch { compaction: bool, ) -> Result> { let actual = read_format.metadata(); - let format_projection = read_format.format_projection(); - let mut actual_schema = flat_projected_columns(actual, format_projection); - for (column_id, target_type) in read_format.json_target_types().iter() { - if let Some(i) = actual_schema - .iter() - .position(|(actual_column_id, _)| actual_column_id == column_id) - { - actual_schema[i].1 = ConcreteDataType::json2(target_type.clone()); - } - } + let actual_schema = flat_projected_columns( + actual, + read_format.format_projection(), + read_format.json_target_types(), + ); let expect_schema = mapper.batch_schema(); if expect_schema == actual_schema @@ -176,7 +171,19 @@ impl FlatCompatBatch { // Same column different type. if expect_data_type != *actual_data_type { - cast_type = Some(expect_data_type.clone()) + ensure!( + !expect_data_type.is_json2() && !actual_data_type.is_json2(), + CompatReaderSnafu { + region_id: expect_metadata.region_id, + reason: format!( + "JSON2 column '{}' must be aligned before FlatCompatBatch, actual: {}, expected: {}", + expect_column.column_schema.name, + actual_data_type, + expect_data_type, + ), + } + ); + cast_type = Some(expect_data_type.clone()); } // Source has this column. index_or_defaults.push(IndexOrDefault::Index { @@ -634,6 +641,7 @@ impl FlatCompatPrimaryKey { #[cfg(test)] mod tests { + use std::collections::BTreeMap; use std::sync::Arc; use api::v1::{OpType, SemanticType}; @@ -645,6 +653,7 @@ mod tests { use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnSchema; + use datatypes::types::json_type::JsonNativeType; use datatypes::value::ValueRef; use mito_codec::row_converter::{ DensePrimaryKeyCodec, PrimaryKeyCodecExt, SparsePrimaryKeyCodec, @@ -814,6 +823,57 @@ mod tests { assert_eq!(expected_batch, result); } + #[test] + fn test_flat_compat_batch_uses_projected_json2_type() -> Result<()> { + let json2 = ConcreteDataType::json2(JsonNativeType::object()); + let actual_metadata = Arc::new(new_metadata( + &[ + ( + 0, + SemanticType::Timestamp, + ConcreteDataType::timestamp_millisecond_datatype(), + ), + (1, SemanticType::Field, json2.clone()), + ], + &[], + )); + let expected_metadata = Arc::new(new_metadata( + &[ + ( + 0, + SemanticType::Timestamp, + ConcreteDataType::timestamp_millisecond_datatype(), + ), + (1, SemanticType::Field, json2), + (2, SemanticType::Field, ConcreteDataType::int64_datatype()), + ], + &[], + )); + let read_columns = ReadColumns::new([0, 1, 2]) + .with_json_target_types(BTreeMap::from([(1, JsonNativeType::Variant)])); + let mapper = FlatProjectionMapper::new_with_read_columns( + &expected_metadata, + vec![0, 1, 2], + read_columns.clone(), + )?; + let read_format = FlatReadFormat::new(actual_metadata, read_columns, None, "test", false)?; + + let compat = FlatCompatBatch::try_new(&mapper, &read_format, false)?.unwrap(); + let json_index = mapper + .batch_schema() + .iter() + .position(|(id, _)| *id == 1) + .unwrap(); + assert!(matches!( + &compat.index_or_defaults[json_index], + IndexOrDefault::Index { + cast_type: None, + .. + } + )); + Ok(()) + } + #[test] fn test_flat_compat_batch_with_read_projection_superset() { let actual_metadata = Arc::new(new_metadata( diff --git a/src/mito2/src/read/flat_projection.rs b/src/mito2/src/read/flat_projection.rs index 5178778ba9..19dc962c56 100644 --- a/src/mito2/src/read/flat_projection.rs +++ b/src/mito2/src/read/flat_projection.rs @@ -27,11 +27,10 @@ use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field}; use datatypes::extension::json::is_json2_extension_type; use datatypes::prelude::{ConcreteDataType, DataType}; use datatypes::schema::{Schema, SchemaRef}; -use datatypes::types::JsonType; -use datatypes::types::json_type::JsonNativeType; use datatypes::value::Value; use datatypes::vectors::Helper; use datatypes::vectors::json::array::JsonArray; +use datatypes::vectors::json::json2_physical_data_type; use snafu::{OptionExt, ResultExt}; use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::storage::ColumnId; @@ -39,7 +38,8 @@ use store_api::storage::ColumnId; use crate::cache::CacheStrategy; use crate::error::{InvalidRequestSnafu, RecordBatchSnafu, Result}; use crate::read::projection::{read_column_ids_from_projection, repeated_vector_with_cache}; -use crate::read::read_columns::ReadColumns; +use crate::read::read_columns::{JsonTargetTypes, ReadColumns}; +use crate::sst::parquet::Json2RewriteTargets; use crate::sst::parquet::flat_format::sst_column_id_indices; use crate::sst::parquet::format::FormatProjection; use crate::sst::{ @@ -94,6 +94,21 @@ impl FlatProjectionMapper { metadata: &RegionMetadataRef, projection: Vec, read_cols: ReadColumns, + ) -> Result { + Self::new_with_json2_rewrite_targets( + metadata, + projection, + read_cols, + &Json2RewriteTargets::default(), + ) + } + + /// Returns a mapper for a compaction read with fixed JSON2 output layouts. + pub(crate) fn new_with_json2_rewrite_targets( + metadata: &RegionMetadataRef, + projection: Vec, + read_cols: ReadColumns, + json2_rewrite_targets: &Json2RewriteTargets, ) -> Result { // If the original projection is empty. let is_empty_projection = projection.is_empty(); @@ -113,10 +128,8 @@ impl FlatProjectionMapper { output_col_ids.push(col.column_id); let mut schema = col.column_schema.clone(); - if let Some(data_type) = - json2_read_datatype(col.column_id, &schema.data_type, &read_cols) - { - schema.data_type = data_type; + if let Some(data_type) = read_cols.json_target_type(col.column_id) { + schema.data_type = ConcreteDataType::json2(data_type.clone()); } col_schemas.push(schema); } @@ -133,16 +146,11 @@ impl FlatProjectionMapper { read_cols.clone(), ); - let mut batch_schema = flat_projected_columns(metadata, &format_projection); + let batch_schema = + flat_projected_columns(metadata, &format_projection, read_cols.json_target_types()); - for (column_id, data_type) in batch_schema.iter_mut() { - if let Some(updated) = json2_read_datatype(*column_id, data_type, &read_cols) { - *data_type = updated; - } - } - - // Safety: We get the column id from the metadata. - let input_arrow_schema = compute_input_arrow_schema(metadata, &batch_schema); + let input_arrow_schema = + compute_input_arrow_schema(metadata, &batch_schema, &read_cols, json2_rewrite_targets); // If projection is empty, we don't output any column. let output_schema = if is_empty_projection { @@ -391,35 +399,6 @@ impl FlatProjectionMapper { } } -fn json2_read_datatype( - column_id: ColumnId, - data_type: &ConcreteDataType, - read_cols: &ReadColumns, -) -> Option { - let json_type = data_type.as_json()?; - if !json_type.is_json2() { - return None; - } - - if let Some(concretized) = read_cols.json_target_type(column_id).cloned() { - return Some(ConcreteDataType::json2(concretized)); - } - - if is_empty_json2_type(json_type) { - return Some(ConcreteDataType::json2(JsonNativeType::Variant)); - } - - None -} - -fn is_empty_json2_type(json_type: &JsonType) -> bool { - match json_type.native_type() { - JsonNativeType::Null => true, - JsonNativeType::Object(fields) if fields.is_empty() => true, - _ => false, - } -} - fn single_value_string_dictionary<'a>( array: &'a Arc, output_type: &ConcreteDataType, @@ -442,12 +421,13 @@ fn single_value_string_dictionary<'a>( (dict_array.values().len() == 1 && dict_array.null_count() == 0).then_some(dict_array) } -/// Returns ids and datatypes of columns of the output batch after applying the `projection`. +/// Returns ids and datatypes of columns after applying the projection and JSON2 target types. /// /// It adds the time index column if it doesn't present in the projection. pub(crate) fn flat_projected_columns( metadata: &RegionMetadata, format_projection: &FormatProjection, + json_target_types: &JsonTargetTypes, ) -> Vec<(ColumnId, ConcreteDataType)> { let time_index = metadata.time_index_column(); let num_columns = if format_projection @@ -460,16 +440,18 @@ pub(crate) fn flat_projected_columns( }; let mut schema = vec![None; num_columns]; for (column_id, index) in &format_projection.column_id_to_projected_index { - // Safety: FormatProjection ensures the id is valid. - schema[*index] = Some(( - *column_id, + let data_type = if let Some(json_type) = json_target_types.get(column_id) { + ConcreteDataType::json2(json_type.clone()) + } else { + // Safety: FormatProjection ensures the id is valid. metadata .column_by_id(*column_id) .unwrap() .column_schema .data_type - .clone(), - )); + .clone() + }; + schema[*index] = Some((*column_id, data_type)); } if num_columns != format_projection.column_id_to_projected_index.len() { schema[num_columns - 1] = Some(( @@ -489,13 +471,25 @@ pub(crate) fn flat_projected_columns( pub(crate) fn compute_input_arrow_schema( metadata: &RegionMetadata, batch_schema: &[(ColumnId, ConcreteDataType)], + read_cols: &ReadColumns, + json2_rewrite_targets: &Json2RewriteTargets, ) -> datatypes::arrow::datatypes::SchemaRef { let mut new_fields = Vec::with_capacity(batch_schema.len() + 3); for (column_id, data_type) in batch_schema { + let data_type = json2_rewrite_targets + .get(column_id) + .map(|x| json2_physical_data_type(&x.target_layout)) + .or_else(|| { + read_cols + .json_target_type(*column_id) + .map(|x| x.as_arrow_type()) + }) + .unwrap_or_else(|| data_type.as_arrow_type()); + let column_metadata = metadata.column_by_id(*column_id).unwrap(); let field = Field::new( &column_metadata.column_schema.name, - data_type.as_arrow_type(), + data_type, column_metadata.column_schema.is_nullable(), ) .with_metadata(column_metadata.column_schema.metadata().clone()); diff --git a/src/mito2/src/read/read_columns.rs b/src/mito2/src/read/read_columns.rs index 42a9f71ddc..89e99fa3ef 100644 --- a/src/mito2/src/read/read_columns.rs +++ b/src/mito2/src/read/read_columns.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::BTreeMap; +use std::hash::Hash; use std::mem; use std::sync::Arc; @@ -46,6 +47,7 @@ impl ReadColumns { } } + /// Attaches query-time JSON2 projection types. pub fn with_json_target_types( mut self, json_target_types: BTreeMap, @@ -66,10 +68,11 @@ impl ReadColumns { self.column_ids_iter().collect() } - pub fn json_target_types(&self) -> &JsonTargetTypes { + pub(crate) fn json_target_types(&self) -> &JsonTargetTypes { &self.json_target_types } + /// Returns the query-time JSON2 projection type for a column. pub fn json_target_type(&self, column_id: ColumnId) -> Option<&JsonNativeType> { self.json_target_types.get(&column_id) } @@ -77,7 +80,6 @@ impl ReadColumns { pub fn estimated_size(&self) -> usize { self.col_ids.capacity() * mem::size_of::() + self.col_ids.len() * mem::size_of::() - + self.json_target_types.len() - * (mem::size_of::() + mem::size_of::()) + + self.json_target_types.len() * (size_of::() + size_of::()) } } diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index df199d4bac..72f1f8df37 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -83,6 +83,7 @@ use crate::sst::index::inverted_index::applier::InvertedIndexApplierRef; use crate::sst::index::inverted_index::applier::builder::InvertedIndexApplierBuilder; #[cfg(feature = "vector_index")] use crate::sst::index::vector_index::applier::{VectorIndexApplier, VectorIndexApplierRef}; +use crate::sst::parquet::Json2RewriteTargets; use crate::sst::parquet::file_range::PreFilterMode; use crate::sst::parquet::reader::ReaderMetrics; @@ -952,6 +953,8 @@ pub struct ScanInput { pub(crate) snapshot_sequence: Option, /// Whether this scan is for compaction. pub(crate) compaction: bool, + /// Compaction-only JSON2 physical rewrite targets. + json2_rewrite_targets: Json2RewriteTargets, /// Counters that should receive query-load metrics. pub(crate) query_stat_counters: Option, #[cfg(feature = "enterprise")] @@ -992,6 +995,7 @@ impl ScanInput { explain_flat_format: false, snapshot_sequence: None, compaction: false, + json2_rewrite_targets: Arc::default(), query_stat_counters: None, #[cfg(feature = "enterprise")] extension_ranges: Vec::new(), @@ -1197,6 +1201,13 @@ impl ScanInput { self } + /// Sets compaction-only JSON2 physical rewrite targets. + #[must_use] + pub(crate) fn with_json2_rewrite_targets(mut self, targets: Json2RewriteTargets) -> Self { + self.json2_rewrite_targets = targets; + self + } + /// Builds memtable ranges to scan by `index`. pub(crate) fn build_mem_ranges(&self, index: RowGroupIndex) -> SmallVec<[MemtableRange; 2]> { let memtable = &self.memtables[index.index]; @@ -1325,6 +1336,7 @@ impl ScanInput { .read_sst(file.clone()) .predicate(predicate) .projection(Some(self.read_cols.clone())) + .json2_rewrite_targets(self.json2_rewrite_targets.clone()) .cache(self.cache_strategy.clone()) .inverted_index_appliers(self.inverted_index_appliers.clone()) .bloom_filter_index_appliers(self.bloom_filter_index_appliers.clone()) diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 23695ca424..ab2e0c737d 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -20,8 +20,10 @@ use std::sync::atomic::{AtomicI64, AtomicU64}; use std::sync::{Arc, LazyLock}; use std::time::Instant; +use arrow_schema::extension::ExtensionType; use common_telemetry::{debug, error, info, warn}; use common_wal::options::WalOptions; +use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; use futures::StreamExt; use futures::future::BoxFuture; use log_store::kafka::log_store::KafkaLogStore; @@ -50,14 +52,14 @@ use crate::config::MitoConfig; use crate::engine::region_hook::RegionHookRef; use crate::error; use crate::error::{ - EmptyRegionDirSnafu, InvalidMetadataSnafu, InvalidRegionOptionsSnafu, ObjectStoreNotFoundSnafu, - RegionCorruptedSnafu, Result, StaleLogEntrySnafu, + DataTypeMismatchSnafu, EmptyRegionDirSnafu, InvalidMetadataSnafu, InvalidRegionOptionsSnafu, + ObjectStoreNotFoundSnafu, RegionCorruptedSnafu, Result, StaleLogEntrySnafu, }; use crate::manifest::action::RegionManifest; use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions}; -use crate::memtable::MemtableBuilderProvider; use crate::memtable::bulk::part::BulkPart; use crate::memtable::time_partition::{TimePartitions, TimePartitionsRef}; +use crate::memtable::{MemtableBuilderProvider, ensure_json2_not_use_time_series_memtable}; use crate::metrics::{CACHE_FILL_DOWNLOADED_FILES, CACHE_FILL_PENDING_FILES}; use crate::region::options::RegionOptions; use crate::region::version::{VersionBuilder, VersionControl, VersionControlRef}; @@ -93,6 +95,41 @@ fn initial_pruned_entry_id(wal_options: &WalOptions) -> EntryId { } } +fn maybe_upgrade_json2_layout(metadata: RegionMetadataRef) -> Result { + let mut upgrades = Vec::new(); + for (index, column) in metadata.column_metadatas.iter().enumerate() { + if !column.column_schema.data_type.is_json2() { + continue; + } + let Some(extension) = column + .column_schema + .extension_type::() + .context(DataTypeMismatchSnafu)? + else { + continue; + }; + if extension.metadata().is_version_2() { + continue; + } + upgrades.push((index, extension.metadata().json_settings().clone())); + } + + if upgrades.is_empty() { + return Ok(metadata); + } + + let mut upgraded = metadata.as_ref().clone(); + for (index, settings) in upgrades { + let extension = Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings))); + upgraded.column_metadatas[index] + .column_schema + .with_extension_type(&extension); + } + + let builder = RegionMetadataBuilder::from_existing(upgraded); + Ok(Arc::new(builder.build().context(InvalidMetadataSnafu)?)) +} + /// A fetcher to retrieve partition expr for a region. /// /// Compatibility: older regions didn't persist `partition_expr` in engine metadata, @@ -324,6 +361,7 @@ impl RegionOpener { options.sst_format = Some(FormatType::PrimaryKey); FormatType::PrimaryKey }; + ensure_json2_not_use_time_series_memtable(&metadata, &options)?; // Create a manifest manager for this region and writes regions to the manifest file. let mut region_manifest_options = RegionManifestOptions::new(config, ®ion_dir, &object_store); @@ -475,8 +513,10 @@ impl RegionOpener { } else { manifest.metadata.clone() }; + let metadata = maybe_upgrade_json2_layout(metadata)?; // Updates the region options with the manifest. sanitize_region_options(&manifest, &mut region_options); + ensure_json2_not_use_time_series_memtable(&metadata, ®ion_options)?; let region_id = self.region_id; let provider = self.provider::(®ion_options.wal_options)?; @@ -1319,23 +1359,31 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; + use arrow_schema::extension::ExtensionType; use common_base::readable_size::ReadableSize; + use common_error::ext::WhateverResult; use common_test_util::temp_dir::create_temp_dir; use common_time::Timestamp; use common_wal::options::{KafkaWalOptions, WalOptions}; use datatypes::arrow::array::{ArrayRef, BinaryArray, Int64Array}; use datatypes::arrow::record_batch::RecordBatch; + use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; + use datatypes::json::JsonSettings; + use datatypes::prelude::ConcreteDataType; + use datatypes::schema::ColumnSchema; + use datatypes::types::json_type::{JsonNativeType, JsonObjectType}; use object_store::ObjectStore; use object_store::services::{Fs, Memory, S3}; use parquet::arrow::ArrowWriter; use parquet::file::metadata::{KeyValue, PageIndexPolicy}; use parquet::file::properties::WriterProperties; + use store_api::metadata::RegionMetadataBuilder; use store_api::region_request::PathType; use store_api::storage::{FileId, RegionId}; use super::{ - initial_pruned_entry_id, preload_parquet_meta_cache_for_files, sanitize_region_options, - supports_open_region_object_storage_requirement, + initial_pruned_entry_id, maybe_upgrade_json2_layout, preload_parquet_meta_cache_for_files, + sanitize_region_options, supports_open_region_object_storage_requirement, }; use crate::cache::CacheManager; use crate::cache::file_cache::{FileType, IndexKey}; @@ -1389,6 +1437,38 @@ mod tests { ); } + #[test] + fn test_upgrade_json2_layout() -> WhateverResult<()> { + let settings = JsonSettings::try_new(vec![], Some(3))?; + let extension = Json2ExtensionType::new(Arc::new(JsonMetadata::new_v1(settings.clone()))); + let mut column = ColumnSchema::new( + "field_0", + ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())), + true, + ); + column.with_extension_type(&extension); + + let mut metadata = sst_region_metadata(); + metadata.column_metadatas[2].column_schema = column; + let builder = RegionMetadataBuilder::from_existing(metadata); + let metadata = Arc::new(builder.build()?); + + let upgraded = maybe_upgrade_json2_layout(metadata)?; + let column = &upgraded.column_metadatas[2].column_schema; + let extension = column.extension_type::()?.unwrap(); + assert!(extension.metadata().is_version_2()); + assert_eq!(&settings, extension.metadata().json_settings()); + + let arrow_schema = upgraded.schema.arrow_schema(); + let field = arrow_schema.field_with_name("field_0").unwrap(); + let extension = field.try_extension_type::().unwrap(); + assert!(extension.metadata().is_version_2()); + + let unchanged = maybe_upgrade_json2_layout(upgraded.clone())?; + assert!(Arc::ptr_eq(&upgraded, &unchanged)); + Ok(()) + } + #[test] #[cfg(not(feature = "test-shared-fs-region-migration"))] fn test_open_requirement_rejects_fs_object_store() { diff --git a/src/mito2/src/sst.rs b/src/mito2/src/sst.rs index f5e77e5be2..ba670ef734 100644 --- a/src/mito2/src/sst.rs +++ b/src/mito2/src/sst.rs @@ -506,13 +506,20 @@ impl SeriesEstimator { mod tests { use std::sync::Arc; + use ::parquet::arrow::AsyncArrowWriter; + use ::parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use ::parquet::basic::LogicalType; + use ::parquet::variant::{VariantArray, VariantType, json_to_variant}; use common_query::prelude::greptime_native_histogram; use datatypes::arrow::array::{ - BinaryArray, DictionaryArray, TimestampMillisecondArray, UInt8Array, UInt32Array, - UInt64Array, + ArrayRef, BinaryArray, DictionaryArray, Int64Array, StringArray, StructArray, + TimestampMillisecondArray, UInt8Array, UInt32Array, UInt64Array, }; use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field, Schema, TimeUnit}; use datatypes::arrow::record_batch::RecordBatch; + use datatypes::extension::json::{Json2ExtensionType, Json2PhysicalLayout}; + use datatypes::vectors::json::array::JsonArray; + use serde_json::json; use super::*; @@ -1005,4 +1012,125 @@ mod tests { err ); } + + fn json2_v2_test_type() -> ArrowDataType { + ArrowDataType::Struct( + vec![ + Arc::new(Field::new("active", ArrowDataType::Boolean, true)), + Arc::new(Field::new("hot", ArrowDataType::Int64, true)), + Arc::new(Field::new("name", ArrowDataType::Utf8, true)), + ] + .into(), + ) + } + + /// Validates the persisted-format foundation for the JSON2 v2 remainder. + /// + /// JSON2 will store `!__remainder__!` as a nested Variant child of its root + /// Struct. Before enabling that layout in production, this test ensures the + /// SST schema wrapper and Arrow writer preserve the Variant extension, + /// encode the Parquet Variant logical type, and round-trip the values + /// without changing the surrounding Struct. + #[tokio::test] + async fn test_nested_variant_survives_sst_writer_schema_roundtrip() + -> Result<(), Box> { + let json: ArrayRef = Arc::new(StringArray::from(vec![ + Some(r#"{}"#), + Some(r#"{"name":"Alice","active":true}"#), + Some(r#"{"nested":{"count":42},"items":[1,"two",null]}"#), + Some(r#"{"\u5b57\u6bb5":"\u503c"}"#), + None, + ])); + let remainder = json_to_variant(&json)?; + let remainder_field = remainder.field("!__remainder__!"); + let remainder_array = ArrayRef::from(remainder); + let hot_field = Field::new("hot", ArrowDataType::Int64, true); + let data_array = Arc::new(StructArray::new( + vec![remainder_field.clone(), hot_field.clone()].into(), + vec![ + remainder_array, + Arc::new(Int64Array::from(vec![ + Some(1), + Some(2), + Some(3), + Some(4), + None, + ])), + ], + None, + )); + let data_field = Field::new( + "data", + ArrowDataType::Struct(vec![remainder_field, hot_field].into()), + true, + ) + .with_extension_type(Json2ExtensionType::default()); + let schema = Arc::new(Schema::new(vec![data_field])); + let source = RecordBatch::try_new(schema.clone(), vec![data_array])?; + + let wrapped = maybe_wrap_schema(&schema)?; + let mut buffer = Vec::new(); + let mut writer = AsyncArrowWriter::try_new(&mut buffer, wrapped, None)?; + writer.write(&source).await?; + writer.close().await?; + + let builder = ParquetRecordBatchReaderBuilder::try_new(bytes::Bytes::from(buffer))?; + let parquet_remainder = + &builder.parquet_schema().root_schema().get_fields()[0].get_fields()[0]; + assert_eq!( + parquet_remainder.get_basic_info().logical_type_ref(), + Some(&LogicalType::Variant { + specification_version: None, + }) + ); + + let ArrowDataType::Struct(children) = builder.schema().field_with_name("data")?.data_type() + else { + unreachable!(); + }; + assert!(children[0].has_valid_extension_type::()); + + let mut reader = builder.build()?; + let result = reader.next().unwrap()?; + assert_eq!(source, result); + let result_field = result.schema().field(0).clone(); + let result = result + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + VariantArray::try_new(result.column(0))?; + let result: ArrayRef = Arc::new(result.clone()); + let result = + JsonArray::from(&result).project_to_v2(&result_field, &json2_v2_test_type())?; + assert_eq!( + json!({"active": true, "hot": 2, "name": "Alice"}), + JsonArray::from(&result).try_get_value(1)? + ); + Ok(()) + } + + /// Ensures future readers retain compatibility with the first JSON2 v2 layout. + #[test] + fn test_read_json2_v2_fixture() -> Result<(), Box> { + let bytes = bytes::Bytes::from_static(include_bytes!("../test-data/json2-v2.parquet")); + let builder = ParquetRecordBatchReaderBuilder::try_new(bytes)?; + let field = builder.schema().field(0).clone(); + assert!(Json2PhysicalLayout::try_from_root(&field)?.is_version_2()); + + let batch = builder.build()?.next().unwrap()?; + let data = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + VariantArray::try_new(data.column(0))?; + let data: ArrayRef = Arc::new(data.clone()); + let data = JsonArray::from(&data).project_to_v2(&field, &json2_v2_test_type())?; + assert_eq!( + json!({"active": true, "hot": 2, "name": "Alice"}), + JsonArray::from(&data).try_get_value(1)? + ); + Ok(()) + } } diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index 0603f9cf43..99646b7bf3 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -14,11 +14,13 @@ //! SST in parquet format. +use std::collections::BTreeMap; use std::sync::Arc; use common_base::readable_size::ReadableSize; +use datatypes::json::JsonSettings; use parquet::file::metadata::ParquetMetaData; -use store_api::storage::FileId; +use store_api::storage::{ColumnId, FileId}; use crate::sst::DEFAULT_WRITE_BUFFER_SIZE; use crate::sst::file::FileTimeRange; @@ -49,6 +51,19 @@ pub const PARQUET_METADATA_KEY: &str = "greptime:metadata"; /// default execution batch size to reduce rebatching and concatenation in the /// query pipeline. pub(crate) const DEFAULT_READ_BATCH_SIZE: usize = 8 * 1024; + +/// JSON2 physical layouts requested by a compaction read. +pub(crate) type Json2RewriteTargets = Arc>; + +/// Fixed JSON2 physical layout used while rewriting compaction input. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct Json2TargetLayout { + /// Logical JSON2 extension metadata attached to the rewritten field. + pub(crate) extension_metadata: String, + /// Settings used to build the fixed physical layout. + pub(crate) target_layout: JsonSettings, +} + /// Default row group size for parquet files. /// /// Keep the existing persisted/on-disk default stable. It intentionally stays diff --git a/src/mito2/src/sst/parquet/flat_format.rs b/src/mito2/src/sst/parquet/flat_format.rs index e381ea76f4..5432c112f4 100644 --- a/src/mito2/src/sst/parquet/flat_format.rs +++ b/src/mito2/src/sst/parquet/flat_format.rs @@ -241,9 +241,8 @@ impl FlatReadFormat { } /// Enables wrapping binary `__primary_key` batches back to a dictionary in [`Self::convert_batch`]. - pub(crate) fn set_pk_as_binary(&mut self) -> Result<()> { - self.pk_dict_wrap_schema = Some(self.output_arrow_schema()?); - Ok(()) + pub(crate) fn set_pk_as_binary(&mut self, output_schema: SchemaRef) { + self.pk_dict_wrap_schema = Some(output_schema); } /// Index of a column in the projected batch by its column id. @@ -306,19 +305,16 @@ impl FlatReadFormat { .project(projection) .context(ComputeArrowSnafu)?; let mut fields = schema.fields().iter().cloned().collect::>(); - for (column_id, target_type) in self.json_target_types().iter() { + for (column_id, target) in self.json_target_types().iter() { let Some(index) = self.parquet_projected_index_by_id(*column_id) else { continue; }; let Some(field) = schema.fields().get(index) else { continue; }; - fields[index] = Arc::new( - field - .as_ref() - .clone() - .with_data_type(ConcreteDataType::json2(target_type.clone()).as_arrow_type()), - ); + let mut field = field.as_ref().clone(); + field.set_data_type(ConcreteDataType::json2(target.clone()).as_arrow_type()); + fields[index] = Arc::new(field); } schema.fields = fields.into(); Ok(Arc::new(schema)) @@ -326,7 +322,7 @@ impl FlatReadFormat { /// Index of a column in the projected schema produced directly by parquet /// reading, before any primary-key-to-flat conversion. - fn parquet_projected_index_by_id(&self, column_id: ColumnId) -> Option { + pub(crate) fn parquet_projected_index_by_id(&self, column_id: ColumnId) -> Option { match &self.parquet_adapter { ParquetAdapter::Flat(p) => p .format_projection @@ -359,7 +355,7 @@ impl FlatReadFormat { } } - /// Gets JSON2 target types keyed by column id. + /// Gets JSON2 read targets. pub(crate) fn json_target_types(&self) -> &JsonTargetTypes { self.read_cols.json_target_types() } @@ -1012,9 +1008,8 @@ mod tests { false, ) .unwrap(); - read_format.set_pk_as_binary().unwrap(); - let output_schema = read_format.output_arrow_schema().unwrap(); + read_format.set_pk_as_binary(output_schema.clone()); let binary_schema = override_pk_field_to_binary(&output_schema); // The __primary_key field must preserve its field_id metadata after diff --git a/src/mito2/src/sst/parquet/format.rs b/src/mito2/src/sst/parquet/format.rs index 4ccdcb4108..1a4325b7a7 100644 --- a/src/mito2/src/sst/parquet/format.rs +++ b/src/mito2/src/sst/parquet/format.rs @@ -53,7 +53,7 @@ use store_api::storage::{ColumnId, NestedPath, SequenceNumber}; use crate::error::{ ConvertVectorSnafu, DecodeSnafu, InvalidRecordBatchSnafu, NewRecordBatchSnafu, Result, }; -use crate::read::read_columns::{JsonTargetTypes, ReadColumns}; +use crate::read::read_columns::ReadColumns; use crate::read::{Batch, BatchBuilder, BatchColumn}; use crate::sst::file::{FileMeta, FileTimeRange}; use crate::sst::parquet::read_columns::{ParquetReadColumn, ParquetReadColumns}; @@ -615,14 +615,13 @@ impl FormatProjection { sst_column_num: usize, cols: ReadColumns, ) -> Self { - let json_target_types = cols.json_target_types().clone(); let mut projected_columns: Vec<_> = cols .col_ids - .into_iter() + .iter() + .copied() .filter_map(|col_id| { id_to_index.get(&col_id).copied().map(|index_of_sst| { - let nested_paths = - json_target_nested_paths(metadata, &json_target_types, col_id); + let nested_paths = json_target_nested_paths(metadata, &cols, col_id); (col_id, index_of_sst, nested_paths) }) }) @@ -704,10 +703,10 @@ impl FormatProjection { fn json_target_nested_paths( metadata: &RegionMetadataRef, - json_target_types: &JsonTargetTypes, + read_columns: &ReadColumns, column_id: ColumnId, ) -> Vec { - let Some(target_type) = json_target_types.get(&column_id) else { + let Some(target_type) = read_columns.json_target_type(column_id) else { return Vec::new(); }; let Some(column) = metadata.column_by_id(column_id) else { diff --git a/src/mito2/src/sst/parquet/json_align/stream.rs b/src/mito2/src/sst/parquet/json_align/stream.rs index fabb3c807b..bfd0b82edc 100644 --- a/src/mito2/src/sst/parquet/json_align/stream.rs +++ b/src/mito2/src/sst/parquet/json_align/stream.rs @@ -12,15 +12,17 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; use std::pin::Pin; use std::task::{Context, Poll}; use datafusion_common::cast_column; 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::datatypes::{DataType, Field, FieldRef, SchemaRef}; use datatypes::arrow::record_batch::RecordBatch; -use datatypes::extension::json::is_json2_extension_type; +use datatypes::extension::json::{JsonMetadata, is_json2_extension_type}; +use datatypes::json::JsonSettings; use datatypes::vectors::json::array::JsonArray; use futures::Stream; use snafu::{ResultExt, ensure}; @@ -28,6 +30,13 @@ use snafu::{ResultExt, ensure}; use crate::error::{ CastColumnSnafu, DataTypeMismatchSnafu, NewRecordBatchSnafu, Result, UnexpectedSnafu, }; +use crate::sst::parquet::Json2TargetLayout; + +#[derive(Debug)] +struct Json2RewriteSettings { + logical_settings: JsonSettings, + target_layout: JsonSettings, +} /// Aligns projected batches to the expected output schema for nested projections. /// @@ -61,6 +70,8 @@ pub struct NestedSchemaAligner { /// Whether all projected roots are present and the stream can pass batches /// through. all_roots_present: bool, + /// JSON2 columns that require semantic source-to-target layout rewriting. + json2_rewrite_targets: HashMap, /// The cache for whether incoming batches already match output schema. is_schema_matched: Option, } @@ -96,9 +107,39 @@ where projected_root_presence, expected_input_col_num, all_roots_present, + json2_rewrite_targets: HashMap::new(), is_schema_matched: None, }) } + + /// Sets JSON2 columns that must be rewritten into the output field layout. + pub(crate) fn with_json2_rewrite_targets( + mut self, + targets: &HashMap, + ) -> Result { + self.json2_rewrite_targets = targets + .iter() + .map(|(name, layout)| { + let metadata = serde_json::from_str::(&layout.extension_metadata) + .map_err(|e| { + UnexpectedSnafu { + reason: format!( + "invalid JSON2 extension metadata for column '{name}': {e}" + ), + } + .build() + })?; + Ok(( + name.clone(), + Json2RewriteSettings { + logical_settings: metadata.into_json_settings(), + target_layout: layout.target_layout.clone(), + }, + )) + }) + .collect::>()?; + Ok(self) + } } impl Stream for NestedSchemaAligner @@ -125,6 +166,7 @@ where &this.output_schema, &this.projected_root_presence, this.expected_input_col_num, + &this.json2_rewrite_targets, ))) } } @@ -140,6 +182,7 @@ fn align_projected_batch( output_schema: &SchemaRef, projected_root_presence: &[bool], expected_input_col_num: usize, + json2_rewrite_targets: &HashMap, ) -> Result { ensure!( rb.columns().len() == expected_input_col_num, @@ -154,6 +197,7 @@ fn align_projected_batch( let mut cols = Vec::with_capacity(projected_root_presence.len()); let mut idx = 0; + let input_schema = rb.schema_ref(); for (field, present) in output_schema.fields().iter().zip(projected_root_presence) { if !present { @@ -161,19 +205,39 @@ fn align_projected_batch( continue; } - cols.push(align_array(rb.column(idx), field)?); + cols.push(align_array( + rb.column(idx), + input_schema.field(idx), + field, + json2_rewrite_targets.get(field.name()), + )?); idx += 1; } RecordBatch::try_new(output_schema.clone(), cols).context(NewRecordBatchSnafu) } -fn align_array(array: &ArrayRef, field: &FieldRef) -> Result { +fn align_array( + array: &ArrayRef, + source: &Field, + field: &FieldRef, + rewrite_settings: Option<&Json2RewriteSettings>, +) -> Result { + if let Some(settings) = rewrite_settings { + return JsonArray::from(array) + .rewrite_to_v2(source, &settings.logical_settings, &settings.target_layout) + .context(DataTypeMismatchSnafu); + } if array.data_type() == field.data_type() { return Ok(array.clone()); } if is_json2_extension_type(field) { + if is_json2_extension_type(source) { + return JsonArray::from(array) + .project_to_v2(source, field.data_type()) + .context(DataTypeMismatchSnafu); + } return JsonArray::from(array) .project_to(field.data_type()) .context(DataTypeMismatchSnafu); @@ -188,6 +252,7 @@ fn align_array(array: &ArrayRef, field: &FieldRef) -> Result { #[cfg(test)] mod tests { + use std::collections::HashMap; use std::sync::Arc; use datatypes::arrow::array::{ @@ -200,6 +265,33 @@ mod tests { use super::*; + #[test] + fn test_aligner_resolves_json2_rewrite_settings() + -> std::result::Result<(), Box> { + let logical_settings = JsonSettings::default(); + let target_layout = JsonSettings::try_new(vec![], Some(0))?; + let rewrite_targets = HashMap::from([( + "j".to_string(), + Json2TargetLayout { + extension_metadata: serde_json::to_string(&JsonMetadata::new( + logical_settings.clone(), + ))?, + target_layout: target_layout.clone(), + }, + )]); + let aligner = NestedSchemaAligner::new( + stream::empty::>(), + vec![], + schema(Vec::::new()), + )? + .with_json2_rewrite_targets(&rewrite_targets)?; + + let settings = &aligner.json2_rewrite_targets["j"]; + assert_eq!(logical_settings, settings.logical_settings); + assert_eq!(target_layout, settings.target_layout); + Ok(()) + } + #[tokio::test] async fn test_aligner_with_all_projected_roots_match() { let output_schema = schema([ diff --git a/src/mito2/src/sst/parquet/read_columns.rs b/src/mito2/src/sst/parquet/read_columns.rs index 68d7cae619..e6a5ab9fab 100644 --- a/src/mito2/src/sst/parquet/read_columns.rs +++ b/src/mito2/src/sst/parquet/read_columns.rs @@ -268,36 +268,33 @@ fn build_parquet_leaves_indices( // children should not enter this fallback path. for col in &projection.cols { let path_matches = &prefix_matched[&col.root_index]; - let needs_remainder = col - .nested_paths - .iter() - .zip(path_matches) - .any(|(path, matched)| { - !*matched || path_points_to_struct(parquet_schema_desc, col.root_index, path) - }); + let mut needs_remainder = false; + for (matched, nested_path) in path_matches.iter().zip(&col.nested_paths) { + if *matched { + if !needs_remainder { + needs_remainder = + path_points_to_struct(parquet_schema_desc, col.root_index, nested_path); + } + continue; + } + + if let Some(leaf_idx) = + find_nearest_variant_parent(parquet_schema_desc, col.root_index, nested_path) + { + matched_leaves.insert(leaf_idx); + matched_roots.insert(col.root_index); + } else { + needs_remainder = true; + } + } + if needs_remainder { let remainder_leaves = find_remainder_leaves(parquet_schema_desc, col.root_index); if !remainder_leaves.is_empty() { matched_leaves.extend(remainder_leaves); matched_roots.insert(col.root_index); - continue; } } - - for (matched, nested_path) in path_matches.iter().zip(&col.nested_paths) { - if *matched { - continue; - } - - let Some(leaf_idx) = - find_nearest_variant_parent(parquet_schema_desc, col.root_index, nested_path) - else { - continue; - }; - - matched_leaves.insert(leaf_idx); - matched_roots.insert(col.root_index); - } } let mut matched_leaves = matched_leaves.into_iter().collect::>(); @@ -551,6 +548,28 @@ mod tests { Ok(()) } + // A nested path under an explicit Variant is stored entirely in that Variant parent. The + // remainder may preserve `opaque: null`, but it cannot contain `opaque.leaf`, so reading the + // nearest Variant parent is sufficient. + #[test] + fn test_v2_variant_parent_path_reads_parent() -> Result<(), ParquetError> { + let parquet = build_test_v2_schema()?; + let projection = + ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths( + vec![vec![ + "j".to_string(), + "opaque".to_string(), + "leaf".to_string(), + ]], + )]); + + let plan = build_projection_plan(&projection, &parquet); + + assert_eq!(vec![true], plan.projected_root_presence); + assert_eq!(ProjectionMask::leaves(&parquet, [4]), plan.mask); + Ok(()) + } + #[test] fn test_merges_mixed_paths() { let parquet_schema_desc = build_test_nested_parquet_schema(); @@ -819,10 +838,18 @@ mod tests { .with_repetition(Repetition::OPTIONAL) .build()?, ); + // Normally there are no other explicit Variant fields exist if a remainder field is present. + // However, when structured values reach JSON2_MAX_STRUCTURED_DEPTH, there are. `opaque` + // models such a deep leaf without building a deeply nested test schema. + let opaque = Arc::new( + Type::primitive_type_builder("opaque", parquet::basic::Type::BYTE_ARRAY) + .with_repetition(Repetition::OPTIONAL) + .build()?, + ); let root = Arc::new( Type::group_type_builder("j") .with_repetition(Repetition::OPTIONAL) - .with_fields(vec![remainder, commit, hot]) + .with_fields(vec![remainder, commit, hot, opaque]) .build()?, ); Ok(SchemaDescriptor::new(Arc::new( diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 3adb335564..63232f6698 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -16,11 +16,14 @@ #[cfg(feature = "vector_index")] use std::collections::BTreeSet; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::time::{Duration, Instant}; use api::v1::SemanticType; +use arrow_schema::extension::{ + EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType, +}; use common_recordbatch::filter::SimpleFilterEvaluator; use common_telemetry::{debug, error, tracing, warn}; use datafusion::physical_plan::PhysicalExpr; @@ -31,8 +34,9 @@ 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_json2_extension_type; +use datatypes::extension::json::{Json2ExtensionType, is_json2_extension_type}; use datatypes::prelude::DataType; +use datatypes::vectors::json::json2_physical_data_type; use futures::StreamExt; use mito_codec::row_converter::build_primary_key_codec; use object_store::ObjectStore; @@ -76,7 +80,6 @@ use crate::sst::index::inverted_index::applier::{ }; #[cfg(feature = "vector_index")] use crate::sst::index::vector_index::applier::VectorIndexApplierRef; -use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; use crate::sst::parquet::file_range::{ FileRangeContext, FileRangeContextRef, PartitionFilterContext, PreFilterMode, RangeBase, }; @@ -94,6 +97,7 @@ use crate::sst::parquet::read_columns::{ProjectionMaskPlan, build_projection_pla use crate::sst::parquet::row_group::ParquetFetchMetrics; use crate::sst::parquet::row_selection::RowGroupSelection; use crate::sst::parquet::stats::RowGroupPruningStats; +use crate::sst::parquet::{DEFAULT_READ_BATCH_SIZE, Json2RewriteTargets, Json2TargetLayout}; use crate::sst::{override_pk_field_to_binary, tag_maybe_to_dictionary_field}; const INDEX_TYPE_FULLTEXT: &str = "fulltext"; @@ -108,6 +112,43 @@ fn should_read_pk_as_binary(parquet_meta: &ParquetMetaData) -> bool { should_read_pk_as_binary_with_limit(parquet_meta, DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT) } +fn apply_json2_rewrite_targets( + read_format: &FlatReadFormat, + targets: &Json2RewriteTargets, +) -> Result { + let schema = read_format.output_arrow_schema()?; + if targets.is_empty() { + return Ok(schema); + } + + let mut schema = schema.as_ref().clone(); + let mut fields = schema.fields().iter().cloned().collect::>(); + for (column_id, layout) in targets.iter() { + let Some(index) = read_format.parquet_projected_index_by_id(*column_id) else { + continue; + }; + let Some(field) = fields.get(index) else { + continue; + }; + let mut field = field.as_ref().clone(); + field.set_data_type(json2_physical_data_type(&layout.target_layout)); + + let mut metadata = field.metadata().clone(); + metadata.insert( + EXTENSION_TYPE_NAME_KEY.to_string(), + Json2ExtensionType::NAME.to_string(), + ); + metadata.insert( + EXTENSION_TYPE_METADATA_KEY.to_string(), + layout.extension_metadata.clone(), + ); + field.set_metadata(metadata); + fields[index] = Arc::new(field); + } + schema.fields = fields.into(); + Ok(Arc::new(schema)) +} + fn should_read_pk_as_binary_with_limit( parquet_meta: &ParquetMetaData, dict_page_size_limit: usize, @@ -163,6 +204,8 @@ pub struct ParquetReaderBuilder { /// `None` reads all columns. Due to schema change, the projection /// can contain columns not in the parquet file. read_cols: Option, + /// Compaction-only JSON2 physical rewrite targets. + json2_rewrite_targets: Json2RewriteTargets, /// Strategy to cache SST data. cache_strategy: CacheStrategy, /// Index appliers. @@ -208,6 +251,7 @@ impl ParquetReaderBuilder { object_store, predicate: None, read_cols: None, + json2_rewrite_targets: Arc::default(), cache_strategy: CacheStrategy::Disabled, inverted_index_appliers: [None, None], bloom_filter_index_appliers: [None, None], @@ -250,6 +294,13 @@ impl ParquetReaderBuilder { self } + /// Attaches fixed JSON2 physical layouts used by compaction readers. + #[must_use] + pub(crate) fn json2_rewrite_targets(mut self, targets: Json2RewriteTargets) -> Self { + self.json2_rewrite_targets = targets; + self + } + /// Attaches the cache to the builder. #[must_use] pub fn cache(mut self, cache: CacheStrategy) -> ParquetReaderBuilder { @@ -563,6 +614,8 @@ impl ParquetReaderBuilder { ); } + let output_schema = apply_json2_rewrite_targets(&read_format, &self.json2_rewrite_targets)?; + // Create ArrowReaderMetadata for async stream building. let mut arrow_reader_options = ArrowReaderOptions::new(); if !read_format @@ -574,7 +627,7 @@ impl ParquetReaderBuilder { // Read `__primary_key` as Binary when it's too large for dictionary // encoding; convert_batch wraps it back to a DictionaryArray. let schema_for_reader = if should_read_pk_as_binary(&parquet_meta) { - read_format.set_pk_as_binary()?; + read_format.set_pk_as_binary(output_schema.clone()); override_pk_field_to_binary(read_format.arrow_schema()) } else { read_format.arrow_schema().clone() @@ -585,7 +638,18 @@ impl ParquetReaderBuilder { ArrowReaderMetadata::try_new(parquet_meta.clone(), arrow_reader_options) .context(ReadDataPartSnafu)?; - let output_schema = read_format.output_arrow_schema()?; + let json2_rewrite_targets = self + .json2_rewrite_targets + .iter() + .map(|(column_id, layout)| { + let column = region_meta + .column_by_id(*column_id) + .context(UnexpectedSnafu { + reason: format!("JSON2 target column by id {column_id} does not exist"), + })?; + Ok((column.column_schema.name.clone(), layout.clone())) + }) + .collect::>>()?; let reader_builder = RowGroupReaderBuilder { file_handle: self.file_handle.clone(), @@ -594,6 +658,7 @@ impl ParquetReaderBuilder { parquet_metadata_size, arrow_metadata, output_schema, + json2_rewrite_targets, object_store: self.object_store.clone(), projection: projection_plan, has_nested_projection, @@ -1814,6 +1879,8 @@ pub(crate) struct RowGroupReaderBuilder { arrow_metadata: ArrowReaderMetadata, /// Projected output schema aligned with `projection.projected_root_presence`. output_schema: SchemaRef, + /// JSON2 columns that must be semantically rewritten into the projected target layout. + json2_rewrite_targets: HashMap, /// Object store as an Operator. object_store: ObjectStore, /// Projection mask. @@ -1975,7 +2042,7 @@ impl RowGroupReaderBuilder { &self, stream: ProjectedRecordBatchStream, ) -> Result { - if !self.has_nested_projection { + if !self.has_nested_projection && self.json2_rewrite_targets.is_empty() { return Ok(stream); } @@ -1984,6 +2051,7 @@ impl RowGroupReaderBuilder { self.projection.projected_root_presence.clone(), self.output_schema.clone(), )? + .with_json2_rewrite_targets(&self.json2_rewrite_targets)? .boxed()) } diff --git a/src/mito2/test-data/json2-v2.parquet b/src/mito2/test-data/json2-v2.parquet new file mode 100644 index 0000000000..a3cfe8f185 Binary files /dev/null and b/src/mito2/test-data/json2-v2.parquet differ diff --git a/src/pipeline/src/etl/transform/transformer/greptime.rs b/src/pipeline/src/etl/transform/transformer/greptime.rs index ca15808c00..c85510e3e4 100644 --- a/src/pipeline/src/etl/transform/transformer/greptime.rs +++ b/src/pipeline/src/etl/transform/transformer/greptime.rs @@ -312,6 +312,7 @@ impl GreptimeTransformer { pub struct ColumnMetadata { column_schema: datatypes::schema::ColumnSchema, semantic_type: SemanticType, + json_settings: OnceCell, } impl From for ColumnMetadata { @@ -338,10 +339,23 @@ impl From for ColumnMetadata { Self { column_schema, semantic_type, + json_settings: OnceCell::new(), } } } +impl ColumnMetadata { + fn json_settings(&self) -> Result<&JsonSettings> { + self.json_settings.get_or_try_init(|| { + if let Some(extension) = self.column_schema.extension_type::()? { + Ok(extension.metadata().json_settings().clone()) + } else { + Ok(parse_legacy_json2_settings(self.column_schema.metadata())?.unwrap_or_default()) + } + }) + } +} + impl TryFrom for ColumnSchema { type Error = api::error::Error; @@ -349,6 +363,7 @@ impl TryFrom for ColumnSchema { let ColumnMetadata { column_schema, semantic_type, + .. } = value; let options = options_from_column_schema(&column_schema); @@ -438,6 +453,7 @@ impl SchemaInfo { Some(ColumnMetadata { column_schema, semantic_type, + json_settings: OnceCell::new(), }) } else { None @@ -512,6 +528,7 @@ fn resolve_schema( ColumnMetadata { column_schema, semantic_type, + json_settings: OnceCell::new(), } }); let key = column.to_string(); @@ -759,16 +776,29 @@ fn resolve_value( } VrlValue::Array(_) | VrlValue::Object(_) => { - let json_settings = schema_info.json_settings_for_column(&column_name, None)?; + let index = index.or_else(|| { + let column = schema_info.find_column_schema_in_table(&column_name)?; + let index = schema_info.schema.len(); + schema_info.schema.push(column); + schema_info.index.insert(column_name.clone(), index); + Some(index) + }); + // TODO(LFC): Default to JSON2 for auto-created tables. + let json2_index = index.filter(|&index| { + matches!( + &schema_info.schema[index].column_schema.data_type, + ConcreteDataType::Json(column_type) if column_type.is_json2() + ) + }); - let value = if let Some(json_settings) = json_settings { + let value = if let Some(index) = json2_index { let value: serde_json::Value = value.try_into().map_err(|e: StdError| { CoerceIncompatibleTypesSnafu { msg: e.to_string() }.build() })?; - let value = json_settings.encode(value)?; + let value = schema_info.schema[index].json_settings()?.encode(value)?; resolve_schema( - index, + Some(index), p_ctx, &column_name, &ConcreteDataType::json2(Default::default()), @@ -856,6 +886,7 @@ fn identity_pipeline_inner( schema_info.schema.push(ColumnMetadata { column_schema, semantic_type: SemanticType::Timestamp, + json_settings: OnceCell::new(), }); let mut opt_map = HashMap::new(); @@ -1035,6 +1066,24 @@ mod tests { use super::*; use crate::{PipelineDefinition, identity_pipeline}; + #[test] + fn test_column_metadata_caches_json_settings() -> Result<()> { + let column = ColumnMetadata { + column_schema: datatypes::schema::ColumnSchema::new( + "data", + ConcreteDataType::json2(Default::default()), + true, + ), + semantic_type: SemanticType::Field, + json_settings: OnceCell::new(), + }; + + let first = column.json_settings()?; + let second = column.json_settings()?; + assert!(std::ptr::eq(first, second)); + Ok(()) + } + #[test] fn test_transform_json2_uses_destination_table_settings() { let table = |name: &str, settings: JsonSettings, sample: serde_json::Value| { diff --git a/src/query/src/optimizer/json_type_concretize.rs b/src/query/src/optimizer/json_type_concretize.rs index c9964e119d..e7f1392dc5 100644 --- a/src/query/src/optimizer/json_type_concretize.rs +++ b/src/query/src/optimizer/json_type_concretize.rs @@ -124,6 +124,13 @@ fn deduce_json_types(plan: &LogicalPlan) -> Result Result bool { + match expr { + Expr::Column(_) => true, + Expr::Alias(alias) => { + matches!(alias.expr.as_ref(), Expr::Column(column) if column.name == alias.name) + } + _ => false, + } +} + fn deduce_json_type(expr: &Expr) -> Result> { let f = match expr { Expr::ScalarFunction(f) if f.name().eq_ignore_ascii_case(JsonGetWithType::NAME) => f, diff --git a/src/sql/src/statements.rs b/src/sql/src/statements.rs index 4f96b66f87..cf0b248675 100644 --- a/src/sql/src/statements.rs +++ b/src/sql/src/statements.rs @@ -40,6 +40,7 @@ use api::v1::SemanticType; use common_sql::default_constraint::parse_column_default_constraint; use common_time::timezone::Timezone; use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; +use datatypes::json::JsonSettings; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{COMMENT_KEY, ColumnDefaultConstraint, ColumnSchema}; use datatypes::types::json_type::JsonNativeType; @@ -163,7 +164,10 @@ pub fn column_to_schema( false }; if is_json2_column { - let settings = column.extensions.build_json_settings()?.unwrap_or_default(); + let settings = column + .extensions + .build_json_settings()? + .unwrap_or_else(JsonSettings::new_v2); let extension = Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings))); column_schema.with_extension_type(&extension); } @@ -643,6 +647,53 @@ mod tests { ); } + #[test] + fn test_new_json2_column_uses_v2_layout() -> std::result::Result<(), Box> + { + let column = Column { + column_def: ColumnDef { + name: "data".into(), + data_type: SqlDataType::Custom( + sqlparser::ast::ObjectName::from(vec!["JSON2".into()]), + vec![], + ), + options: vec![], + }, + extensions: ColumnExtensions::default(), + }; + + let schema = column_to_schema(&column, "ts", None)?; + let metadata: serde_json::Value = + serde_json::from_str(schema.metadata().get("ARROW:extension:metadata").unwrap())?; + assert_eq!(Some(2), metadata["layout_version"].as_u64()); + assert_eq!( + Some(100), + metadata["json_settings"]["max_auto_expanded_paths"].as_u64() + ); + + let mut hinted = column; + hinted + .extensions + .set_json_settings(datatypes::json::JsonSettings::try_new( + vec![datatypes::json::JsonTypeHint { + path: vec!["kind".to_string()], + data_type: ConcreteDataType::string_datatype(), + nullable: true, + default_constraint: None, + inverted_index: false, + }], + None, + )?)?; + let schema = column_to_schema(&hinted, "ts", None)?; + let metadata: serde_json::Value = + serde_json::from_str(schema.metadata().get("ARROW:extension:metadata").unwrap())?; + assert_eq!( + Some(100), + metadata["json_settings"]["max_auto_expanded_paths"].as_u64() + ); + Ok(()) + } + #[test] pub fn test_column_to_schema_timestamp_with_timezone() { let column = Column { diff --git a/src/sql/src/statements/create.rs b/src/sql/src/statements/create.rs index f3fbc90836..0db002c672 100644 --- a/src/sql/src/statements/create.rs +++ b/src/sql/src/statements/create.rs @@ -17,7 +17,7 @@ use std::fmt::{Display, Formatter}; use common_catalog::consts::FILE_ENGINE; use common_sql::default_constraint::parse_column_default_constraint; -use datatypes::json::JsonSettings; +use datatypes::json::{JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS, JsonSettings}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ ColumnDefaultConstraint, FulltextOptions, SkippingIndexOptions, VectorDistanceMetric, @@ -369,7 +369,12 @@ impl ColumnExtensions { }) }) .collect::>>()?; - let settings = JsonSettings::try_new(type_hints, options.max_auto_expanded_paths)?; + let settings = JsonSettings::try_new( + type_hints, + options + .max_auto_expanded_paths + .or(Some(JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS)), + )?; Ok(Some(settings)) } @@ -1017,6 +1022,31 @@ ENGINE=mito } } + #[test] + fn test_parse_json2_max_auto_expanded_paths_option() -> Result<()> { + let sql = r#"CREATE TABLE traces ( + log_json_data JSON2 ( + status_code INT64 NOT NULL, + max_auto_expanded_paths = 1 + ), + ts TIMESTAMP TIME INDEX + )"#; + let result = ParserContext::create_with_dialect( + sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + )?; + let Statement::CreateTable(create_table) = &result[0] else { + unreachable!() + }; + let settings = create_table.columns[0] + .extensions + .build_json_settings()? + .unwrap(); + assert_eq!(settings.max_auto_expanded_paths(), Some(1)); + Ok(()) + } + #[test] fn test_display_json2_type_hints_quotes_numeric_segments() { let sql = r#"CREATE TABLE traces ( diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index f21b459c2a..a716284c75 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -3360,19 +3360,21 @@ CREATE TABLE b ( let output = execute_sql(&instance, "SHOW CREATE TABLE b").await.data; let expected = r#" -+-------+----------------------------------+ -| Table | Create Table | -+-------+----------------------------------+ -| b | CREATE TABLE IF NOT EXISTS "b" ( | -| | "j" JSON2 NULL, | -| | "ts" TIMESTAMP(3) NOT NULL, | -| | TIME INDEX ("ts") | -| | ) | -| | | -| | ENGINE=mito | -| | WITH( | -| | append_mode = 'true' | -| | ) | -+-------+----------------------------------+"#; ++-------+-----------------------------------+ +| Table | Create Table | ++-------+-----------------------------------+ +| b | CREATE TABLE IF NOT EXISTS "b" ( | +| | "j" JSON2( | +| | max_auto_expanded_paths = 100 | +| | ) NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true' | +| | ) | ++-------+-----------------------------------+"#; check_output_stream(output, expected).await; } diff --git a/tests/cases/standalone/common/types/json/json2.result b/tests/cases/standalone/common/types/json/json2.result index 6afa0a759c..32fa31a76c 100644 --- a/tests/cases/standalone/common/types/json/json2.result +++ b/tests/cases/standalone/common/types/json/json2.result @@ -126,38 +126,38 @@ select j.a, j.a.x from json2_table order by ts; | {"b":-2} | | | {"b":3} | | | {"b":-4} | | -| | | +| {} | | | | | | {"b":"s7"} | | | {"b":8} | | -| {"b":null,"x":true} | true | -| {"b":10,"x":null} | | +| {"x":true} | true | +| {"b":10} | | +-----------------------------------+-------------------------------------+ select j, j.a from json2_table order by ts; -+--------------------------------------------------------------------+-----------------------------------+ -| j | json_get(json2_table.j,Utf8("a")) | -+--------------------------------------------------------------------+-----------------------------------+ -| {"a":{"b":1},"c":"s1","d":[{"e":{"f":0.1,"g":null}}]} | {"b":1} | -| {"a":{"b":-2},"c":"s2","d":[{"e":{"f":0.2,"g":null}}]} | {"b":-2} | -| {"a":{"b":3},"c":"s3","d":null} | {"b":3} | -| {"a":{"b":-4},"c":null,"d":[{"e":{"f":null,"g":-0.4}}]} | {"b":-4} | -| {"a":null,"c":"s5","d":null} | | -| {"a":null,"c":"s6","d":null} | | -| {"a":{"b":"s7"},"c":[1],"d":[{"e":{"g":-0.7}}]} | {"b":"s7"} | -| {"a":{"b":8},"c":"s8","d":null} | {"b":8} | -| {"a":{"b":null,"x":true},"c":"s9","d":[{"e":{"g":-0.9}}],"y":null} | {"b":null,"x":true} | -| {"a":{"b":10,"x":null},"c":null,"d":null,"y":false} | {"b":10,"x":null} | -+--------------------------------------------------------------------+-----------------------------------+ ++--------------------------------------------------+-----------------------------------+ +| j | json_get(json2_table.j,Utf8("a")) | ++--------------------------------------------------+-----------------------------------+ +| {"a":{"b":1},"c":"s1","d":[{"e":{"f":0.1}}]} | {"b":1} | +| {"a":{"b":-2},"c":"s2","d":[{"e":{"f":0.2}}]} | {"b":-2} | +| {"a":{"b":3},"c":"s3"} | {"b":3} | +| {"a":{"b":-4},"d":[{"e":{"g":-0.4}}]} | {"b":-4} | +| {"a":{},"c":"s5"} | {} | +| {"c":"s6"} | | +| {"a":{"b":"s7"},"c":[1],"d":[{"e":{"g":-0.7}}]} | {"b":"s7"} | +| {"a":{"b":8},"c":"s8"} | {"b":8} | +| {"a":{"x":true},"c":"s9","d":[{"e":{"g":-0.9}}]} | {"x":true} | +| {"a":{"b":10},"y":false} | {"b":10} | ++--------------------------------------------------+-----------------------------------+ select j from json2_table where j.a.b = 1; -+-------------------------------------------------------+ -| j | -+-------------------------------------------------------+ -| {"a":{"b":1},"c":"s1","d":[{"e":{"f":0.1,"g":null}}]} | -+-------------------------------------------------------+ ++----------------------------------------------+ +| j | ++----------------------------------------------+ +| {"a":{"b":1},"c":"s1","d":[{"e":{"f":0.1}}]} | ++----------------------------------------------+ select j.c, j.y from json2_table order by ts; @@ -178,37 +178,37 @@ select j.c, j.y from json2_table order by ts; select j from json2_table order by ts; -+--------------------------------------------------------------------+ -| j | -+--------------------------------------------------------------------+ -| {"a":{"b":1},"c":"s1","d":[{"e":{"f":0.1,"g":null}}]} | -| {"a":{"b":-2},"c":"s2","d":[{"e":{"f":0.2,"g":null}}]} | -| {"a":{"b":3},"c":"s3","d":null} | -| {"a":{"b":-4},"c":null,"d":[{"e":{"f":null,"g":-0.4}}]} | -| {"a":null,"c":"s5","d":null} | -| {"a":null,"c":"s6","d":null} | -| {"a":{"b":"s7"},"c":[1],"d":[{"e":{"g":-0.7}}]} | -| {"a":{"b":8},"c":"s8","d":null} | -| {"a":{"b":null,"x":true},"c":"s9","d":[{"e":{"g":-0.9}}],"y":null} | -| {"a":{"b":10,"x":null},"c":null,"d":null,"y":false} | -+--------------------------------------------------------------------+ ++--------------------------------------------------+ +| j | ++--------------------------------------------------+ +| {"a":{"b":1},"c":"s1","d":[{"e":{"f":0.1}}]} | +| {"a":{"b":-2},"c":"s2","d":[{"e":{"f":0.2}}]} | +| {"a":{"b":3},"c":"s3"} | +| {"a":{"b":-4},"d":[{"e":{"g":-0.4}}]} | +| {"a":{},"c":"s5"} | +| {"c":"s6"} | +| {"a":{"b":"s7"},"c":[1],"d":[{"e":{"g":-0.7}}]} | +| {"a":{"b":8},"c":"s8"} | +| {"a":{"x":true},"c":"s9","d":[{"e":{"g":-0.9}}]} | +| {"a":{"b":10},"y":false} | ++--------------------------------------------------+ select * from json2_table order by ts; -+-------------------------+--------------------------------------------------------------------+ -| ts | j | -+-------------------------+--------------------------------------------------------------------+ -| 1970-01-01T00:00:00.001 | {"a":{"b":1},"c":"s1","d":[{"e":{"f":0.1,"g":null}}]} | -| 1970-01-01T00:00:00.002 | {"a":{"b":-2},"c":"s2","d":[{"e":{"f":0.2,"g":null}}]} | -| 1970-01-01T00:00:00.003 | {"a":{"b":3},"c":"s3","d":null} | -| 1970-01-01T00:00:00.004 | {"a":{"b":-4},"c":null,"d":[{"e":{"f":null,"g":-0.4}}]} | -| 1970-01-01T00:00:00.005 | {"a":null,"c":"s5","d":null} | -| 1970-01-01T00:00:00.006 | {"a":null,"c":"s6","d":null} | -| 1970-01-01T00:00:00.007 | {"a":{"b":"s7"},"c":[1],"d":[{"e":{"g":-0.7}}]} | -| 1970-01-01T00:00:00.008 | {"a":{"b":8},"c":"s8","d":null} | -| 1970-01-01T00:00:00.009 | {"a":{"b":null,"x":true},"c":"s9","d":[{"e":{"g":-0.9}}],"y":null} | -| 1970-01-01T00:00:00.010 | {"a":{"b":10,"x":null},"c":null,"d":null,"y":false} | -+-------------------------+--------------------------------------------------------------------+ ++-------------------------+--------------------------------------------------+ +| ts | j | ++-------------------------+--------------------------------------------------+ +| 1970-01-01T00:00:00.001 | {"a":{"b":1},"c":"s1","d":[{"e":{"f":0.1}}]} | +| 1970-01-01T00:00:00.002 | {"a":{"b":-2},"c":"s2","d":[{"e":{"f":0.2}}]} | +| 1970-01-01T00:00:00.003 | {"a":{"b":3},"c":"s3"} | +| 1970-01-01T00:00:00.004 | {"a":{"b":-4},"d":[{"e":{"g":-0.4}}]} | +| 1970-01-01T00:00:00.005 | {"a":{},"c":"s5"} | +| 1970-01-01T00:00:00.006 | {"c":"s6"} | +| 1970-01-01T00:00:00.007 | {"a":{"b":"s7"},"c":[1],"d":[{"e":{"g":-0.7}}]} | +| 1970-01-01T00:00:00.008 | {"a":{"b":8},"c":"s8"} | +| 1970-01-01T00:00:00.009 | {"a":{"x":true},"c":"s9","d":[{"e":{"g":-0.9}}]} | +| 1970-01-01T00:00:00.010 | {"a":{"b":10},"y":false} | ++-------------------------+--------------------------------------------------+ select count(*) from (select j from json2_table group by j); @@ -228,88 +228,88 @@ select count(*) from (select distinct j from json2_table); select ts, j from (select ts, j from json2_table) order by ts; -+-------------------------+--------------------------------------------------------------------+ -| ts | j | -+-------------------------+--------------------------------------------------------------------+ -| 1970-01-01T00:00:00.001 | {"a":{"b":1},"c":"s1","d":[{"e":{"f":0.1,"g":null}}]} | -| 1970-01-01T00:00:00.002 | {"a":{"b":-2},"c":"s2","d":[{"e":{"f":0.2,"g":null}}]} | -| 1970-01-01T00:00:00.003 | {"a":{"b":3},"c":"s3","d":null} | -| 1970-01-01T00:00:00.004 | {"a":{"b":-4},"c":null,"d":[{"e":{"f":null,"g":-0.4}}]} | -| 1970-01-01T00:00:00.005 | {"a":null,"c":"s5","d":null} | -| 1970-01-01T00:00:00.006 | {"a":null,"c":"s6","d":null} | -| 1970-01-01T00:00:00.007 | {"a":{"b":"s7"},"c":[1],"d":[{"e":{"g":-0.7}}]} | -| 1970-01-01T00:00:00.008 | {"a":{"b":8},"c":"s8","d":null} | -| 1970-01-01T00:00:00.009 | {"a":{"b":null,"x":true},"c":"s9","d":[{"e":{"g":-0.9}}],"y":null} | -| 1970-01-01T00:00:00.010 | {"a":{"b":10,"x":null},"c":null,"d":null,"y":false} | -+-------------------------+--------------------------------------------------------------------+ ++-------------------------+--------------------------------------------------+ +| ts | j | ++-------------------------+--------------------------------------------------+ +| 1970-01-01T00:00:00.001 | {"a":{"b":1},"c":"s1","d":[{"e":{"f":0.1}}]} | +| 1970-01-01T00:00:00.002 | {"a":{"b":-2},"c":"s2","d":[{"e":{"f":0.2}}]} | +| 1970-01-01T00:00:00.003 | {"a":{"b":3},"c":"s3"} | +| 1970-01-01T00:00:00.004 | {"a":{"b":-4},"d":[{"e":{"g":-0.4}}]} | +| 1970-01-01T00:00:00.005 | {"a":{},"c":"s5"} | +| 1970-01-01T00:00:00.006 | {"c":"s6"} | +| 1970-01-01T00:00:00.007 | {"a":{"b":"s7"},"c":[1],"d":[{"e":{"g":-0.7}}]} | +| 1970-01-01T00:00:00.008 | {"a":{"b":8},"c":"s8"} | +| 1970-01-01T00:00:00.009 | {"a":{"x":true},"c":"s9","d":[{"e":{"g":-0.9}}]} | +| 1970-01-01T00:00:00.010 | {"a":{"b":10},"y":false} | ++-------------------------+--------------------------------------------------+ select json_get(j, '') from json2_table order by ts; -+--------------------------------------------------------------------+ -| json_get(json2_table.j,Utf8("")) | -+--------------------------------------------------------------------+ -| {"a":{"b":1},"c":"s1","d":[{"e":{"f":0.1,"g":null}}]} | -| {"a":{"b":-2},"c":"s2","d":[{"e":{"f":0.2,"g":null}}]} | -| {"a":{"b":3},"c":"s3","d":null} | -| {"a":{"b":-4},"c":null,"d":[{"e":{"f":null,"g":-0.4}}]} | -| {"a":null,"c":"s5","d":null} | -| {"a":null,"c":"s6","d":null} | -| {"a":{"b":"s7"},"c":[1],"d":[{"e":{"g":-0.7}}]} | -| {"a":{"b":8},"c":"s8","d":null} | -| {"a":{"b":null,"x":true},"c":"s9","d":[{"e":{"g":-0.9}}],"y":null} | -| {"a":{"b":10,"x":null},"c":null,"d":null,"y":false} | -+--------------------------------------------------------------------+ ++--------------------------------------------------+ +| json_get(json2_table.j,Utf8("")) | ++--------------------------------------------------+ +| {"a":{"b":1},"c":"s1","d":[{"e":{"f":0.1}}]} | +| {"a":{"b":-2},"c":"s2","d":[{"e":{"f":0.2}}]} | +| {"a":{"b":3},"c":"s3"} | +| {"a":{"b":-4},"d":[{"e":{"g":-0.4}}]} | +| {"a":{},"c":"s5"} | +| {"c":"s6"} | +| {"a":{"b":"s7"},"c":[1],"d":[{"e":{"g":-0.7}}]} | +| {"a":{"b":8},"c":"s8"} | +| {"a":{"x":true},"c":"s9","d":[{"e":{"g":-0.9}}]} | +| {"a":{"b":10},"y":false} | ++--------------------------------------------------+ select json_get(j, '$') from json2_table order by ts; -+--------------------------------------------------------------------+ -| json_get(json2_table.j,Utf8("$")) | -+--------------------------------------------------------------------+ -| {"a":{"b":1},"c":"s1","d":[{"e":{"f":0.1,"g":null}}]} | -| {"a":{"b":-2},"c":"s2","d":[{"e":{"f":0.2,"g":null}}]} | -| {"a":{"b":3},"c":"s3","d":null} | -| {"a":{"b":-4},"c":null,"d":[{"e":{"f":null,"g":-0.4}}]} | -| {"a":null,"c":"s5","d":null} | -| {"a":null,"c":"s6","d":null} | -| {"a":{"b":"s7"},"c":[1],"d":[{"e":{"g":-0.7}}]} | -| {"a":{"b":8},"c":"s8","d":null} | -| {"a":{"b":null,"x":true},"c":"s9","d":[{"e":{"g":-0.9}}],"y":null} | -| {"a":{"b":10,"x":null},"c":null,"d":null,"y":false} | -+--------------------------------------------------------------------+ ++--------------------------------------------------+ +| json_get(json2_table.j,Utf8("$")) | ++--------------------------------------------------+ +| {"a":{"b":1},"c":"s1","d":[{"e":{"f":0.1}}]} | +| {"a":{"b":-2},"c":"s2","d":[{"e":{"f":0.2}}]} | +| {"a":{"b":3},"c":"s3"} | +| {"a":{"b":-4},"d":[{"e":{"g":-0.4}}]} | +| {"a":{},"c":"s5"} | +| {"c":"s6"} | +| {"a":{"b":"s7"},"c":[1],"d":[{"e":{"g":-0.7}}]} | +| {"a":{"b":8},"c":"s8"} | +| {"a":{"x":true},"c":"s9","d":[{"e":{"g":-0.9}}]} | +| {"a":{"b":10},"y":false} | ++--------------------------------------------------+ select json_get(j, '.') from json2_table order by ts; -+--------------------------------------------------------------------+ -| json_get(json2_table.j,Utf8(".")) | -+--------------------------------------------------------------------+ -| {"a":{"b":1},"c":"s1","d":[{"e":{"f":0.1,"g":null}}]} | -| {"a":{"b":-2},"c":"s2","d":[{"e":{"f":0.2,"g":null}}]} | -| {"a":{"b":3},"c":"s3","d":null} | -| {"a":{"b":-4},"c":null,"d":[{"e":{"f":null,"g":-0.4}}]} | -| {"a":null,"c":"s5","d":null} | -| {"a":null,"c":"s6","d":null} | -| {"a":{"b":"s7"},"c":[1],"d":[{"e":{"g":-0.7}}]} | -| {"a":{"b":8},"c":"s8","d":null} | -| {"a":{"b":null,"x":true},"c":"s9","d":[{"e":{"g":-0.9}}],"y":null} | -| {"a":{"b":10,"x":null},"c":null,"d":null,"y":false} | -+--------------------------------------------------------------------+ ++--------------------------------------------------+ +| json_get(json2_table.j,Utf8(".")) | ++--------------------------------------------------+ +| {"a":{"b":1},"c":"s1","d":[{"e":{"f":0.1}}]} | +| {"a":{"b":-2},"c":"s2","d":[{"e":{"f":0.2}}]} | +| {"a":{"b":3},"c":"s3"} | +| {"a":{"b":-4},"d":[{"e":{"g":-0.4}}]} | +| {"a":{},"c":"s5"} | +| {"c":"s6"} | +| {"a":{"b":"s7"},"c":[1],"d":[{"e":{"g":-0.7}}]} | +| {"a":{"b":8},"c":"s8"} | +| {"a":{"x":true},"c":"s9","d":[{"e":{"g":-0.9}}]} | +| {"a":{"b":10},"y":false} | ++--------------------------------------------------+ select json_get(j, '$.') from json2_table order by ts; -+--------------------------------------------------------------------+ -| json_get(json2_table.j,Utf8("$.")) | -+--------------------------------------------------------------------+ -| {"a":{"b":1},"c":"s1","d":[{"e":{"f":0.1,"g":null}}]} | -| {"a":{"b":-2},"c":"s2","d":[{"e":{"f":0.2,"g":null}}]} | -| {"a":{"b":3},"c":"s3","d":null} | -| {"a":{"b":-4},"c":null,"d":[{"e":{"f":null,"g":-0.4}}]} | -| {"a":null,"c":"s5","d":null} | -| {"a":null,"c":"s6","d":null} | -| {"a":{"b":"s7"},"c":[1],"d":[{"e":{"g":-0.7}}]} | -| {"a":{"b":8},"c":"s8","d":null} | -| {"a":{"b":null,"x":true},"c":"s9","d":[{"e":{"g":-0.9}}],"y":null} | -| {"a":{"b":10,"x":null},"c":null,"d":null,"y":false} | -+--------------------------------------------------------------------+ ++--------------------------------------------------+ +| json_get(json2_table.j,Utf8("$.")) | ++--------------------------------------------------+ +| {"a":{"b":1},"c":"s1","d":[{"e":{"f":0.1}}]} | +| {"a":{"b":-2},"c":"s2","d":[{"e":{"f":0.2}}]} | +| {"a":{"b":3},"c":"s3"} | +| {"a":{"b":-4},"d":[{"e":{"g":-0.4}}]} | +| {"a":{},"c":"s5"} | +| {"c":"s6"} | +| {"a":{"b":"s7"},"c":[1],"d":[{"e":{"g":-0.7}}]} | +| {"a":{"b":8},"c":"s8"} | +| {"a":{"x":true},"c":"s9","d":[{"e":{"g":-0.9}}]} | +| {"a":{"b":10},"y":false} | ++--------------------------------------------------+ select j.a.b + 1 from json2_table order by ts; @@ -346,10 +346,10 @@ select j.d from json2_table order by ts; +-----------------------------------+ | json_get(json2_table.j,Utf8("d")) | +-----------------------------------+ -| [{"e":{"f":0.1,"g":null}}] | -| [{"e":{"f":0.2,"g":null}}] | +| [{"e":{"f":0.1}}] | +| [{"e":{"f":0.2}}] | | | -| [{"e":{"f":null,"g":-0.4}}] | +| [{"e":{"g":-0.4}}] | | | | | | [{"e":{"g":-0.7}}] | @@ -417,3 +417,116 @@ drop table json2_variant_null; Affected Rows: 0 +create table json2_finite_paths ( + ts timestamp time index, + j json2( + max_auto_expanded_paths = 1, + hint string + ) +) +with ( + 'append_mode' = 'true', + 'sst_format' = 'flat' +); + +Affected Rows: 0 + +show create table json2_finite_paths; + ++--------------------+---------------------------------------------------+ +| Table | Create Table | ++--------------------+---------------------------------------------------+ +| json2_finite_paths | CREATE TABLE IF NOT EXISTS "json2_finite_paths" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "j" JSON2( | +| | max_auto_expanded_paths = 1, | +| | "hint" STRING NULL | +| | ) NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true', | +| | sst_format = 'flat' | +| | ) | ++--------------------+---------------------------------------------------+ + +insert into json2_finite_paths values + (1, '{"hint":"h1","alpha":1,"conflict":1}'), + (2, '{"hint":"h2","alpha":2,"conflict":"text"}'); + +Affected Rows: 2 + +admin flush_table('json2_finite_paths'); + ++-----------------------------------------+ +| ADMIN flush_table('json2_finite_paths') | ++-----------------------------------------+ +| 0 | ++-----------------------------------------+ + +insert into json2_finite_paths values + (3, '{"hint":"h3","beta":3,"conflict":true}'), + (4, '{"hint":"h4","beta":4,"conflict":"other"}'); + +Affected Rows: 2 + +admin flush_table('json2_finite_paths'); + ++-----------------------------------------+ +| ADMIN flush_table('json2_finite_paths') | ++-----------------------------------------+ +| 0 | ++-----------------------------------------+ + +select + ts, + j, + j.hint, + j.alpha::bigint as alpha, + j.beta::bigint as beta, + j.conflict +from json2_finite_paths +order by ts; + ++-------------------------+-------------------------------------------+---------------------------------------------+-------+------+-------------------------------------------------+ +| ts | j | json_get(json2_finite_paths.j,Utf8("hint")) | alpha | beta | json_get(json2_finite_paths.j,Utf8("conflict")) | ++-------------------------+-------------------------------------------+---------------------------------------------+-------+------+-------------------------------------------------+ +| 1970-01-01T00:00:00.001 | {"alpha":1,"conflict":1,"hint":"h1"} | h1 | 1 | | 1 | +| 1970-01-01T00:00:00.002 | {"alpha":2,"conflict":"text","hint":"h2"} | h2 | 2 | | text | +| 1970-01-01T00:00:00.003 | {"beta":3,"conflict":true,"hint":"h3"} | h3 | | 3 | true | +| 1970-01-01T00:00:00.004 | {"beta":4,"conflict":"other","hint":"h4"} | h4 | | 4 | other | ++-------------------------+-------------------------------------------+---------------------------------------------+-------+------+-------------------------------------------------+ + +admin compact_table('json2_finite_paths'); + ++-------------------------------------------+ +| ADMIN compact_table('json2_finite_paths') | ++-------------------------------------------+ +| 0 | ++-------------------------------------------+ + +select + ts, + j, + j.hint, + j.alpha::bigint as alpha, + j.beta::bigint as beta, + j.conflict +from json2_finite_paths +order by ts; + ++-------------------------+-------------------------------------------+---------------------------------------------+-------+------+-------------------------------------------------+ +| ts | j | json_get(json2_finite_paths.j,Utf8("hint")) | alpha | beta | json_get(json2_finite_paths.j,Utf8("conflict")) | ++-------------------------+-------------------------------------------+---------------------------------------------+-------+------+-------------------------------------------------+ +| 1970-01-01T00:00:00.001 | {"alpha":1,"conflict":1,"hint":"h1"} | h1 | 1 | | 1 | +| 1970-01-01T00:00:00.002 | {"alpha":2,"conflict":"text","hint":"h2"} | h2 | 2 | | text | +| 1970-01-01T00:00:00.003 | {"beta":3,"conflict":true,"hint":"h3"} | h3 | | 3 | true | +| 1970-01-01T00:00:00.004 | {"beta":4,"conflict":"other","hint":"h4"} | h4 | | 4 | other | ++-------------------------+-------------------------------------------+---------------------------------------------+-------+------+-------------------------------------------------+ + +drop table json2_finite_paths; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2.sql b/tests/cases/standalone/common/types/json/json2.sql index 406cc4af54..090e1f4fd2 100644 --- a/tests/cases/standalone/common/types/json/json2.sql +++ b/tests/cases/standalone/common/types/json/json2.sql @@ -105,3 +105,53 @@ from json2_variant_null order by ts; drop table json2_variant_null; + +create table json2_finite_paths ( + ts timestamp time index, + j json2( + max_auto_expanded_paths = 1, + hint string + ) +) +with ( + 'append_mode' = 'true', + 'sst_format' = 'flat' +); + +show create table json2_finite_paths; + +insert into json2_finite_paths values + (1, '{"hint":"h1","alpha":1,"conflict":1}'), + (2, '{"hint":"h2","alpha":2,"conflict":"text"}'); + +admin flush_table('json2_finite_paths'); + +insert into json2_finite_paths values + (3, '{"hint":"h3","beta":3,"conflict":true}'), + (4, '{"hint":"h4","beta":4,"conflict":"other"}'); + +admin flush_table('json2_finite_paths'); + +select + ts, + j, + j.hint, + j.alpha::bigint as alpha, + j.beta::bigint as beta, + j.conflict +from json2_finite_paths +order by ts; + +admin compact_table('json2_finite_paths'); + +select + ts, + j, + j.hint, + j.alpha::bigint as alpha, + j.beta::bigint as beta, + j.conflict +from json2_finite_paths +order by ts; + +drop table json2_finite_paths; diff --git a/tests/cases/standalone/common/types/json/json2_type_hints.result b/tests/cases/standalone/common/types/json/json2_type_hints.result index 785fa76942..d9a384e084 100644 --- a/tests/cases/standalone/common/types/json/json2_type_hints.result +++ b/tests/cases/standalone/common/types/json/json2_type_hints.result @@ -20,6 +20,7 @@ SHOW CREATE TABLE json2_type_hints; | json2_type_hints | CREATE TABLE IF NOT EXISTS "json2_type_hints" ( | | | "ts" TIMESTAMP(3) NOT NULL, | | | "j" JSON2( | +| | max_auto_expanded_paths = 100, | | | "user"."age" BIGINT NOT NULL DEFAULT 18, | | | "user"."name" STRING NULL DEFAULT 'unknown', | | | "user"."active" BOOLEAN NULL, | diff --git a/tests/compatibility/cases/legacy_json2_non_append_table/verify.result b/tests/compatibility/cases/legacy_json2_non_append_table/verify.result index 2cc3c230e2..078eefb94f 100644 --- a/tests/compatibility/cases/legacy_json2_non_append_table/verify.result +++ b/tests/compatibility/cases/legacy_json2_non_append_table/verify.result @@ -50,6 +50,22 @@ INSERT INTO t_legacy_json2_non_append_table (ts, j) VALUES Affected Rows: 1 +ADMIN FLUSH_TABLE('t_legacy_json2_non_append_table'); + ++-------------------------------------------------------+ +| ADMIN flush_table('t_legacy_json2_non_append_table') | ++-------------------------------------------------------+ +| 0 | ++-------------------------------------------------------+ + +ADMIN compact_table('t_legacy_json2_non_append_table'); + ++---------------------------------------------------------+ +| ADMIN compact_table('t_legacy_json2_non_append_table') | ++---------------------------------------------------------+ +| 0 | ++---------------------------------------------------------+ + SELECT ts, j.a AS a, j.nested.s AS nested_s FROM t_legacy_json2_non_append_table ORDER BY ts; diff --git a/tests/compatibility/cases/legacy_json2_non_append_table/verify.sql b/tests/compatibility/cases/legacy_json2_non_append_table/verify.sql index 2dbc7fc729..1cdc0015da 100644 --- a/tests/compatibility/cases/legacy_json2_non_append_table/verify.sql +++ b/tests/compatibility/cases/legacy_json2_non_append_table/verify.sql @@ -11,6 +11,10 @@ SHOW CREATE TABLE t_legacy_json2_non_append_table; INSERT INTO t_legacy_json2_non_append_table (ts, j) VALUES ('2026-07-08 00:02:00+0000', '{"a": 3, "nested": {"s": "current-3"}}'); +ADMIN FLUSH_TABLE('t_legacy_json2_non_append_table'); + +ADMIN compact_table('t_legacy_json2_non_append_table'); + SELECT ts, j.a AS a, j.nested.s AS nested_s FROM t_legacy_json2_non_append_table ORDER BY ts;