diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index 3105816393..b90e5a9404 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -23,6 +23,7 @@ use std::time::Instant; use bytes::Bytes; use common_telemetry::{debug, error, info}; use datatypes::arrow::datatypes::SchemaRef; +use datatypes::extension::json::is_structured_json_field; use partition::expr::PartitionExpr; use smallvec::{SmallVec, smallvec}; use snafu::ResultExt; @@ -43,6 +44,7 @@ use crate::error::{ }; 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, @@ -805,6 +807,14 @@ fn memtable_flat_sources( let num_ranges = ranges.len(); let mut input_iters = Vec::with_capacity(num_ranges); let mut current_ranges = Vec::new(); + + let has_json2 = schema.fields().iter().any(is_structured_json_field); + let mut json_align_schemas = if has_json2 { + Some(Vec::with_capacity(num_ranges)) + } else { + None + }; + for (_range_id, range) in ranges { if let Some(encoded) = range.encoded() { let max_sequence = range.stats().max_sequence(); @@ -812,6 +822,14 @@ 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); + } + let iter = range.build_record_batch_iter(None, None)?; input_iters.push(iter); let range_rows = range.num_rows(); @@ -839,20 +857,33 @@ fn memtable_flat_sources( .max() .unwrap_or(0); + 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( &schema, options.append_mode, options.merge_mode(), field_column_start, - std::mem::replace(&mut input_iters, Vec::with_capacity(num_ranges)), + input_iters, )?; - flat_sources.sources.push(( - FlatSource::new_iter(schema.clone(), maybe_dedup), - max_sequence, - )); + flat_sources + .sources + .push((FlatSource::new_iter(schema, 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 + }; } } @@ -865,6 +896,10 @@ fn memtable_flat_sources( input_iters.len(), 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()) @@ -888,6 +923,24 @@ 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 diff --git a/src/mito2/src/memtable.rs b/src/mito2/src/memtable.rs index 879bc88f10..7e34cb2af3 100644 --- a/src/mito2/src/memtable.rs +++ b/src/mito2/src/memtable.rs @@ -23,6 +23,7 @@ use std::time::Duration; pub use bulk::part::EncodedBulkPart; use bytes::Bytes; use common_time::Timestamp; +use datatypes::arrow::datatypes::SchemaRef; use datatypes::arrow::record_batch::RecordBatch; use mito_codec::key_values::KeyValue; pub use mito_codec::key_values::KeyValues; @@ -557,6 +558,11 @@ pub trait IterBuilder: Send + Sync { .fail() } + /// Returns a cheap schema hint for record batches yielded by this builder. + fn record_batch_schema_hint(&self) -> Option { + None + } + /// Returns the [EncodedRange] if the range is already encoded into SST. fn encoded_range(&self) -> Option { None @@ -729,6 +735,11 @@ impl MemtableRange { .fail() } + /// Returns a cheap schema hint for record batches yielded by this range. + pub fn record_batch_schema_hint(&self) -> Option { + self.context.builder.record_batch_schema_hint() + } + /// Returns whether the iterator is a record batch iterator. pub fn is_record_batch(&self) -> bool { self.context.builder.is_record_batch() diff --git a/src/mito2/src/memtable/bulk.rs b/src/mito2/src/memtable/bulk.rs index 04e0a5e3da..bba9cea5bc 100644 --- a/src/mito2/src/memtable/bulk.rs +++ b/src/mito2/src/memtable/bulk.rs @@ -16,6 +16,7 @@ pub(crate) mod chunk_reader; pub mod context; +pub(crate) mod json_align; pub mod part; pub mod part_reader; mod row_group_reader; @@ -46,6 +47,7 @@ use tokio::sync::Semaphore; 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, @@ -62,7 +64,6 @@ use crate::read::flat_merge::FlatMergeIterator; use crate::region::options::MergeMode; use crate::sst::parquet::flat_format::field_column_start; use crate::sst::parquet::{DEFAULT_READ_BATCH_SIZE, DEFAULT_ROW_GROUP_SIZE}; -use crate::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; /// Default merge threshold for triggering compaction. const DEFAULT_MERGE_THRESHOLD: usize = 16; @@ -384,8 +385,6 @@ pub struct BulkMemtable { min_timestamp: AtomicI64, max_sequence: AtomicU64, num_rows: AtomicUsize, - /// Cached flat SST arrow schema for memtable compaction. - flat_arrow_schema: SchemaRef, /// Compactor for merging bulk parts compactor: Arc>, /// Dispatcher for scheduling compaction tasks @@ -618,12 +617,6 @@ impl Memtable for BulkMemtable { } fn fork(&self, id: MemtableId, metadata: &RegionMetadataRef) -> MemtableRef { - // Computes the new flat schema based on the new metadata. - let flat_arrow_schema = to_flat_sst_arrow_schema( - metadata, - &FlatSchemaOptions::from_encoding(metadata.primary_key_encoding), - ); - Arc::new(Self { id, config: self.config.clone(), @@ -634,7 +627,6 @@ impl Memtable for BulkMemtable { min_timestamp: AtomicI64::new(i64::MAX), max_sequence: AtomicU64::new(0), num_rows: AtomicUsize::new(0), - flat_arrow_schema, compactor: Arc::new(Mutex::new(MemtableCompactor::new( metadata.region_id, id, @@ -661,7 +653,6 @@ impl Memtable for BulkMemtable { .should_merge_parts(self.config.merge_threshold); if should_merge { compactor.merge_parts( - &self.flat_arrow_schema, &self.parts, &self.metadata, !self.append_mode, @@ -685,11 +676,6 @@ impl BulkMemtable { merge_mode: MergeMode, ) -> Self { let config = config.sanitize(); - let flat_arrow_schema = to_flat_sst_arrow_schema( - &metadata, - &FlatSchemaOptions::from_encoding(metadata.primary_key_encoding), - ); - let region_id = metadata.region_id; Self { id, @@ -701,7 +687,6 @@ impl BulkMemtable { min_timestamp: AtomicI64::new(i64::MAX), max_sequence: AtomicU64::new(0), num_rows: AtomicUsize::new(0), - flat_arrow_schema, compactor: Arc::new(Mutex::new(MemtableCompactor::new(region_id, id, config))), compact_dispatcher, append_mode, @@ -768,7 +753,6 @@ impl BulkMemtable { metadata: self.metadata.clone(), parts: self.parts.clone(), config: self.config.clone(), - flat_arrow_schema: self.flat_arrow_schema.clone(), compactor: self.compactor.clone(), append_mode: self.append_mode, merge_mode: self.merge_mode, @@ -832,6 +816,10 @@ impl IterBuilder for BulkRangeIterBuilder { Ok(Box::new(iter)) } + fn record_batch_schema_hint(&self) -> Option { + Some(self.part.schema()) + } + fn encoded_range(&self) -> Option { None } @@ -864,6 +852,10 @@ impl IterBuilder for MultiBulkRangeIterBuilder { } } + fn record_batch_schema_hint(&self) -> Option { + self.part.schemas().next() + } + fn encoded_range(&self) -> Option { None } @@ -905,6 +897,10 @@ impl IterBuilder for EncodedBulkRangeIterBuilder { } } + fn record_batch_schema_hint(&self) -> Option { + Some(self.part.schema()) + } + fn encoded_range(&self) -> Option { Some(EncodedRange { data: self.part.data().clone(), @@ -1022,6 +1018,20 @@ impl PartToMerge { } } + /// Returns the Arrow schema of the record batches contained in this [`PartToMerge`]. + fn arrow_schema(&self) -> SchemaRef { + match self { + PartToMerge::Bulk { part, .. } => part.schema(), + // A MultiBulkPart is built from batches that have already been aligned, so + // all contained batches are expected to share the same arrow schema. + PartToMerge::Multi { part, .. } => part + .schemas() + .next() + .expect("MultiBulkPart must contain at least one record batch"), + PartToMerge::Encoded { part, .. } => part.schema(), + } + } + /// Creates a record batch iterator for this part. fn create_iterator( self, @@ -1065,7 +1075,6 @@ impl MemtableCompactor { /// Merges parts (bulk and encoded) and then encodes the result. fn merge_parts( &mut self, - arrow_schema: &SchemaRef, bulk_parts: &RwLock, metadata: &RegionMetadataRef, dedup: bool, @@ -1105,7 +1114,6 @@ impl MemtableCompactor { .map(|group| { Self::merge_parts_group( group, - arrow_schema, metadata, dedup, merge_mode, @@ -1139,7 +1147,6 @@ impl MemtableCompactor { /// Merges a group of parts into a single part (either MultiBulkPart or EncodedBulkPart). fn merge_parts_group( parts_to_merge: Vec, - arrow_schema: &SchemaRef, metadata: &RegionMetadataRef, dedup: bool, merge_mode: MergeMode, @@ -1183,10 +1190,12 @@ impl MemtableCompactor { true, )?); - // Creates iterators for all parts to merge. + let aligner = Json2Aligner::try_new(parts_to_merge.iter().map(PartToMerge::arrow_schema))?; + let iterators: Vec = parts_to_merge .into_iter() .filter_map(|part| part.create_iterator(context.clone()).ok().flatten()) + .map(|iter| aligner.wrap_iter(iter)) .collect(); if iterators.is_empty() { @@ -1194,10 +1203,9 @@ impl MemtableCompactor { } let merged_iter = - FlatMergeIterator::new(arrow_schema.clone(), iterators, DEFAULT_READ_BATCH_SIZE)?; + FlatMergeIterator::new(aligner.schema().clone(), iterators, DEFAULT_READ_BATCH_SIZE)?; let boxed_iter: BoxedRecordBatchIterator = if dedup { - // Applies deduplication based on merge mode match merge_mode { MergeMode::LastRow => { let dedup_iter = FlatDedupIterator::new(merged_iter, FlatLastRow::new(false)); @@ -1205,7 +1213,7 @@ impl MemtableCompactor { } MergeMode::LastNonNull => { let field_column_start = - field_column_start(metadata, arrow_schema.fields().len()); + field_column_start(metadata, aligner.schema().fields().len()); let dedup_iter = FlatDedupIterator::new( merged_iter, @@ -1226,7 +1234,7 @@ impl MemtableCompactor { let mut metrics = BulkPartEncodeMetrics::default(); let encoded_part = encoder.encode_record_batch_iter( boxed_iter, - arrow_schema.clone(), + aligner.schema().clone(), min_timestamp, max_timestamp, max_sequence, @@ -1277,8 +1285,6 @@ struct MemCompactTask { parts: Arc>, /// Configuration for the bulk memtable. config: BulkMemtableConfig, - /// Cached flat SST arrow schema - flat_arrow_schema: SchemaRef, /// Compactor for merging bulk parts compactor: Arc>, /// Whether the append mode is enabled @@ -1298,7 +1304,6 @@ impl MemCompactTask { .should_merge_parts(self.config.merge_threshold); if should_merge { compactor.merge_parts( - &self.flat_arrow_schema, &self.parts, &self.metadata, !self.append_mode, @@ -1406,13 +1411,25 @@ impl MemtableBuilder for BulkMemtableBuilder { #[cfg(test)] mod tests { + use api::helper::encode_json_value; + use api::v1::value::ValueData; + use api::v1::{Mutation, Row, Rows, SemanticType}; + use datatypes::data_type::ConcreteDataType; + use datatypes::extension::json::{JsonExtensionType, JsonMetadata}; + use datatypes::json::value::JsonValue; + use datatypes::schema::ColumnSchema; + use datatypes::types::json_type::{JsonNativeType, JsonObjectType}; use mito_codec::row_converter::build_primary_key_codec; + use serde_json::json; + use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder, RegionMetadataRef}; use super::*; use crate::memtable::bulk::part::BulkPartConverter; use crate::read::scan_region::PredicateGroup; use crate::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; - use crate::test_util::memtable_util::{build_key_values_with_ts_seq_values, metadata_for_test}; + use crate::test_util::memtable_util::{ + build_key_values_with_ts_seq_values, metadata_for_test, region_metadata_to_row_schema, + }; fn create_bulk_part_with_converter( k0: &str, @@ -1541,6 +1558,138 @@ mod tests { } } + #[test] + fn test_bulk_memtable_compact_parts_with_json2() { + let metadata = mock_metadata_with_json2(); + + let config = BulkMemtableConfig { + merge_threshold: 2, + encode_row_threshold: 1, + encode_bytes_threshold: 1, + ..Default::default() + }; + let memtable = BulkMemtable::new( + 999, + config, + metadata.clone(), + None, + None, + true, + MergeMode::LastRow, + ); + memtable.set_unordered_part_threshold(0); + + let part1 = mock_bulk_part_with_json2(&metadata, vec![1000, 2000], 100).unwrap(); + let part2 = mock_bulk_part_with_json2(&metadata, vec![3000, 4000], 200).unwrap(); + + memtable.write_bulk(part1).unwrap(); + memtable.write_bulk(part2).unwrap(); + memtable.compact(false).unwrap(); + + let stats = memtable.stats(); + assert_eq!(4, stats.num_rows); + assert_eq!(201, stats.max_sequence); + + let predicate_group = PredicateGroup::new(&metadata, &[]).unwrap(); + let opts = RangesOptions::default().with_predicate(predicate_group); + let ranges = memtable.ranges(None, opts).unwrap(); + + assert_eq!(1, ranges.ranges.len()); + let total_rows: usize = ranges.ranges.values().map(|r| r.stats().num_rows()).sum(); + assert_eq!(4, total_rows); + } + + fn mock_metadata_with_json2() -> RegionMetadataRef { + let col_meta_1 = ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 0, + }; + + let data_type = ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())); + let mut col_schema = ColumnSchema::new("data", data_type, true); + let extension = JsonExtensionType::new(Arc::new(JsonMetadata::default())); + col_schema.with_extension_type(&extension).unwrap(); + + let col_meta_2 = ColumnMetadata { + column_schema: col_schema, + semantic_type: SemanticType::Field, + column_id: 1, + }; + let mut builder = RegionMetadataBuilder::new(RegionId::new(123, 789)); + builder + .push_column_metadata(col_meta_1) + .push_column_metadata(col_meta_2) + .primary_key(vec![]); + Arc::new(builder.build().unwrap()) + } + + fn mock_bulk_part_with_json2( + metadata: &RegionMetadataRef, + timestamps: Vec, + sequence: u64, + ) -> Result { + let capacity = timestamps.len(); + let primary_key_codec = build_primary_key_codec(metadata); + let json_type = JsonNativeType::Object(JsonObjectType::from([ + ("id".to_string(), JsonNativeType::i64()), + ( + "payload".to_string(), + JsonNativeType::Object(JsonObjectType::from([( + "message".to_string(), + JsonNativeType::String, + )])), + ), + ])); + let mut options = FlatSchemaOptions::from_encoding(metadata.primary_key_encoding); + options + .concretized_json_types + .insert("data".to_string(), json_type.as_arrow_type()); + let schema = to_flat_sst_arrow_schema(metadata, &options); + + let mut converter = + BulkPartConverter::new(metadata, schema, capacity, primary_key_codec, true); + + let rows = timestamps + .into_iter() + .map(|ts| { + 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 val2 = api::v1::Value { + value_data: Some(value_data), + }; + Row { + values: vec![val1, val2], + } + }) + .collect(); + + let mutation = Mutation { + op_type: 1, + sequence, + rows: Some(Rows { + schema: region_metadata_to_row_schema(metadata), + rows, + }), + write_hint: None, + }; + let key_values = KeyValues::new(metadata.as_ref(), mutation).unwrap(); + + converter.append_key_values(&key_values)?; + converter.convert() + } + #[test] fn test_bulk_memtable_ranges_with_projection() { let metadata = metadata_for_test(); diff --git a/src/mito2/src/memtable/bulk/json_align.rs b/src/mito2/src/memtable/bulk/json_align.rs new file mode 100644 index 0000000000..12d7c81fb5 --- /dev/null +++ b/src/mito2/src/memtable/bulk/json_align.rs @@ -0,0 +1,403 @@ +// 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::data_type::DataType; +use datatypes::extension::json::is_structured_json_field; +use datatypes::types::JsonType; +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: HashMap = base_schema + .fields() + .iter() + .enumerate() + .filter(|&(_idx, field)| is_structured_json_field(field)) + .map(|(idx, field)| (idx, JsonType::from(field.data_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; + } + merged + .merge(&JsonType::from(schema.field(*idx).data_type())) + .context(DataTypeMismatchSnafu)?; + } + } + + // 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)) + .try_align(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_structured_json_field(base_field); + let is_json2 = is_structured_json_field(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, Int64Array, StringViewArray, StructArray, UInt64Array, + }; + use datatypes::arrow::datatypes::{DataType, Field, Fields, Schema}; + use datatypes::extension::json::{JsonExtensionType, JsonMetadata}; + + 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::new(Arc::new(JsonMetadata::default()))), + ); + 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_structured_json_field(&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_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(JsonExtensionType::new(Arc::new(JsonMetadata::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 ed7b8addf1..92cf2666df 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -32,17 +32,15 @@ use datatypes::arrow; use datatypes::arrow::array::{ Array, ArrayRef, BinaryArray, BooleanArray, StringDictionaryBuilder, UInt8Array, UInt64Array, }; -use datatypes::arrow::compute::{SortColumn, SortOptions}; +use datatypes::arrow::compute::{SortColumn, SortOptions, concat_batches}; use datatypes::arrow::datatypes::{ DataType as ArrowDataType, Field, Schema, SchemaRef, UInt32Type, }; use datatypes::data_type::DataType; use datatypes::extension::json::is_structured_json_field; use datatypes::prelude::{MutableVector, Vector}; -use datatypes::types::JsonType; use datatypes::value::ValueRef; use datatypes::vectors::Helper; -use datatypes::vectors::json::array::JsonArray; use mito_codec::key_values::{KeyValue, KeyValues}; use mito_codec::row_converter::{PrimaryKeyCodec, SortField, build_primary_key_codec_with_fields}; use parquet::arrow::ArrowWriter; @@ -57,11 +55,12 @@ use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; use store_api::storage::{ColumnId, FileId, SequenceNumber, SequenceRange}; use crate::error::{ - self, ColumnNotFoundSnafu, ComputeArrowSnafu, ConvertValueSnafu, CreateDefaultSnafu, - DataTypeMismatchSnafu, EncodeMemtableSnafu, EncodeSnafu, InvalidMetadataSnafu, - InvalidRequestSnafu, NewRecordBatchSnafu, Result, + 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}; @@ -150,6 +149,10 @@ impl TryFrom<&BulkPart> for BulkWalEntry { } impl BulkPart { + pub(crate) fn schema(&self) -> SchemaRef { + self.batch.schema() + } + pub(crate) fn estimated_size(&self) -> usize { record_batch_estimated_size(&self.batch) } @@ -435,10 +438,12 @@ impl UnorderedPart { // Get the schema from the first part let schema = self.parts[0].batch.schema(); let concatenated = if schema.fields().iter().any(is_structured_json_field) { - let (schema, batches) = align_parts(&self.parts)?; - arrow::compute::concat_batches(&schema, &batches).context(ComputeArrowSnafu)? + 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 { - arrow::compute::concat_batches(&schema, self.parts.iter().map(|x| &x.batch)) + concat_batches(&schema, self.parts.iter().map(|x| &x.batch)) .context(ComputeArrowSnafu)? }; @@ -477,73 +482,6 @@ impl UnorderedPart { } } -/// Align the JSON columns in [BulkPart]s, to unified Arrow arrays. So that we can compute (concat, -/// sort, etc.) on them. -fn align_parts(parts: &[BulkPart]) -> Result<(SchemaRef, Vec)> { - debug_assert!( - !parts.is_empty() - && parts - .windows(2) - .all(|w| w[0].batch.schema_ref().fields().len() - == w[1].batch.schema_ref().fields().len()) - ); - - let first = &parts[0]; - let base_schema = first.batch.schema_ref(); - let rest = &parts[1..]; - - let mut merged_types = HashMap::new(); - let mut aligned_fields = Vec::with_capacity(base_schema.fields().len()); - for (i, field) in base_schema.fields().iter().enumerate() { - if is_structured_json_field(field) { - let mut merged = JsonType::from(field.data_type()); - rest.iter() - .try_fold(&mut merged, |acc, x| { - acc.merge(&JsonType::from(x.batch.schema_ref().field(i).data_type()))?; - Ok(acc) - }) - .context(DataTypeMismatchSnafu)?; - merged_types.insert(i, merged.as_arrow_type()); - - aligned_fields.push(Arc::new( - Field::new( - field.name().clone(), - merged.as_arrow_type(), - field.is_nullable(), - ) - .with_metadata(field.metadata().clone()), - )); - } else { - aligned_fields.push(field.clone()) - }; - } - let aligned_schema = Arc::new(Schema::new_with_metadata( - aligned_fields, - base_schema.metadata().clone(), - )); - - let mut aligned_batches = Vec::with_capacity(parts.len()); - for part in parts { - let mut columns = Vec::with_capacity(part.batch.num_columns()); - for (i, column) in part.batch.columns().iter().enumerate() { - if let Some(expect) = merged_types.get(&i) { - columns.push( - JsonArray::from(column) - .try_align(expect) - .context(ConvertValueSnafu)?, - ); - } else { - columns.push(column.clone()); - } - } - aligned_batches.push( - RecordBatch::try_new(aligned_schema.clone(), columns).context(NewRecordBatchSnafu)?, - ); - } - - Ok((aligned_schema, aligned_batches)) -} - /// More accurate estimation of the size of a record batch. pub fn record_batch_estimated_size(batch: &RecordBatch) -> usize { batch @@ -1045,17 +983,27 @@ pub fn convert_bulk_part( pub struct EncodedBulkPart { data: Bytes, metadata: BulkPartMeta, + /// Cached Arrow schema to avoid rebuilding it from parquet metadata. + schema: SchemaRef, } impl EncodedBulkPart { - pub fn new(data: Bytes, metadata: BulkPartMeta) -> Self { - Self { data, metadata } + pub fn new(data: Bytes, metadata: BulkPartMeta, schema: SchemaRef) -> Self { + Self { + data, + metadata, + schema, + } } pub fn metadata(&self) -> &BulkPartMeta { &self.metadata } + pub(crate) fn schema(&self) -> SchemaRef { + self.schema.clone() + } + /// Returns the size of the encoded data in bytes pub(crate) fn size_bytes(&self) -> usize { self.data.len() @@ -1227,8 +1175,9 @@ impl BulkPartEncoder { metrics: &mut BulkPartEncodeMetrics, ) -> Result> { let mut buf = Vec::with_capacity(4096); - let mut writer = ArrowWriter::try_new(&mut buf, arrow_schema, self.writer_props.clone()) - .context(EncodeMemtableSnafu)?; + let mut writer = + ArrowWriter::try_new(&mut buf, arrow_schema.clone(), self.writer_props.clone()) + .context(EncodeMemtableSnafu)?; let mut total_rows = 0; let mut series_estimator = SeriesEstimator::default(); @@ -1276,6 +1225,7 @@ impl BulkPartEncoder { num_series, max_sequence, }, + schema: arrow_schema, })) } @@ -1290,7 +1240,7 @@ impl BulkPartEncoder { let file_metadata = { let mut writer = - ArrowWriter::try_new(&mut buf, arrow_schema, self.writer_props.clone()) + ArrowWriter::try_new(&mut buf, arrow_schema.clone(), self.writer_props.clone()) .context(EncodeMemtableSnafu)?; writer.write(&part.batch).context(EncodeMemtableSnafu)?; writer.finish().context(EncodeMemtableSnafu)? @@ -1310,6 +1260,7 @@ impl BulkPartEncoder { num_series: part.estimated_series_count() as u64, max_sequence: part.sequence, }, + schema: arrow_schema, })) } } @@ -1587,6 +1538,10 @@ impl MultiBulkPart { self.total_rows } + pub(crate) fn schemas(&self) -> impl Iterator + '_ { + self.batches.iter().map(|batch| batch.schema()) + } + /// Returns the minimum timestamp. pub fn min_timestamp(&self) -> i64 { self.min_timestamp diff --git a/src/mito2/src/memtable/bulk/row_group_reader.rs b/src/mito2/src/memtable/bulk/row_group_reader.rs index 36e8e8aac1..100701514b 100644 --- a/src/mito2/src/memtable/bulk/row_group_reader.rs +++ b/src/mito2/src/memtable/bulk/row_group_reader.rs @@ -15,6 +15,7 @@ use std::sync::Arc; use bytes::Bytes; +use datatypes::extension::json::is_structured_json_field; use parquet::arrow::ProjectionMask; use parquet::arrow::arrow_reader::{ ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReader, @@ -43,8 +44,25 @@ impl MemtableRowGroupReaderBuilder { data: Bytes, ) -> error::Result { // Create ArrowReaderMetadata for building the reader. - let arrow_reader_options = - ArrowReaderOptions::new().with_schema(context.read_format().arrow_schema().clone()); + let mut arrow_reader_options = ArrowReaderOptions::new(); + + // JSON2 columns in region metadata are not concretized with nested + // fields here, so let parquet use its embedded Arrow schema instead. + // + // TODO: Pass the encoded part's concrete schema into this builder. It + // avoids parsing the Arrow schema from parquet metadata while keeping + // JSON2 nested fields exact. + if !context + .read_format() + .arrow_schema() + .fields() + .iter() + .any(is_structured_json_field) + { + arrow_reader_options = + arrow_reader_options.with_schema(context.read_format().arrow_schema().clone()); + } + let arrow_metadata = ArrowReaderMetadata::try_new(parquet_metadata.clone(), arrow_reader_options) .context(ReadDataPartSnafu)?; diff --git a/tests-integration/tests/json2.rs b/tests-integration/tests/json2.rs new file mode 100644 index 0000000000..22bcc26266 --- /dev/null +++ b/tests-integration/tests/json2.rs @@ -0,0 +1,356 @@ +// 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::path::Path; + +use sqlx::{Connection, Executor, MySqlConnection, Row}; +use tests_integration::test_util::{StorageType, setup_mysql_server}; + +type Json2Rows = Vec<( + String, + Option, + Option, + Option, + Option, +)>; + +#[tokio::test(flavor = "multi_thread")] +async fn test_json2_single_mem_range_flush() { + common_telemetry::init_default_ut_logging(); + + let (mut guard, server) = + setup_mysql_server(StorageType::File, "test_json2_single_range_flush").await; + let addr = server.bind_addr().unwrap(); + let mut conn = MySqlConnection::connect(&format!("mysql://{addr}/public")) + .await + .unwrap(); + + create_json2_table(&mut conn, "json2_single_mem_range", 100) + .await + .unwrap(); + + // One INSERT statement produces one bulk part, so flush sees exactly one + // mem range. The rows still carry different JSON2 shapes, exercising + // alignment inside that range. + conn.execute( + r#" +INSERT INTO json2_single_mem_range VALUES + (1, 'host1', '{"payload":{"id":1},"metric":1}'), + (2, 'host2', '{"payload":{"name":"n2"},"flag":true}') +"#, + ) + .await + .unwrap(); + conn.execute("ADMIN FLUSH_TABLE('json2_single_mem_range')") + .await + .unwrap(); + + assert_eq!( + vec![ + ( + "host1".to_string(), + Some("1".to_string()), + None, + Some("1".to_string()), + None, + ), + ( + "host2".to_string(), + None, + Some("n2".to_string()), + None, + Some("true".to_string()) + ), + ], + query_json2_rows(&mut conn, "json2_single_mem_range") + .await + .unwrap() + ); + + let _ = server.shutdown().await; + guard.remove_all().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_json2_multi_mem_range_flush() { + common_telemetry::init_default_ut_logging(); + + let (mut guard, server) = + setup_mysql_server(StorageType::File, "test_json2_multi_range_flush").await; + let addr = server.bind_addr().unwrap(); + let mut conn = MySqlConnection::connect(&format!("mysql://{addr}/public")) + .await + .unwrap(); + + create_json2_table(&mut conn, "json2_multiple_mem_ranges", 100) + .await + .unwrap(); + + // Separate INSERT statements produce separate bulk parts. The high merge + // threshold keeps them unmerged, so flush must align JSON2 schemas across + // multiple mem ranges. + conn.execute( + r#"INSERT INTO json2_multiple_mem_ranges VALUES + (1, 'host1', '{"payload":{"id":1},"metric":1}')"#, + ) + .await + .unwrap(); + conn.execute( + r#"INSERT INTO json2_multiple_mem_ranges VALUES + (2, 'host2', '{"payload":{"name":"n2"},"flag":true}')"#, + ) + .await + .unwrap(); + conn.execute("ADMIN FLUSH_TABLE('json2_multiple_mem_ranges')") + .await + .unwrap(); + + assert_eq!( + vec![ + ( + "host1".to_string(), + Some("1".to_string()), + None, + Some("1".to_string()), + None, + ), + ( + "host2".to_string(), + None, + Some("n2".to_string()), + None, + Some("true".to_string()) + ), + ], + query_json2_rows(&mut conn, "json2_multiple_mem_ranges") + .await + .unwrap() + ); + + conn.execute( + r#"INSERT INTO json2_multiple_mem_ranges VALUES + (3, 'host3', '{"payload":{"id":3,"name":"n3"},"metric":3,"flag":false}')"#, + ) + .await + .unwrap(); + conn.execute( + r#"INSERT INTO json2_multiple_mem_ranges VALUES + (4, 'host4', '{"payload":{"extra":"e4"},"metric":4}')"#, + ) + .await + .unwrap(); + conn.execute("ADMIN FLUSH_TABLE('json2_multiple_mem_ranges')") + .await + .unwrap(); + assert_eq!(2, count_parquet_files(guard.home_guard.temp_dir.path())); + + conn.execute("ADMIN COMPACT_TABLE('json2_multiple_mem_ranges')") + .await + .unwrap(); + + assert_eq!( + vec![ + ( + "host1".to_string(), + Some("1".to_string()), + None, + Some("1".to_string()), + None, + ), + ( + "host2".to_string(), + None, + Some("n2".to_string()), + None, + Some("true".to_string()) + ), + ( + "host3".to_string(), + Some("3".to_string()), + Some("n3".to_string()), + Some("3".to_string()), + Some("false".to_string()) + ), + ("host4".to_string(), None, None, Some("4".to_string()), None,), + ], + query_json2_rows(&mut conn, "json2_multiple_mem_ranges") + .await + .unwrap() + ); + + let _ = server.shutdown().await; + guard.remove_all().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_json2_multi_row_insert() { + common_telemetry::init_default_ut_logging(); + + const NUM_ROWS: usize = 1024; + const TABLE_NAME: &str = "json2_multi_row_insert"; + + let (mut guard, server) = + setup_mysql_server(StorageType::File, "test_json2_multi_row_insert").await; + let addr = server.bind_addr().unwrap(); + let mut conn = MySqlConnection::connect(&format!("mysql://{addr}/public")) + .await + .unwrap(); + + create_json2_compaction_table(&mut conn, TABLE_NAME) + .await + .unwrap(); + + for i in 0..NUM_ROWS { + let json = json2_payload(i); + let sql = format!( + r#"INSERT INTO {TABLE_NAME} VALUES + ({}, 'host{}', '{}')"#, + i + 1, + i, + json + ); + conn.execute(sql.as_str()).await.unwrap(); + } + + assert_eq!( + NUM_ROWS as i64, + count_table_rows(&mut conn, TABLE_NAME).await.unwrap() + ); + + let _ = server.shutdown().await; + guard.remove_all().await; +} + +async fn create_json2_table( + conn: &mut MySqlConnection, + table_name: &str, + merge_threshold: usize, +) -> sqlx::Result<()> { + conn.execute( + format!( + r#" +CREATE TABLE {table_name} ( + ts TIMESTAMP TIME INDEX, + host STRING PRIMARY KEY, + j JSON2 +) WITH ( + 'append_mode' = 'true', + 'sst_format' = 'flat', + 'memtable.type' = 'bulk', + 'memtable.bulk.merge_threshold' = '{merge_threshold}', + 'memtable.bulk.encode_row_threshold' = '1000000' +) +"# + ) + .as_str(), + ) + .await?; + Ok(()) +} + +async fn create_json2_compaction_table( + conn: &mut MySqlConnection, + table_name: &str, +) -> sqlx::Result<()> { + conn.execute( + format!( + r#" +CREATE TABLE {table_name} ( + ts TIMESTAMP TIME INDEX, + host STRING PRIMARY KEY, + j JSON2 +) WITH ( + 'append_mode' = 'true', + 'sst_format' = 'flat', + 'memtable.type' = 'bulk', + 'memtable.bulk.merge_threshold' = '8', + 'memtable.bulk.encode_row_threshold' = '64', + 'memtable.bulk.encode_bytes_threshold' = '100000000' +) +"# + ) + .as_str(), + ) + .await?; + Ok(()) +} + +async fn query_json2_rows(conn: &mut MySqlConnection, table_name: &str) -> sqlx::Result { + let rows = sqlx::query( + format!( + r#" +SELECT + host, + j.payload.id::STRING AS payload_id, + j.payload.name::STRING AS payload_name, + j.metric::STRING AS metric, + j.flag::STRING AS flag +FROM {table_name} +ORDER BY ts +"# + ) + .as_str(), + ) + .fetch_all(conn) + .await?; + + Ok(rows + .into_iter() + .map(|row| { + ( + row.get::("host"), + row.get::, _>("payload_id"), + row.get::, _>("payload_name"), + row.get::, _>("metric"), + row.get::, _>("flag"), + ) + }) + .collect()) +} + +async fn count_table_rows(conn: &mut MySqlConnection, table_name: &str) -> sqlx::Result { + let row = sqlx::query(format!("SELECT COUNT(*) AS count FROM {table_name}").as_str()) + .fetch_one(conn) + .await?; + Ok(row.get("count")) +} + +fn json2_payload(i: usize) -> String { + match i % 4 { + 0 => format!(r#"{{"payload":{{"id":{i}}},"metric":{i}}}"#), + 1 => format!(r#"{{"payload":{{"name":"n{i}"}},"flag":true}}"#), + 2 => format!(r#"{{"payload":{{"score":{}.5}},"tags":["a","b"]}}"#, i), + _ => format!(r#"{{"payload":{{"extra":"e{i}"}},"metric":{i},"flag":false}}"#), + } +} + +fn count_parquet_files(path: &Path) -> usize { + let Ok(entries) = std::fs::read_dir(path) else { + return 0; + }; + + entries + .filter_map(Result::ok) + .map(|entry| entry.path()) + .map(|path| { + if path.is_dir() { + count_parquet_files(&path) + } else if path.extension().is_some_and(|ext| ext == "parquet") { + 1 + } else { + 0 + } + }) + .sum() +} diff --git a/tests-integration/tests/main.rs b/tests-integration/tests/main.rs index e8b8232ad8..f14538bc88 100644 --- a/tests-integration/tests/main.rs +++ b/tests-integration/tests/main.rs @@ -18,6 +18,7 @@ mod grpc; #[macro_use] mod http; +mod json2; mod jsonbench; #[macro_use] mod sql;