fix: preserve Arrow Field metadata across C Data exports

native/core/src/execution/jni_api.rs

@@ -753,6 +753,7 @@ fn prepare_output(
     let schema_addrs = unsafe { schema_addrs.get_elements(env, ReleaseMode::NoCopyBack)? };
     let schema_addrs = &*schema_addrs;
 
+    let output_schema = output_batch.schema();
     let results = output_batch.columns();
     let num_rows = output_batch.num_rows();
 
@@ -779,6 +780,7 @@ fn prepare_output(
         let mut i = 0;
         while i < results.len() {
             let array_ref = results.get(i).ok_or(CometError::IndexOutOfBounds(i))?;
+            let field = output_schema.field(i);
 
             if array_ref.offset() != 0 {
                 // https://github.com/apache/datafusion-comet/issues/2051
@@ -795,11 +797,11 @@ fn prepare_output(
 
                 new_array
                     .to_data()
-                    .move_to_spark(array_addrs[i], schema_addrs[i])?;
+                    .move_to_spark(field, array_addrs[i], schema_addrs[i])?;
             } else {
                 array_ref
                     .to_data()
-                    .move_to_spark(array_addrs[i], schema_addrs[i])?;
+                    .move_to_spark(field, array_addrs[i], schema_addrs[i])?;
             }
             i += 1;
         }

native/core/src/execution/utils.rs

@@ -19,28 +19,34 @@
 use crate::execution::operators::ExecutionError;
 use arrow::{
     array::ArrayData,
+    datatypes::Field,
     ffi::{FFI_ArrowArray, FFI_ArrowSchema},
 };
 
 pub trait SparkArrowConvert {
     /// Move Arrow Arrays to C data interface.
-    fn move_to_spark(&self, array: i64, schema: i64) -> Result<(), ExecutionError>;
+    fn move_to_spark(&self, field: &Field, array: i64, schema: i64) -> Result<(), ExecutionError>;
 }
 
 impl SparkArrowConvert for ArrayData {
     /// Move this ArrowData to pointers of Arrow C data interface.
-    fn move_to_spark(&self, array: i64, schema: i64) -> Result<(), ExecutionError> {
+    fn move_to_spark(&self, field: &Field, array: i64, schema: i64) -> Result<(), ExecutionError> {
         let array_ptr = array as *mut FFI_ArrowArray;
         let schema_ptr = schema as *mut FFI_ArrowSchema;
 
         let array_align = std::mem::align_of::<FFI_ArrowArray>();
         let schema_align = std::mem::align_of::<FFI_ArrowSchema>();
+        let ffi_array = FFI_ArrowArray::new(self);
+        // Spark owns the top-level name and nullability. Preserve the existing anonymous schema
+        // shape while carrying logical extension metadata from the RecordBatch field.
+        let ffi_schema =
+            FFI_ArrowSchema::try_from(self.data_type())?.with_metadata(field.metadata())?;
 
         // Check if the pointer alignment is correct.
         if array_ptr.align_offset(array_align) != 0 || schema_ptr.align_offset(schema_align) != 0 {
             unsafe {
-                std::ptr::write_unaligned(array_ptr, FFI_ArrowArray::new(self));
-                std::ptr::write_unaligned(schema_ptr, FFI_ArrowSchema::try_from(self.data_type())?);
+                std::ptr::write_unaligned(array_ptr, ffi_array);
+                std::ptr::write_unaligned(schema_ptr, ffi_schema);
             }
         } else {
             // SAFETY: `array_ptr` and `schema_ptr` are aligned correctly.
@@ -55,8 +61,8 @@ impl SparkArrowConvert for ArrayData {
                 "move_to_spark: schema_ptr not aligned"
             );
             unsafe {
-                std::ptr::write(array_ptr, FFI_ArrowArray::new(self));
-                std::ptr::write(schema_ptr, FFI_ArrowSchema::try_from(self.data_type())?);
+                std::ptr::write(array_ptr, ffi_array);
+                std::ptr::write(schema_ptr, ffi_schema);
             }
         }
 
@@ -65,3 +71,43 @@ impl SparkArrowConvert for ArrayData {
 }
 
 pub use datafusion_comet_common::bytes_to_i128;
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow::{
+        array::{Array, Int32Array},
+        datatypes::DataType,
+    };
+    use std::{collections::HashMap, mem::MaybeUninit};
+
+    #[test]
+    fn test_move_to_spark_preserves_field_metadata() {
+        let field = Field::new("v", DataType::Int32, true).with_metadata(HashMap::from([(
+            "ARROW:extension:name".to_string(),
+            "example.logical-type".to_string(),
+        )]));
+        let data = Int32Array::from(vec![Some(1), None]).into_data();
+        let mut ffi_array = MaybeUninit::<FFI_ArrowArray>::uninit();
+        let mut ffi_schema = MaybeUninit::<FFI_ArrowSchema>::uninit();
+
+        data.move_to_spark(
+            &field,
+            ffi_array.as_mut_ptr() as i64,
+            ffi_schema.as_mut_ptr() as i64,
+        )
+        .unwrap();
+
+        let ffi_array = unsafe { ffi_array.assume_init() };
+        let ffi_schema = unsafe { ffi_schema.assume_init() };
+        let exported = Field::try_from(&ffi_schema).unwrap();
+
+        assert_eq!(exported.name(), "");
+        assert_eq!(exported.data_type(), field.data_type());
+        assert!(!exported.is_nullable());
+        assert_eq!(exported.metadata(), field.metadata());
+
+        drop(ffi_array);
+        drop(ffi_schema);
+    }
+}

native/core/src/parquet/mod.rs

@@ -308,8 +308,11 @@ pub extern "system" fn Java_org_apache_comet_parquet_Native_currentColumnBatch(
             .ok_or_else(|| CometError::Execution {
                 source: ExecutionError::GeneralError("There is no more data to read".to_string()),
             });
-        let data = batch_reader?.column(column_idx as usize).into_data();
-        data.move_to_spark(array_addr, schema_addr)
+        let batch = batch_reader?;
+        let schema = batch.schema();
+        let field = schema.field(column_idx as usize);
+        let data = batch.column(column_idx as usize).into_data();
+        data.move_to_spark(field, array_addr, schema_addr)
             .map_err(|e| e.into())
     })
 }

spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala

@@ -400,7 +400,18 @@ object Utils extends CometTypeShim with Logging {
                 targetRoot = VectorSchemaRoot.create(sourceRoot.getSchema, allocator)
                 targetRoot.allocateNew()
               }
-              VectorSchemaRootAppender.append(targetRoot, sourceRoot)
+              try {
+                VectorSchemaRootAppender.append(targetRoot, sourceRoot)
+              } catch {
+                case e: IllegalArgumentException =>
+                  logWarning(
+                    "Arrow batches cannot be appended during BroadcastExchange coalescing; " +
+                      "skipping coalesce",
+                    e)
+                  targetRoot.close()
+                  targetRoot = null
+                  return (buffers, 0L, 0L)
+              }
               totalRows += sourceRoot.getRowCount
               batchCount += 1
             }

spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala

@@ -25,9 +25,10 @@ import org.scalatest.Tag
 import org.apache.spark.sql.CometTestBase
 import org.apache.spark.sql.catalyst.TableIdentifier
 import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation
-import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometSortMergeJoinExec}
+import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometSortMergeJoinExec, CometUnionExec}
 import org.apache.spark.sql.execution.adaptive.AQEShuffleReadExec
 import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{ArrayType, IntegerType, MetadataBuilder, StructField, StructType}
 
 import org.apache.comet.CometConf
 
@@ -657,6 +658,80 @@ class CometJoinSuite extends CometTestBase {
     }
   }
 
+  test("Broadcast coalescing handles union children with different nullability") {
+    withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+      withParquetTable(Seq((1, 10), (2, 20), (3, 30)), "t") {
+        val (_, cometPlan) = checkSparkAnswerAndOperator(
+          sql("""
+              |SELECT /*+ BROADCAST(b) */ p._1, b.v
+              |FROM t p JOIN (
+              |  SELECT _1 AS k, 99 AS v FROM t
+              |  UNION ALL
+              |  SELECT _1 AS k, _2 + 1 AS v FROM t
+              |) b ON p._1 = b.k
+              |""".stripMargin),
+          Seq(
+            classOf[CometBroadcastExchangeExec],
+            classOf[CometBroadcastHashJoinExec],
+            classOf[CometUnionExec]))
+
+        val broadcast = collect(cometPlan) { case b: CometBroadcastExchangeExec => b }.head
+        assert(broadcast.metrics("numCoalescedBatches").value > 0L)
+        assert(broadcast.metrics("numCoalescedRows").value == 6L)
+      }
+    }
+  }
+
+  test("Broadcast coalescing falls back for array field metadata mismatch") {
+    withSQLConf(
+      SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+      SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key -> "false",
+      SQLConf.IGNORE_MISSING_PARQUET_FIELD_ID.key -> "true") {
+      withTempPath { dir =>
+        val path = dir.getCanonicalPath
+        Seq((1, Seq(10)), (2, Seq(20))).toDF("k", "v").coalesce(1).write.parquet(path)
+
+        def readWithFieldId(fieldId: Long) = {
+          val metadata = new MetadataBuilder()
+            .putLong("parquet.field.id", fieldId)
+            .build()
+          val schema = StructType(
+            Seq(
+              StructField("k", IntegerType, nullable = true),
+              StructField(
+                "v",
+                ArrayType(IntegerType, containsNull = true),
+                nullable = true,
+                metadata)))
+          spark.read.schema(schema).parquet(path)
+        }
+
+        withTempView("metadata_left", "metadata_right") {
+          readWithFieldId(1).createOrReplaceTempView("metadata_left")
+          readWithFieldId(2).createOrReplaceTempView("metadata_right")
+
+          val (_, cometPlan) = checkSparkAnswerAndOperator(
+            sql("""
+                |SELECT /*+ BROADCAST(u) */ p.k, u.v
+                |FROM metadata_left p JOIN (
+                |  SELECT k, v FROM metadata_left
+                |  UNION ALL
+                |  SELECT k, v FROM metadata_right
+                |) u ON p.k = u.k
+                |""".stripMargin),
+            Seq(
+              classOf[CometBroadcastExchangeExec],
+              classOf[CometBroadcastHashJoinExec],
+              classOf[CometUnionExec]))
+
+          val broadcast = collect(cometPlan) { case b: CometBroadcastExchangeExec => b }.head
+          assert(broadcast.metrics("numCoalescedBatches").value == 0L)
+          assert(broadcast.metrics("numCoalescedRows").value == 0L)
+        }
+      }
+    }
+  }
+
   // Reproducer for SPARK-43113: full outer SMJ with a join filter that references
   // a nullable column should not match when the filter evaluates to NULL.
   test("SPARK-43113: Full outer SMJ with NULL in join filter") {

spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala

@@ -48,6 +48,7 @@ import org.apache.spark.sql.types._
 import com.google.common.primitives.UnsignedLong
 
 import org.apache.comet.CometConf
+import org.apache.comet.vector.CometVector
 
 abstract class ParquetReadSuite extends CometTestBase {
   import testImplicits._
@@ -1664,7 +1665,28 @@ abstract class ParquetReadSuite extends CometTestBase {
           .mode("overwrite")
           .parquet(dir.getCanonicalPath)
         val df = spark.read.schema(readSchema).parquet(dir.getCanonicalPath)
-        checkSparkAnswerAndOperator(df)
+        val (_, cometPlan) = checkSparkAnswerAndOperator(df)
+        val scan = collect(cometPlan) { case scan: CometNativeScanExec => scan }.head
+        val fieldIds = scan
+          .executeColumnar()
+          .mapPartitions { batches =>
+            batches.map { batch =>
+              try {
+                batch
+                  .column(0)
+                  .asInstanceOf[CometVector]
+                  .getValueVector
+                  .getField
+                  .getMetadata
+                  .get(CometParquetUtils.PARQUET_FIELD_ID_META_KEY)
+              } finally {
+                batch.close()
+              }
+            }
+          }
+          .collect()
+          .toSet
+        assert(fieldIds == Set("0"))
       }
     }
   }