fix(mito): preserve mixed JSON2 types during compaction (#9135)

* fix(mito): preserve mixed JSON2 types during compaction

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* 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 <ratuthomm@gmail.com>

---------

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>
(cherry picked from commit 9696882217)
(cherry picked from commit 6ae3d829e4)
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

@@ -52,6 +52,8 @@ pub(crate) type Json2RewritePlans = HashMap<String, Json2RewritePlan>;
 #[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,
         };

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::<Vec<_>>();
+        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()

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::<Vec<_>>()
         );
+        // 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(())
     }