From acd16af6e2d8ae6876f4a095234e1e7e41ac16e8 Mon Sep 17 00:00:00 2001
From: shyjsarah <44659226+shyjsarah@users.noreply.github.com>
Date: Mon, 12 Oct 2026 02:55:18 -0700
Subject: [PATCH] format_table: fix partition extraction for reordered columns

---
 .../datafusion/tests/rest_format_table_dml.rs | 43 ++++++++++++++++++-
 .../paimon/src/table/format_table_writer.rs   | 10 ++---
 2 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/crates/integrations/datafusion/tests/rest_format_table_dml.rs b/crates/integrations/datafusion/tests/rest_format_table_dml.rs
index 3494b2ba9..487331091 100644
--- a/crates/integrations/datafusion/tests/rest_format_table_dml.rs
+++ b/crates/integrations/datafusion/tests/rest_format_table_dml.rs
@@ -47,6 +47,13 @@ const WAREHOUSE: &str = "test_warehouse";
 /// A Format Table with one file in each of `dt=a` and `dt=b`, whose partitions a mock REST
 /// catalog manages. Its data directory is the root of `temp_dir`.
 async fn catalog_managed_format_table(temp_dir: &TempDir) -> (SQLContext, mock_server::RESTServer) {
+    catalog_managed_format_table_with_schema(temp_dir, format_table_schema(true)).await
+}
+
+async fn catalog_managed_format_table_with_schema(
+    temp_dir: &TempDir,
+    schema: Schema,
+) -> (SQLContext, mock_server::RESTServer) {
     let server = start_mock_server(
         WAREHOUSE.to_string(),
         temp_dir.path().to_string_lossy().into_owned(),
@@ -60,7 +67,7 @@ async fn catalog_managed_format_table(temp_dir: &TempDir) -> (SQLContext, mock_s
     server.add_table_with_schema(
         DATABASE,
         TABLE,
-        format_table_schema(true),
+        schema,
         &format!("file://{}", temp_dir.path().display()),
     );
     server.set_table_external(DATABASE, TABLE, false);
@@ -85,6 +92,19 @@ async fn catalog_managed_format_table(temp_dir: &TempDir) -> (SQLContext, mock_s
     (context, server)
 }
 
+fn partition_last_format_table_schema() -> Schema {
+    Schema::builder()
+        .column("id", DataType::BigInt(BigIntType::new()))
+        .column("name", DataType::VarChar(VarCharType::string_type()))
+        .column("dt", DataType::VarChar(VarCharType::new(255).unwrap()))
+        .partition_keys(["dt"])
+        .option("type", "format-table")
+        .option("file.format", "parquet")
+        .option("metastore.partitioned-table", "true")
+        .build()
+        .unwrap()
+}
+
 /// The same table with its partitions discovered from the directory layout, in a filesystem
 /// catalog. Its data directory is the table directory below the warehouse.
 async fn directory_partitioned_format_table(
@@ -256,6 +276,27 @@ async fn test_insert_into_catalog_managed_format_table_registers_partition() {
     );
 }
 
+#[cfg(not(windows))]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn test_insert_into_catalog_managed_format_table_with_partition_column_last() {
+    let temp_dir = tempfile::tempdir().unwrap();
+    let (context, _server) =
+        catalog_managed_format_table_with_schema(&temp_dir, partition_last_format_table_schema())
+            .await;
+
+    context
+        .sql(&format!(
+            "INSERT INTO {TABLE_NAME} VALUES (11, 'eleven', 'c')"
+        ))
+        .await
+        .unwrap()
+        .collect()
+        .await
+        .unwrap();
+
+    assert_eq!(ids(&context, TABLE_NAME).await, [1, 2, 3, 11]);
+}
+
 #[cfg(not(windows))]
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn test_append_keeps_visible_files_when_statistics_report_fails() {
diff --git a/crates/paimon/src/table/format_table_writer.rs b/crates/paimon/src/table/format_table_writer.rs
index 09aada676..aaf8dfc90 100644
--- a/crates/paimon/src/table/format_table_writer.rs
+++ b/crates/paimon/src/table/format_table_writer.rs
@@ -243,7 +243,7 @@ impl FormatTableWriter {
 
         let mut groups: HashMap<Vec<u8>, Vec<usize>> = HashMap::new();
         let mut specs = HashMap::new();
-        let partition_fields = self.table.schema().partition_fields();
+        let schema_fields = self.table.schema().fields();
         let options = CoreOptions::new(self.table.schema().options());
         let computer = crate::spec::PartitionComputer::new(
             self.table.schema().partition_keys(),
@@ -252,12 +252,8 @@ impl FormatTableWriter {
             options.legacy_partition_name(),
         )?;
         for row_index in 0..batch.num_rows() {
-            let partition = BinaryRow::from_arrow(
-                &batch,
-                row_index,
-                &self.partition_indices,
-                &partition_fields,
-            )?;
+            let partition =
+                BinaryRow::from_arrow(&batch, row_index, &self.partition_indices, schema_fields)?;
             let key = partition.to_serialized_bytes();
             groups.entry(key.clone()).or_default().push(row_index);
             if let std::collections::hash_map::Entry::Vacant(entry) = specs.entry(key) {