From a0c8f8978b8f51db4a61746efa20ad4d27217a23 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 30 Aug 2026 01:11:51 +0800 Subject: [PATCH 1/4] fix: preserve Arrow Field metadata across C Data exports --- native/core/src/execution/jni_api.rs | 6 ++-- native/core/src/execution/utils.rs | 52 ++++++++++++++++++++++++---- native/core/src/parquet/mod.rs | 6 ++-- 3 files changed, 54 insertions(+), 10 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 86780d0c908..4cb50b0baf3 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -761,6 +761,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(); @@ -787,6 +788,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 @@ -803,11 +805,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; } diff --git a/native/core/src/execution/utils.rs b/native/core/src/execution/utils.rs index 6195e3f0aea..3e02dc229a0 100644 --- a/native/core/src/execution/utils.rs +++ b/native/core/src/execution/utils.rs @@ -19,28 +19,45 @@ use crate::execution::operators::ExecutionError; use arrow::{ array::ArrayData, + datatypes::Field, + error::ArrowError, ffi::{FFI_ArrowArray, FFI_ArrowSchema}, }; +fn ffi_schema_for_field(field: &Field) -> Result { + if field.name().contains('\0') { + // ArrowSchema names are NUL-terminated C strings. Spark owns the logical output name, so + // substitute only the exported name while retaining the Field's type and metadata. + let field = field + .clone() + .with_name(field.name().replace('\0', "\u{fffd}")); + FFI_ArrowSchema::try_from(&field) + } else { + FFI_ArrowSchema::try_from(field) + } +} + 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::(); let schema_align = std::mem::align_of::(); + let ffi_array = FFI_ArrowArray::new(self); + let ffi_schema = ffi_schema_for_field(field)?; // 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 +72,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 +82,26 @@ impl SparkArrowConvert for ArrayData { } pub use datafusion_comet_common::bytes_to_i128; + +#[cfg(test)] +mod tests { + use super::*; + use arrow::datatypes::DataType; + use std::collections::HashMap; + + #[test] + fn test_ffi_schema_preserves_field_and_sanitizes_nul_name() { + let field = Field::new("v\0tail", DataType::Int32, true).with_metadata(HashMap::from([( + "ARROW:extension:name".to_string(), + "example.logical-type".to_string(), + )])); + + let ffi_schema = ffi_schema_for_field(&field).unwrap(); + let exported = Field::try_from(&ffi_schema).unwrap(); + + assert_eq!(exported.name(), "v\u{fffd}tail"); + assert_eq!(exported.data_type(), field.data_type()); + assert_eq!(exported.is_nullable(), field.is_nullable()); + assert_eq!(exported.metadata(), field.metadata()); + } +} diff --git a/native/core/src/parquet/mod.rs b/native/core/src/parquet/mod.rs index cfa03220c10..ea61fe54ac7 100644 --- a/native/core/src/parquet/mod.rs +++ b/native/core/src/parquet/mod.rs @@ -308,8 +308,10 @@ 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 field = batch.schema().field(column_idx as usize).clone(); + let data = batch.column(column_idx as usize).into_data(); + data.move_to_spark(&field, array_addr, schema_addr) .map_err(|e| e.into()) }) } From 7ef3bbd15261d6af6c19486e35d99e34930567ec Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 30 Aug 2026 02:06:31 +0800 Subject: [PATCH 2/4] fix: skip broadcast coalescing for schema mismatches --- .../apache/spark/sql/comet/util/Utils.scala | 7 +++++ .../apache/comet/exec/CometJoinSuite.scala | 26 ++++++++++++++++++- 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index 769d8058de5..6b95eaa975a 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala @@ -396,6 +396,13 @@ object Utils extends CometTypeShim with Logging { } while (reader.loadNextBatch()) { val sourceRoot = reader.getVectorSchemaRoot + if (targetRoot != null && targetRoot.getSchema != sourceRoot.getSchema) { + logWarning( + "Arrow schemas differ during BroadcastExchange coalescing; skipping coalesce") + targetRoot.close() + targetRoot = null + return (buffers, 0L, 0L) + } if (targetRoot == null) { targetRoot = VectorSchemaRoot.create(sourceRoot.getSchema, allocator) targetRoot.allocateNew() diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index 2da0a4c4d84..bf1c2373403 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -25,7 +25,7 @@ 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 @@ -657,6 +657,30 @@ class CometJoinSuite extends CometTestBase { } } + test("Broadcast coalescing falls back when union children have 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 == 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") { From 754cc2f016017fa21d4f0ac79425761108ac21bb Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 30 Aug 2026 08:55:59 +0800 Subject: [PATCH 3/4] fix: narrow FFI schema export to field metadata --- native/core/src/execution/utils.rs | 50 +++++++++++-------- native/core/src/parquet/mod.rs | 5 +- .../apache/spark/sql/comet/util/Utils.scala | 7 --- .../apache/comet/exec/CometJoinSuite.scala | 6 +-- .../comet/parquet/ParquetReadSuite.scala | 24 ++++++++- 5 files changed, 57 insertions(+), 35 deletions(-) diff --git a/native/core/src/execution/utils.rs b/native/core/src/execution/utils.rs index 3e02dc229a0..4102bfab345 100644 --- a/native/core/src/execution/utils.rs +++ b/native/core/src/execution/utils.rs @@ -20,23 +20,9 @@ use crate::execution::operators::ExecutionError; use arrow::{ array::ArrayData, datatypes::Field, - error::ArrowError, ffi::{FFI_ArrowArray, FFI_ArrowSchema}, }; -fn ffi_schema_for_field(field: &Field) -> Result { - if field.name().contains('\0') { - // ArrowSchema names are NUL-terminated C strings. Spark owns the logical output name, so - // substitute only the exported name while retaining the Field's type and metadata. - let field = field - .clone() - .with_name(field.name().replace('\0', "\u{fffd}")); - FFI_ArrowSchema::try_from(&field) - } else { - FFI_ArrowSchema::try_from(field) - } -} - pub trait SparkArrowConvert { /// Move Arrow Arrays to C data interface. fn move_to_spark(&self, field: &Field, array: i64, schema: i64) -> Result<(), ExecutionError>; @@ -51,7 +37,10 @@ impl SparkArrowConvert for ArrayData { let array_align = std::mem::align_of::(); let schema_align = std::mem::align_of::(); let ffi_array = FFI_ArrowArray::new(self); - let ffi_schema = ffi_schema_for_field(field)?; + // 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 { @@ -86,22 +75,39 @@ pub use datafusion_comet_common::bytes_to_i128; #[cfg(test)] mod tests { use super::*; - use arrow::datatypes::DataType; - use std::collections::HashMap; + use arrow::{ + array::{Array, Int32Array}, + datatypes::DataType, + }; + use std::{collections::HashMap, mem::MaybeUninit}; #[test] - fn test_ffi_schema_preserves_field_and_sanitizes_nul_name() { - let field = Field::new("v\0tail", DataType::Int32, true).with_metadata(HashMap::from([( + 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::::uninit(); + let mut ffi_schema = MaybeUninit::::uninit(); - let ffi_schema = ffi_schema_for_field(&field).unwrap(); + 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(), "v\u{fffd}tail"); + assert_eq!(exported.name(), ""); assert_eq!(exported.data_type(), field.data_type()); - assert_eq!(exported.is_nullable(), field.is_nullable()); + assert!(!exported.is_nullable()); assert_eq!(exported.metadata(), field.metadata()); + + drop(ffi_array); + drop(ffi_schema); } } diff --git a/native/core/src/parquet/mod.rs b/native/core/src/parquet/mod.rs index ea61fe54ac7..8abaf96d779 100644 --- a/native/core/src/parquet/mod.rs +++ b/native/core/src/parquet/mod.rs @@ -309,9 +309,10 @@ pub extern "system" fn Java_org_apache_comet_parquet_Native_currentColumnBatch( source: ExecutionError::GeneralError("There is no more data to read".to_string()), }); let batch = batch_reader?; - let field = batch.schema().field(column_idx as usize).clone(); + 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) + data.move_to_spark(field, array_addr, schema_addr) .map_err(|e| e.into()) }) } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index 6b95eaa975a..769d8058de5 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala @@ -396,13 +396,6 @@ object Utils extends CometTypeShim with Logging { } while (reader.loadNextBatch()) { val sourceRoot = reader.getVectorSchemaRoot - if (targetRoot != null && targetRoot.getSchema != sourceRoot.getSchema) { - logWarning( - "Arrow schemas differ during BroadcastExchange coalescing; skipping coalesce") - targetRoot.close() - targetRoot = null - return (buffers, 0L, 0L) - } if (targetRoot == null) { targetRoot = VectorSchemaRoot.create(sourceRoot.getSchema, allocator) targetRoot.allocateNew() diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index bf1c2373403..d66ec520dbe 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -657,7 +657,7 @@ class CometJoinSuite extends CometTestBase { } } - test("Broadcast coalescing falls back when union children have different nullability") { + 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( @@ -675,8 +675,8 @@ class CometJoinSuite extends CometTestBase { classOf[CometUnionExec])) val broadcast = collect(cometPlan) { case b: CometBroadcastExchangeExec => b }.head - assert(broadcast.metrics("numCoalescedBatches").value == 0L) - assert(broadcast.metrics("numCoalescedRows").value == 0L) + assert(broadcast.metrics("numCoalescedBatches").value > 0L) + assert(broadcast.metrics("numCoalescedRows").value == 6L) } } } diff --git a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala index 53a57ed2abe..8553ef23290 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala +++ b/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")) } } } From ef4b599373c5f9002b142cd9bb7d526d46d40d95 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Mon, 31 Aug 2026 00:14:32 +0800 Subject: [PATCH 4/4] fix: fall back on incompatible broadcast batches --- .../apache/spark/sql/comet/util/Utils.scala | 13 ++++- .../apache/comet/exec/CometJoinSuite.scala | 51 +++++++++++++++++++ 2 files changed, 63 insertions(+), 1 deletion(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index 769d8058de5..f07848b2b38 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/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 } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index d66ec520dbe..7f114aa160e 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -28,6 +28,7 @@ import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation 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 @@ -681,6 +682,56 @@ class CometJoinSuite extends CometTestBase { } } + 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") {