diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 4d4094b1472..ea114ba23d2 100644 --- a/native/core/src/execution/jni_api.rs +++ b/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; } diff --git a/native/core/src/execution/utils.rs b/native/core/src/execution/utils.rs index 6195e3f0aea..4102bfab345 100644 --- a/native/core/src/execution/utils.rs +++ b/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::(); let schema_align = std::mem::align_of::(); + 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::::uninit(); + let mut ffi_schema = MaybeUninit::::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); + } +} diff --git a/native/core/src/parquet/mod.rs b/native/core/src/parquet/mod.rs index cfa03220c10..8abaf96d779 100644 --- a/native/core/src/parquet/mod.rs +++ b/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()) }) } 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 2da0a4c4d84..7f114aa160e 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/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") { 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")) } } }