rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-07-07 06:20:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(mito): failed to compact memtable with json2 (#8297)

Browse Source 
* fix(json2): failed to compact memtable

* fix: cargo clippy

* refactor: align schema with json2 filed in flush

* chore: add unit test for json aligner

* chore: add json2 integration test

* fix: cr by codex

* fix: use parquet schema for encoded JSON2 memtable parts

* Use is_structured_json_field to determine whether the field is of JSON2 type.

* fix: cargo clippy

* fix: only align structured json fields

* chore: assert bulk JSON2 aligner input schemas in debug
This commit is contained in:
fys
2026-06-26 16:52:06 +08:00
committed by GitHub
parent e89d591cb4
commit cce70b9427
8 changed files with 1064 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
src/mito2/src/flush.rs
View File

@@ -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<Vec<SchemaRef>>,
input_iters: Vec<BoxedRecordBatchIterator>,
) -> Result<(SchemaRef, Vec<BoxedRecordBatchIterator>)> {
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

11
src/mito2/src/memtable.rs
View File

@@ -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<SchemaRef> {
None
}
/// Returns the [EncodedRange] if the range is already encoded into SST.
fn encoded_range(&self) -> Option<EncodedRange> {
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<SchemaRef> {
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()

209
src/mito2/src/memtable/bulk.rs
View File

@@ -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<Mutex<MemtableCompactor>>,
/// 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<SchemaRef> {
Some(self.part.schema())
}
fn encoded_range(&self) -> Option<EncodedRange> {
None
}
@@ -864,6 +852,10 @@ impl IterBuilder for MultiBulkRangeIterBuilder {
}
}
fn record_batch_schema_hint(&self) -> Option<SchemaRef> {
self.part.schemas().next()
}
fn encoded_range(&self) -> Option<EncodedRange> {
None
}
@@ -905,6 +897,10 @@ impl IterBuilder for EncodedBulkRangeIterBuilder {
}
}
fn record_batch_schema_hint(&self) -> Option<SchemaRef> {
Some(self.part.schema())
}
fn encoded_range(&self) -> Option<EncodedRange> {
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<BulkParts>,
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<PartToMerge>,
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<BoxedRecordBatchIterator> = 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<RwLock<BulkParts>>,
/// Configuration for the bulk memtable.
config: BulkMemtableConfig,
/// Cached flat SST arrow schema
flat_arrow_schema: SchemaRef,
/// Compactor for merging bulk parts
compactor: Arc<Mutex<MemtableCompactor>>,
/// 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<i64>,
sequence: u64,
) -> Result<BulkPart> {
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();

403
src/mito2/src/memtable/bulk/json_align.rs Normal file
View File

@@ -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<I>(input_schemas: I) -> Result<Self>
where
I: IntoIterator<Item = SchemaRef>,
{
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<usize, JsonType> = 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<RecordBatch> {
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<I>(&self, batches: I) -> Result<Vec<RecordBatch>>
where
I: IntoIterator<Item = RecordBatch>,
{
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::<StructArray>()
.unwrap();
let id_values = data_with_id
.column(0)
.as_any()
.downcast_ref::<Int64Array>()
.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::<StructArray>()
.unwrap();
let missing_ids = data_with_name.column(0);
let name_values = data_with_name
.column(1)
.as_any()
.downcast_ref::<StringViewArray>()
.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::<Result<Vec<_>>>().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<Field> {
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<Field>) -> SchemaRef {
Arc::new(Schema::new(vec![
Arc::new(Field::new("ts", DataType::Int64, false)),
json_field,
]))
}
fn id_field() -> Arc<Field> {
Arc::new(Field::new("id", DataType::Int64, true))
}
fn name_field() -> Arc<Field> {
Arc::new(Field::new("name", DataType::Utf8View, true))
}
fn struct_array(fields: Fields, columns: Vec<ArrayRef>) -> ArrayRef {
Arc::new(StructArray::new(fields, columns, None))
}
}

117
src/mito2/src/memtable/bulk/part.rs
View File

@@ -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<RecordBatch>)> {
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<Option<EncodedBulkPart>> {
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<Item = SchemaRef> + '_ {
self.batches.iter().map(|batch| batch.schema())
}
/// Returns the minimum timestamp.
pub fn min_timestamp(&self) -> i64 {
self.min_timestamp

22
src/mito2/src/memtable/bulk/row_group_reader.rs
View File

@@ -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<Self> {
// 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)?;

356
tests-integration/tests/json2.rs Normal file
View File

@@ -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<String>,
Option<String>,
Option<String>,
Option<String>,
)>;
#[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<Json2Rows> {
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::<String, _>("host"),
row.get::<Option<String>, _>("payload_id"),
row.get::<Option<String>, _>("payload_name"),
row.get::<Option<String>, _>("metric"),
row.get::<Option<String>, _>("flag"),
)
})
.collect())
}
async fn count_table_rows(conn: &mut MySqlConnection, table_name: &str) -> sqlx::Result<i64> {
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()
}

1
tests-integration/tests/main.rs
View File

@@ -18,6 +18,7 @@
mod grpc;
#[macro_use]
mod http;
mod json2;
mod jsonbench;
#[macro_use]
mod sql;