(paimon-mosaic) branch main updated: perf(core): build batch column map in linear time (#90)

[lzljs3620320]

This is an automated email from the ASF dual-hosted git repository.

JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-mosaic.git


The following commit(s) were added to refs/heads/main by this push:
     new f138a0e  perf(core): build batch column map in linear time (#90)
f138a0e is described below

commit f138a0ecb1d48bfd33e3596e5e521f69cde77f1c
Author: jianguotian <841604...@qq.com>
AuthorDate: Sun Sep 20 20:03:51 2026 +0800

    perf(core): build batch column map in linear time (#90)
---
 core/src/writer.rs | 26 +++++++++++++++++++++-----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/core/src/writer.rs b/core/src/writer.rs
index 3a3f922..6de041e 100644
--- a/core/src/writer.rs
+++ b/core/src/writer.rs
@@ -149,15 +149,19 @@ pub struct MosaicWriter<S: OutputFile> {
     state: WriterState,
 }
 
+fn batch_column_map(schema: &MosaicSchema) -> Vec<usize> {
+    let mut batch_col_map = vec![0; schema.original_order.len()];
+    for (input_index, &sorted_index) in 
schema.original_order.iter().enumerate() {
+        batch_col_map[sorted_index] = input_index;
+    }
+    batch_col_map
+}
+
 impl<S: OutputFile> MosaicWriter<S> {
     pub fn new(out: S, schema: &Schema, options: WriterOptions) -> 
io::Result<Self> {
         let mosaic_schema = MosaicSchema::from_arrow(schema, 
options.num_buckets)
             .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
-        let batch_col_map: Vec<usize> = mosaic_schema
-            .columns
-            .iter()
-            .map(|col| schema.index_of(&col.name).unwrap())
-            .collect();
+        let batch_col_map = batch_column_map(&mosaic_schema);
         Self::from_mosaic_schema_with_map(out, mosaic_schema, options, 
batch_col_map)
     }
 
@@ -856,6 +860,18 @@ mod tests {
         assert!(state.bytes.is_empty());
     }
 
+    #[test]
+    fn test_batch_column_map_inverts_original_order() {
+        let arrow_schema = Schema::new(vec![
+            Field::new("zebra", DataType::Int32, true),
+            Field::new("alpha", DataType::Int32, true),
+            Field::new("middle", DataType::Int32, true),
+        ]);
+        let mosaic_schema = MosaicSchema::from_arrow(&arrow_schema, 
2).unwrap();
+
+        assert_eq!(batch_column_map(&mosaic_schema), vec![1, 2, 0]);
+    }
+
     #[test]
     fn test_write_simple_file() {
         let arrow_schema = Schema::new(vec![