From 27dbbd87eb83dbd86184b7df2e6ac96e2a344a4a Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Mon, 14 Sep 2026 03:38:53 +0000 Subject: [PATCH] fix(mito): preserve mixed JSON2 types during compaction (#9135) * fix(mito): preserve mixed JSON2 types during compaction Signed-off-by: Lei, HUANG * test(mito): assert all restored JSON2 rows in unordered merge test The regression test for aligning JSON2 layouts across unordered bulk parts only round-tripped the first part's `a` values. A merge that dropped or corrupted the second source's `b` values would still pass. Assert the merged batch has four rows and that the restored values of the second part (`{"b": 3}`, `{"b": 4}`) survive the merge, so the test covers both opaque sources as its comment claims. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG (cherry picked from commit 96968822172b36f437b9b94b6f2e106da93b6ec8) (cherry picked from commit 6ae3d829e49028d9c2675ea6500d850bade065f4) Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> (cherry picked from commit e0c87713993b156c64c71d94940848afa87d54fb) Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/mito2/src/compaction/json2.rs | 17 +++- src/mito2/src/engine/scan_test.rs | 161 ++++++++++++++++++++++++++++++ src/mito2/src/memtable/bulk.rs | 13 ++- 3 files changed, 187 insertions(+), 4 deletions(-) diff --git a/src/mito2/src/compaction/json2.rs b/src/mito2/src/compaction/json2.rs index fa7ec52f71..0c1d1af450 100644 --- a/src/mito2/src/compaction/json2.rs +++ b/src/mito2/src/compaction/json2.rs @@ -52,6 +52,8 @@ pub(crate) type Json2RewritePlans = HashMap; #[derive(Clone)] struct Json2LeafPathStats { rows: u64, + /// Number of input schemas that explicitly represent this leaf. + sources: usize, data_type: JsonNativeType, is_type_conflicted: bool, } @@ -59,8 +61,8 @@ struct Json2LeafPathStats { /// 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 +/// shared by every input SST schema are ranked once to produce a fixed layout; paths that may +/// reside 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 @@ -144,6 +146,14 @@ pub(crate) fn collect_json2_rewrite_plans( collect_json2_path_stats(field, *rows, &hint_paths, &mut stats)?; } + // A leaf absent from an input's explicit schema may have arbitrary values + // (including scalar ancestors) in its remainder. Only common explicit leaves + // are safe to promote without inspecting rows. Keep unsafe paths opaque for + // the entire output SST so existing readers never miss remainder values when + // projecting an explicit leaf. + for stat in stats.values_mut() { + stat.is_type_conflicted |= stat.sources != schemas.len(); + } 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)?; @@ -187,12 +197,14 @@ fn collect_json2_path_stats<'a>( path, Json2LeafPathStats { rows, + sources: 1, data_type, is_type_conflicted: false, }, ); continue; }; + stat.sources += 1; if stat.data_type != data_type { stat.is_type_conflicted = true; } else { @@ -368,6 +380,7 @@ mod tests { )?; let stat = |rows, data_type, is_type_conflicted| Json2LeafPathStats { rows, + sources: 1, data_type, is_type_conflicted, }; diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index 2011ef3a20..8280611390 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -369,6 +369,167 @@ async fn test_json2_v1_region_reopen_and_compaction() -> WhateverResult<()> { Ok(()) } +#[tokio::test] +async fn test_json2_mixed_subject_compaction_preserves_values() -> WhateverResult<()> { + let request = CreateRequestBuilder::new() + .field_datatype(ConcreteDataType::json2(JsonNativeType::Object( + JsonObjectType::new(), + ))) + .insert_option("append_mode", "true") + .insert_option("memtable.type", "bulk") + .insert_option("sst_format", "flat") + .build(); + 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 { + min_compaction_interval: std::time::Duration::from_secs(3600), + ..Default::default() + }) + .await; + let region_id = RegionId::new(1027, 0); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await?; + // #9133: the second SST stores a heterogeneous subject in its remainder. + let values = [ + json!({"subject": {"cid": "first", "uri": "at://first"}}), + json!({"subject": "did:plc:x"}), + json!({"subject": {"cid": "last", "uri": "at://last"}}), + json!({"subject": {"cid": 42}}), + ]; + for (offset, batch) in [(0, &values[..1]), (1, &values[1..])] { + test_util::put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: batch + .iter() + .enumerate() + .map(|(i, value)| { + row(vec![ + ValueData::StringValue("tag".into()), + ValueData::JsonValue(encode_json_value(JsonValue::from(value.clone()))), + ValueData::TimestampMillisecondValue((offset + i) as i64 * 1000), + ]) + }) + .collect(), + }, + ) + .await; + test_util::flush_region(&engine, region_id, None).await; + } + let old_ids = engine + .scanner(region_id, ScanRequest::default()) + .await? + .file_ids(); + assert_eq!(2, old_ids.len()); + for _ in 0..2 { + engine + .handle_request( + region_id, + RegionRequest::Compact(RegionCompactRequest { + options: compact_request::Options::StrictWindow(StrictWindow { + window_seconds: 86400, + }), + ..Default::default() + }), + ) + .await?; + let scanner = engine + .scanner( + region_id, + ScanRequest { + projection: Some(vec![1]), + ..Default::default() + }, + ) + .await?; + assert_eq!( + 1, + scanner.num_files(), + "a swallowed merge failure must not pass" + ); + assert!(scanner.file_ids().iter().all(|id| !old_ids.contains(id))); + 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(); + for i in 0..array.len() { + actual.push(JsonArray::from(&array).try_get_value(i)?); + } + } + assert_eq!(values.as_slice(), actual.as_slice()); + reopen_region(&engine, region_id, table_dir.clone(), true, HashMap::new()).await; + let scanner = engine + .scanner( + region_id, + ScanRequest { + projection: Some(vec![1]), + json_type_hint: HashMap::from([( + "field_0".into(), + JsonNativeType::Object(JsonObjectType::from([( + "subject".into(), + JsonNativeType::String, + )])), + )]), + ..Default::default() + }, + ) + .await?; + let batches = RecordBatches::try_collect(scanner.scan().await?).await?; + let mut subjects = Vec::new(); + for batch in batches.iter() { + let array = batch.column_by_name("field_0").unwrap().clone(); + for i in 0..array.len() { + subjects.push(JsonArray::from(&array).try_get_value(i)?); + } + } + let expected = values.iter().map(|value| { + let subject = &value["subject"]; + json!({"subject": subject.as_str().map(str::to_string).unwrap_or_else(|| subject.to_string())}) + }).collect::>(); + assert_eq!( + expected, subjects, + "projection must read values spilled to remainder" + ); + let scanner = engine + .scanner( + region_id, + ScanRequest { + projection: Some(vec![1]), + json_type_hint: HashMap::from([( + "field_0".into(), + JsonNativeType::Object(JsonObjectType::from([( + "subject".into(), + JsonNativeType::Object(JsonObjectType::from([( + "cid".into(), + JsonNativeType::String, + )])), + )])), + )]), + ..Default::default() + }, + ) + .await?; + let batches = RecordBatches::try_collect(scanner.scan().await?).await?; + let mut cids = Vec::new(); + for batch in batches.iter() { + let array = batch.column_by_name("field_0").unwrap().clone(); + for i in 0..array.len() { + cids.push(JsonArray::from(&array).try_get_value(i)?["subject"]["cid"].clone()); + } + } + assert_eq!( + vec![json!("first"), json!(null), json!("last"), json!("42")], + cids + ); + } + Ok(()) +} + #[tokio::test] async fn test_flush_aligns_different_json2_layouts() -> WhateverResult<()> { let mut request = CreateRequestBuilder::new() diff --git a/src/mito2/src/memtable/bulk.rs b/src/mito2/src/memtable/bulk.rs index 8e1346b992..6cb0feaecf 100644 --- a/src/mito2/src/memtable/bulk.rs +++ b/src/mito2/src/memtable/bulk.rs @@ -1807,7 +1807,7 @@ mod tests { let second = mock_bulk_part_with_json2_values( &metadata, vec![3000, 4000], - vec![json!({"b": 3}), json!({"b": 4})], + vec![json!({"a": 3, "b": 3}), json!({"a": 4, "b": 4})], 200, )?; let parts = vec![ @@ -1879,9 +1879,18 @@ mod tests { unreachable!() }; assert_eq!( - vec![JSON2_REMAINDER_FIELD_NAME, "a", "b"], + vec![JSON2_REMAINDER_FIELD_NAME], fields.iter().map(|x| x.name().as_str()).collect::>() ); + // Neither path is explicit in every source; keep both opaque without losing values. + assert_eq!(4, batch.num_rows()); + let json = datatypes::vectors::json::array::JsonArray::from(batch.column(0)); + let restored = json.project_to_v2(schema.field(0), &ArrowDataType::Binary)?; + let restored = datatypes::vectors::json::array::JsonArray::from(&restored); + assert_eq!(json!({"a": 1}), restored.try_get_value(0)?); + assert_eq!(json!({"a": 2}), restored.try_get_value(1)?); + assert_eq!(json!({"b": 3}), restored.try_get_value(2)?); + assert_eq!(json!({"b": 4}), restored.try_get_value(3)?); Ok(()) }