From 2a8d068bc15a320b82e1a44625c36d0437c023c7 Mon Sep 17 00:00:00 2001
From: luofucong <luofc@foxmail.com>
Date: Mon, 11 May 2026 09:50:57 +0800
Subject: [PATCH] resolve PR comments

Signed-off-by: luofucong <luofc@foxmail.com>
---
 src/mito2/src/compaction.rs         | 3 ++-
 src/mito2/src/test_util/sst_util.rs | 9 ++++++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/mito2/src/compaction.rs b/src/mito2/src/compaction.rs
index a9e2d0668e..1d1151177d 100644
--- a/src/mito2/src/compaction.rs
+++ b/src/mito2/src/compaction.rs
@@ -992,7 +992,8 @@ struct CompactionSstReaderBuilder<'a> {
 }
 
 impl CompactionSstReaderBuilder<'_> {
-    /// Builds [BoxedRecordBatchStream] that reads all SST files and yields batches in flat format for compaction.
+    /// Build a [FlatSource] that yields Arrow `RecordBatch`s from reading all the input SST files,
+    /// for compaction. The schema of the [FlatSource] is unified.
     async fn build_flat_sst_reader(self) -> Result<FlatSource> {
         let scan_input = self.build_scan_input()?.with_compaction(true);
 
diff --git a/src/mito2/src/test_util/sst_util.rs b/src/mito2/src/test_util/sst_util.rs
index 56cf8d8441..3f635c6e94 100644
--- a/src/mito2/src/test_util/sst_util.rs
+++ b/src/mito2/src/test_util/sst_util.rs
@@ -17,6 +17,7 @@
 use std::sync::Arc;
 
 use api::v1::{OpType, SemanticType};
+use arrow_schema::Schema;
 use common_time::Timestamp;
 use datatypes::arrow::array::{
     ArrayRef, BinaryDictionaryBuilder, RecordBatch, StringDictionaryBuilder,
@@ -307,7 +308,13 @@ pub fn new_record_batch_with_custom_sequence(
 
 /// Creates a FlatSource from flat format RecordBatches.
 pub(crate) fn new_flat_source_from_record_batches(batches: Vec<RecordBatch>) -> FlatSource {
-    FlatSource::new_iter(batches[0].schema(), Box::new(batches.into_iter().map(Ok)))
+    FlatSource::new_iter(
+        batches
+            .first()
+            .map(|x| x.schema())
+            .unwrap_or_else(|| Arc::new(Schema::empty())),
+        Box::new(batches.into_iter().map(Ok)),
+    )
 }
 
 /// Creates a new region metadata for testing SSTs with binary datatype.