fix: reorder MySQL result columns to match projected schema (#568)

When SqlTable is created with new_with_schema, the projected schema
field order may differ from MySQL's physical column order. After
upgrading to DataFusion 52, the new BatchCoalescer assumes column
order matches the plan schema and panics on type mismatch.

Reorder the RecordBatch columns by name to match the projected
schema's field order in rows_to_arrow.

Author: hozan23

core/src/sql/arrow_sql_gen/mysql.rs

@@ -58,6 +58,9 @@ pub enum Error {
 
     #[snafu(display("No column name for index: {index}"))]
     NoColumnNameForIndex { index: usize },
+
+    #[snafu(display("Projected schema field \"{field_name}\" not found in query result"))]
+    FieldNotFoundInResult { field_name: String },
 }
 
 pub type Result<T, E = Error> = std::result::Result<T, E>;
@@ -567,8 +570,57 @@ pub fn rows_to_arrow(rows: &[Row], projected_schema: &Option<SchemaRef>) -> Resu
         .collect::<Vec<ArrayRef>>();
     let arrow_fields = arrow_fields.into_iter().flatten().collect::<Vec<Field>>();
     let options = &RecordBatchOptions::new().with_row_count(Some(rows.len()));
-    RecordBatch::try_new_with_options(Arc::new(Schema::new(arrow_fields)), columns, options)
-        .map_err(|err| Error::FailedToBuildRecordBatch { source: err })
+    let batch =
+        RecordBatch::try_new_with_options(Arc::new(Schema::new(arrow_fields)), columns, options)
+            .map_err(|err| Error::FailedToBuildRecordBatch { source: err })?;
+
+    // When SqlTable is created with new_with_schema, the projected schema field
+    // order may differ from MySQL's physical column order. Reorder the result
+    // columns to match the projected schema. After upgrading to DataFusion 52,
+    // the new BatchCoalescer assumes column order matches the plan schema and
+    // panics on type mismatch.
+    if let Some(schema) = projected_schema {
+        let batch_schema = batch.schema();
+        let needs_reorder = schema
+            .fields()
+            .iter()
+            .zip(batch_schema.fields())
+            .any(|(expected, actual)| expected.name() != actual.name());
+
+        if needs_reorder {
+            let mut reordered_columns = Vec::with_capacity(schema.fields().len());
+            let mut reordered_fields = Vec::with_capacity(schema.fields().len());
+
+            for field in schema.fields() {
+                let idx = batch_schema.index_of(field.name()).map_err(|_| {
+                    Error::FieldNotFoundInResult {
+                        field_name: field.name().clone(),
+                    }
+                })?;
+
+                reordered_columns.push(batch.column(idx).clone());
+                reordered_fields.push(batch_schema.field(idx).clone());
+            }
+
+            for (idx, field) in batch_schema.fields().iter().enumerate() {
+                if schema.index_of(field.name()).is_err() {
+                    reordered_columns.push(batch.column(idx).clone());
+                    reordered_fields.push(field.as_ref().clone());
+                }
+            }
+
+            let options = &RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
+
+            return RecordBatch::try_new_with_options(
+                Arc::new(Schema::new(reordered_fields)),
+                reordered_columns,
+                options,
+            )
+            .map_err(|err| Error::FailedToBuildRecordBatch { source: err });
+        }
+    }
+
+    Ok(batch)
 }
 
 #[allow(clippy::unnecessary_wraps)]

core/tests/mysql/mod.rs

@@ -821,6 +821,83 @@ async fn test_arrow_mysql_roundtrip(
     .await;
 }
 
+/// When SqlTable is created with new_with_schema, the projected schema field
+/// order may differ from MySQL's physical column order. rows_to_arrow must
+/// reorder the result columns to match the projected schema. After upgrading
+/// to DataFusion 52, the new BatchCoalescer assumes column order matches the
+/// plan schema and panics on type mismatch.
+async fn test_mysql_projected_schema_column_reorder(port: usize) {
+    let create_table_stmt = "
+CREATE TABLE reorder_table (
+    a INT,
+    b VARCHAR(50),
+    c DOUBLE
+);
+        ";
+    let insert_table_stmt = "
+INSERT INTO reorder_table (a, b, c) VALUES (1, 'hello', 3.14);
+        ";
+
+    // Schema with reversed column order compared to MySQL's physical order
+    let reordered_schema = Arc::new(Schema::new(vec![
+        Field::new("c", DataType::Float64, true),
+        Field::new("b", DataType::Utf8, true),
+        Field::new("a", DataType::Int32, true),
+    ]));
+
+    let ctx = SessionContext::new();
+    let pool = common::get_mysql_connection_pool(port)
+        .await
+        .expect("MySQL connection pool should be created");
+
+    let db_conn = pool
+        .connect_direct()
+        .await
+        .expect("Connection should be established");
+
+    let _ = db_conn
+        .execute(create_table_stmt, &[])
+        .await
+        .expect("MySQL table should be created");
+
+    let _ = db_conn
+        .execute(insert_table_stmt, &[])
+        .await
+        .expect("MySQL table data should be inserted");
+
+    let sqltable_pool: Arc<
+        dyn DbConnectionPool<mysql_async::Conn, &'static (dyn ToValue + Sync)>
+            + Send
+            + Sync
+            + 'static,
+    > = Arc::new(pool);
+
+    // Use new_with_schema with a different column order than the physical table
+    let table = SqlTable::new_with_schema(
+        "mysql",
+        &sqltable_pool,
+        reordered_schema.clone(),
+        "reorder_table",
+    );
+
+    ctx.register_table("reorder_table", Arc::new(table))
+        .expect("Table should be registered");
+
+    let df = ctx
+        .sql("SELECT * FROM reorder_table")
+        .await
+        .expect("DataFrame should be created from query");
+
+    let record_batch = df.collect().await.expect("RecordBatch should be collected");
+    assert_eq!(record_batch.len(), 1);
+
+    let batch = &record_batch[0];
+    // Verify columns are in the projected schema order (c, b, a)
+    assert_eq!(batch.schema().field(0).name(), "c");
+    assert_eq!(batch.schema().field(1).name(), "b");
+    assert_eq!(batch.schema().field(2).name(), "a");
+}
+
 #[rstest]
 #[test_log::test(tokio::test)]
 async fn test_mysql_arrow_oneway() {
@@ -836,6 +913,7 @@ async fn test_mysql_arrow_oneway() {
     test_mysql_decimal_types_to_decimal128(port).await;
     test_mysql_decimal_types_to_decimal256(port).await;
     test_mysql_zero_date_type(port).await;
+    test_mysql_projected_schema_column_reorder(port).await;
 
     mysql_container.remove().await.expect("container to stop");
 }