Skip to content

Commit 754cc2f

Browse files
peterxclipeterxcli
committed
fix: narrow FFI schema export to field metadata
1 parent 7ef3bbd commit 754cc2f

5 files changed

Lines changed: 57 additions & 35 deletions

File tree

‎native/core/src/execution/utils.rs‎

Lines changed: 28 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -20,23 +20,9 @@ use crate::execution::operators::ExecutionError;
2020
use arrow::{
2121
array::ArrayData,
2222
datatypes::Field,
23-
error::ArrowError,
2423
ffi::{FFI_ArrowArray, FFI_ArrowSchema},
2524
};
2625

27-
fn ffi_schema_for_field(field: &Field) -> Result<FFI_ArrowSchema, ArrowError> {
28-
if field.name().contains('\0') {
29-
// ArrowSchema names are NUL-terminated C strings. Spark owns the logical output name, so
30-
// substitute only the exported name while retaining the Field's type and metadata.
31-
let field = field
32-
.clone()
33-
.with_name(field.name().replace('\0', "\u{fffd}"));
34-
FFI_ArrowSchema::try_from(&field)
35-
} else {
36-
FFI_ArrowSchema::try_from(field)
37-
}
38-
}
39-
4026
pub trait SparkArrowConvert {
4127
/// Move Arrow Arrays to C data interface.
4228
fn move_to_spark(&self, field: &Field, array: i64, schema: i64) -> Result<(), ExecutionError>;
@@ -51,7 +37,10 @@ impl SparkArrowConvert for ArrayData {
5137
let array_align = std::mem::align_of::<FFI_ArrowArray>();
5238
let schema_align = std::mem::align_of::<FFI_ArrowSchema>();
5339
let ffi_array = FFI_ArrowArray::new(self);
54-
let ffi_schema = ffi_schema_for_field(field)?;
40+
// Spark owns the top-level name and nullability. Preserve the existing anonymous schema
41+
// shape while carrying logical extension metadata from the RecordBatch field.
42+
let ffi_schema =
43+
FFI_ArrowSchema::try_from(self.data_type())?.with_metadata(field.metadata())?;
5544

5645
// Check if the pointer alignment is correct.
5746
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;
8675
#[cfg(test)]
8776
mod tests {
8877
use super::*;
89-
use arrow::datatypes::DataType;
90-
use std::collections::HashMap;
78+
use arrow::{
79+
array::{Array, Int32Array},
80+
datatypes::DataType,
81+
};
82+
use std::{collections::HashMap, mem::MaybeUninit};
9183

9284
#[test]
93-
fn test_ffi_schema_preserves_field_and_sanitizes_nul_name() {
94-
let field = Field::new("v\0tail", DataType::Int32, true).with_metadata(HashMap::from([(
85+
fn test_move_to_spark_preserves_field_metadata() {
86+
let field = Field::new("v", DataType::Int32, true).with_metadata(HashMap::from([(
9587
"ARROW:extension:name".to_string(),
9688
"example.logical-type".to_string(),
9789
)]));
90+
let data = Int32Array::from(vec![Some(1), None]).into_data();
91+
let mut ffi_array = MaybeUninit::<FFI_ArrowArray>::uninit();
92+
let mut ffi_schema = MaybeUninit::<FFI_ArrowSchema>::uninit();
9893

99-
let ffi_schema = ffi_schema_for_field(&field).unwrap();
94+
data.move_to_spark(
95+
&field,
96+
ffi_array.as_mut_ptr() as i64,
97+
ffi_schema.as_mut_ptr() as i64,
98+
)
99+
.unwrap();
100+
101+
let ffi_array = unsafe { ffi_array.assume_init() };
102+
let ffi_schema = unsafe { ffi_schema.assume_init() };
100103
let exported = Field::try_from(&ffi_schema).unwrap();
101104

102-
assert_eq!(exported.name(), "v\u{fffd}tail");
105+
assert_eq!(exported.name(), "");
103106
assert_eq!(exported.data_type(), field.data_type());
104-
assert_eq!(exported.is_nullable(), field.is_nullable());
107+
assert!(!exported.is_nullable());
105108
assert_eq!(exported.metadata(), field.metadata());
109+
110+
drop(ffi_array);
111+
drop(ffi_schema);
106112
}
107113
}

‎native/core/src/parquet/mod.rs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -309,9 +309,10 @@ pub extern "system" fn Java_org_apache_comet_parquet_Native_currentColumnBatch(
309309
source: ExecutionError::GeneralError("There is no more data to read".to_string()),
310310
});
311311
let batch = batch_reader?;
312-
let field = batch.schema().field(column_idx as usize).clone();
312+
let schema = batch.schema();
313+
let field = schema.field(column_idx as usize);
313314
let data = batch.column(column_idx as usize).into_data();
314-
data.move_to_spark(&field, array_addr, schema_addr)
315+
data.move_to_spark(field, array_addr, schema_addr)
315316
.map_err(|e| e.into())
316317
})
317318
}

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

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -396,13 +396,6 @@ object Utils extends CometTypeShim with Logging {
396396
}
397397
while (reader.loadNextBatch()) {
398398
val sourceRoot = reader.getVectorSchemaRoot
399-
if (targetRoot != null && targetRoot.getSchema != sourceRoot.getSchema) {
400-
logWarning(
401-
"Arrow schemas differ during BroadcastExchange coalescing; skipping coalesce")
402-
targetRoot.close()
403-
targetRoot = null
404-
return (buffers, 0L, 0L)
405-
}
406399
if (targetRoot == null) {
407400
targetRoot = VectorSchemaRoot.create(sourceRoot.getSchema, allocator)
408401
targetRoot.allocateNew()

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -657,7 +657,7 @@ class CometJoinSuite extends CometTestBase {
657657
}
658658
}
659659

660-
test("Broadcast coalescing falls back when union children have different nullability") {
660+
test("Broadcast coalescing handles union children with different nullability") {
661661
withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
662662
withParquetTable(Seq((1, 10), (2, 20), (3, 30)), "t") {
663663
val (_, cometPlan) = checkSparkAnswerAndOperator(
@@ -675,8 +675,8 @@ class CometJoinSuite extends CometTestBase {
675675
classOf[CometUnionExec]))
676676

677677
val broadcast = collect(cometPlan) { case b: CometBroadcastExchangeExec => b }.head
678-
assert(broadcast.metrics("numCoalescedBatches").value == 0L)
679-
assert(broadcast.metrics("numCoalescedRows").value == 0L)
678+
assert(broadcast.metrics("numCoalescedBatches").value > 0L)
679+
assert(broadcast.metrics("numCoalescedRows").value == 6L)
680680
}
681681
}
682682
}

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

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ import org.apache.spark.sql.types._
4848
import com.google.common.primitives.UnsignedLong
4949

5050
import org.apache.comet.CometConf
51+
import org.apache.comet.vector.CometVector
5152

5253
abstract class ParquetReadSuite extends CometTestBase {
5354
import testImplicits._
@@ -1664,7 +1665,28 @@ abstract class ParquetReadSuite extends CometTestBase {
16641665
.mode("overwrite")
16651666
.parquet(dir.getCanonicalPath)
16661667
val df = spark.read.schema(readSchema).parquet(dir.getCanonicalPath)
1667-
checkSparkAnswerAndOperator(df)
1668+
val (_, cometPlan) = checkSparkAnswerAndOperator(df)
1669+
val scan = collect(cometPlan) { case scan: CometNativeScanExec => scan }.head
1670+
val fieldIds = scan
1671+
.executeColumnar()
1672+
.mapPartitions { batches =>
1673+
batches.map { batch =>
1674+
try {
1675+
batch
1676+
.column(0)
1677+
.asInstanceOf[CometVector]
1678+
.getValueVector
1679+
.getField
1680+
.getMetadata
1681+
.get(CometParquetUtils.PARQUET_FIELD_ID_META_KEY)
1682+
} finally {
1683+
batch.close()
1684+
}
1685+
}
1686+
}
1687+
.collect()
1688+
.toSet
1689+
assert(fieldIds == Set("0"))
16681690
}
16691691
}
16701692
}

0 commit comments

Comments
 (0)