andygrove opened a new pull request, #6339: URL: https://github.com/apache/datafusion-comet/pull/6339
## Which issue does this PR close? Closes #6288. ## Rationale for this change Arrow Java's C Data import ignores `ArrowArray.offset` at every level ([apache/arrow-java#88](https://github.com/apache/arrow-java/issues/88)). arrow-rs folds a slice into the buffers for every Spark type except boolean, which keeps its bit offset, so the JVM reads a sliced boolean's values and nulls from bit 0. The fix for #2051 only checked the top level of each output column. A boolean inside a sliced struct, list or map still came back wrong, and the inputs to the JVM UDF bridge were never checked at all. The issue has repros through `LIMIT ... OFFSET`, a grouped aggregate emitting more than one batch, a boolean `ScalaUDF` and a dispatched `regexp_replace`. There is one more path that the issue doesn't list. `collect_list` returns a single run of rows as a slice, because arrow's `concat` of one array returns that array's slice, and that slice keeps its offsets. So a group whose only run starts past row 0 produced a `list<struct<boolean>>` with a sliced boolean two levels down. ## What changes are included in this PR? - `zero_offsets` in `native/common/src/ffi_offsets.rs` gives every level of an array offset 0, including struct children, list and map values, and dictionary values. When no level has an offset, which is almost always, it returns the array unchanged. A sliced boolean has its values bitmap re-sliced to start at bit 0. That shares the buffer when the offset is a whole number of bytes, and copies only the bitmap otherwise. Every other buffer is shared. `FFI_ArrowArray::new` already realigns the validity bitmap to the exported offset. - `move_to_spark` applies it, so it now covers every executed batch and decoded shuffle block. This replaces the top-level `take` in `prepare_output`, and a top-level boolean now costs a bitmap slice instead of a `take`. - `JvmScalarUdfExpr::evaluate` applies it to the arrays it hands the JVM UDF bridge. - `ffi.md` gains a short "Array Offsets" section, so that a new native to JVM export path calls it too. ## How are these changes tested? - Six new Scala tests. `CometExecSuite` covers a struct sliced by `LIMIT ... OFFSET` (with and without nulls), a `named_struct` over sliced aggregate output, and `collect_list` of structs. `CometCodegenSuite` covers a boolean `ScalaUDF` argument, `regexp_replace(IF(b, s, 'zz'), ...)` through the dispatcher, and a struct UDF input with a sliced boolean child. All six fail on `main` with wrong results and pass with this change, on Spark 4.1, 3.5 (Scala 2.12) and 3.4. - Seven unit tests for `zero_offsets` cover a top-level boolean at byte-aligned and unaligned offsets, a struct child, `list<struct<boolean>>`, map values, dictionary values, the copy fallback for any other type with an offset, and the unchanged fast path. The six that rewrite something export through `FFI_ArrowArray`, check that every level has offset 0, and import the array back. A `move_to_spark` test checks the export path itself. Removing the rewrite fails all six, keeping only a top-level check fails exactly the four nested ones, and zeroing the offset without re-slicing the bitmap fails every boolean case. - Locally on Spark 4.1: `CometExecSuite`, `CometCodegenSuite`, `CometUdfBridgeSuite`, `CometAggregateSuite`, `CometTopKSuite`, `CometArrayExpressionSuite`, `CometMapExpressionSuite` and `CometNativeShuffleSuite` all pass. -- 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]
