andygrove commented on code in PR #5552:
URL: https://github.com/apache/datafusion-comet/pull/5552#discussion_r3887457382
##########
native/core/src/execution/utils.rs:
##########
@@ -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() {
Review Comment:
This test mostly exercises `FFI_ArrowSchema::try_from` and the three-line
NUL replacement. It does not go through `move_to_spark`, either alignment
branch, or the actual JVM boundary.
The behavior the PR exists to establish, that field metadata survives from a
native `RecordBatch` into a Spark-side `CometVector`, has no coverage yet.
Would it be worth adding a round trip in `NativeUtilSuite` that asserts
`getValueVector.getField.getMetadata` on the imported vector? Without one, a
future change back to a `DataType`-derived schema would leave every test here
green.
##########
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<FFI_ArrowSchema, ArrowError> {
Review Comment:
The PR body says Spark owns the logical output name, and what #5546 needs
from this boundary is the Arrow extension metadata. Given that, is there a
reason to export the whole `Field` rather than just the metadata? Something
like:
```rust
Field::new("", field.data_type().clone(),
true).with_metadata(field.metadata().clone())
```
would carry what the Variant work needs while leaving the observable shape
of the boundary identical to today.
That matters because the name and the nullability are the two parts that
cost something here. The name is why the NUL sanitizer had to be invented at
all. Before this PR the exported name was always `""` and could never contain a
NUL, so this change creates the failure mode and then patches it. The
nullability is what broke broadcast coalescing in the earlier review round, and
now needs a guard in `Utils.scala` plus a regression test that asserts the
optimization is off.
If a later subtask in #5546 genuinely needs the real name or nullability,
could you say which one in the description? Then the extra surface is clearly
paying for something. If not, narrowing to metadata would let the NUL helper
and the Scala guard both go away.
One more thing if the name export does stay. `FFI_ArrowSchema::try_from`
recurses into children and calls `CString::new` on every nested field name too,
so a struct or list whose *child* name contains a NUL still fails with
`ArrowError::CDataInterface`. If NUL names are reachable enough to handle at
the top level, should the sanitization walk the whole field tree so both cases
behave the same way?
##########
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) {
Review Comment:
Two things here.
The check runs once per batch, but the schema is fixed for the life of an
`ArrowStreamReader`. `getVectorSchemaRoot` reads the schema message when the
stream is initialized, so this could move up to just after the reader is
constructed and run once per buffer instead. `Schema.equals` walks every field
and every metadata map, so on a wide relation with many broadcast batches this
is repeated work for no added coverage.
On the guard itself, this compares full `Schema` equality, which now covers
names and field metadata as well as nullability. So it fires on more than the
union nullability case that motivated it. A broadcast over a union whose
branches project differently named source columns would trip it, since Spark
takes output names from the first child but each native child plan names its
own output fields. Same for a union where one branch carries Variant metadata
and the other does not, which is exactly where #5546 is heading.
How common do you expect that to be on real plans? Right now losing
coalescing surfaces only as a `logWarning`, with no metric or `EXPLAIN` signal
a user could act on.
##########
spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala:
##########
@@ -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
Review Comment:
Good regression coverage for the crash, and I like that it pins the exact
shape from the earlier review round.
Worth flagging that asserting zero here pins the optimization *off* for this
query. If the export gets narrowed to metadata only as suggested on the Rust
side, this case should coalesce normally again and the assertion would want to
flip.
##########
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();
Review Comment:
`.field(idx).clone()` deep-clones the metadata `HashMap` for every column of
every batch on the scan path. Binding the schema first lets you pass a
reference and skip the clone entirely:
```rust
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())
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
