andygrove opened a new pull request, #7405:
URL: https://github.com/apache/datafusion-comet/pull/7405
## Which issue does this PR close?
No issue is open for this. #1115 raised the cost of sending the schema with
every batch for batches going from the JVM to native, and #4572 removed it
there by moving that direction to the Arrow C Stream Interface. This PR removes
it for batches going from native to the JVM.
## Rationale for this change
`executePlan` hands each output batch to the JVM as one `ArrowArray` and one
`ArrowSchema` per column. A plan's output schema rarely changes between
batches, but native exported it with every batch, building each column's format
string, name, metadata and, for a nested type, its children. The JVM then
allocated a struct for it, parsed it back into a `Field`, and made a JNI call
to release it. Measured with the benchmark in this PR, that costs about 0.5 µs
of CPU per flat column and 1.2 µs per `struct<bigint,string>` column in every
batch, on top of exporting and importing the arrays themselves.
Exporting the output as an Arrow C Stream, the way #4572 handles input,
would also send the schema once. But Arrow Java's `ArrowArrayStreamReader`
loads every batch through `Data.importIntoVectorSchemaRoot`, which builds a
temporary `StructVector` and then loads its buffers into a reused root. In a
standalone measurement of both paths (arrow-rs 59.3, Arrow Java 18.3), that
made the stream slower than the current path for single-column and nested
outputs. A stream would also turn native errors into `IOException` text, the
problem #6243 worked around for input, and it would reuse one root for every
batch, which changes who owns the batches `CometExecIterator` returns. Keeping
`executePlan` and dropping only the repeated schema leaves errors, metrics and
batch ownership as they are.
## What changes are included in this PR?
- Native (`jni_api.rs`): the `ExecutionContext` records each output column's
data type and field metadata as `executePlan` last exported them. These two are
what `move_to_spark` builds the `FFI_ArrowSchema` from. `prepare_plan_output`
compares each batch against them, and while every column matches it writes only
the arrays, with the new `move_array_to_spark`, which still applies
`zero_offsets`. Otherwise it writes the schemas too, records the new ones, and
sets `OUTPUT_SCHEMA_EXPORTED` (bit 32) in the row count it returns. A batch
with a dictionary anywhere in its types always carries its schema and clears
the record, because the JVM gives each dictionary it imports a new id and keeps
the values under that id.
- JVM: `NativeUtil.getNextPlanBatch` allocates the `ArrowSchema` structs
once and passes the same, still empty, structs with every batch. When native
sets the flag it imports the batch as `getNextBatch` does, which consumes the
structs, and keeps the batch's fields. Otherwise it imports the arrays with
those fields through the new `ArrowImporter.importVector(array, field,
provider)`. `close()` frees the schema structs that are left. Decoded shuffle
blocks still carry a schema each.
- The internal config `spark.comet.exec.cacheOutputSchema.enabled`, on by
default, turns this off. The JVM and native both read it when the plan is
created, so they agree. It is there for the benchmark and as a fallback.
- `ffi.md` describes the protocol.
- The second commit adds `CometPlanOutputBenchmark`, which compares the two
settings. Each case reads one column of a Parquet table and projects it 1, 10
or 100 times, so the native plan does almost nothing per output column beyond
exporting it, and the JVM only counts the rows of each batch. One more case
reads 100 distinct columns. Before timing, the benchmark checks that both
settings return every row and that the plan sent its schema once per task with
the cache on. Besides wall time it reports each setting's process CPU time per
batch: for a plan with no JVM input the native plan runs on a Tokio worker
while the task thread imports the batches, so wall time hides the saving until
the JVM side is the slower one, but a busy executor has no idle core to hide it
on.
## How are these changes tested?
New `NativeUtilSuite` tests check that:
- batches that arrive without a schema are imported with the fields sent
last, and are passed the same schema structs;
- the structs native filled before failing are released;
- the schema structs left at the end of the output are freed by `close()`;
- a batch without a schema before any schema was sent fails with an error
naming the problem;
- a native plan whose output takes at least 10 batches sends its schema once
per task, and still matches Spark.
New Rust tests cover `move_array_to_spark` read back through a schema
exported for an earlier array, the column comparison (type and metadata count,
name and nullability don't), and `has_dictionary` on nested types.
Locally, `NativeUtilSuite`, `CometExecIteratorLifecycleSuite`,
`CometNativeSuite`, `SparkErrorConverterSuite`, `CometArrowStreamSuite`,
`CometNativeShuffleSuite` and `CometExecSuite` pass (366 tests), and so does
`cargo clippy --all-targets --workspace -- -D warnings`.
`CometPlanOutputBenchmark` on macOS (Apple M3 Max, JDK 17.0.10, `local[1]`,
an optimized native build without LTO). Wall time is the best of each case's
runs, with the cache off, on, and off again to show drift. CPU is the median
process CPU time per batch over seven alternating samples. Each case has 1025
batches of 8192 rows, except the last, which has 129.
| Output columns | Off | On | Off (repeat) | CPU per batch, off | CPU per
batch, on |
| --- | ---: | ---: | ---: | ---: | ---: |
| 100 × `bigint`, 1 column read | 214 ms | 178 ms | 232 ms | 315 µs | 265 µs
|
| 100 × `string`, 1 column read | 269 ms | 211 ms | 271 ms | 395 µs | 341 µs
|
| 100 × `struct<bigint,string>`, 1 column read | 627 ms | 514 ms | 622 ms |
802 µs | 680 µs |
| 100 distinct `bigint`, 100 columns read | 479 ms | 482 ms | 477 ms | 4354
µs | 4321 µs |
An earlier run gave 1.26x, 1.27x and 1.28x for the first three rows. With 1
or 10 output columns the two settings are within run-to-run noise on this
machine: the expected saving is 0.5 to 12 µs on batches that cost 80 to 275 µs
of CPU. In the last row the saving, 34 µs per batch, is under 1% of decoding
100 columns.
To run it:
```shell
SPARK_GENERATE_BENCHMARK_FILES=1 make
benchmark-org.apache.spark.sql.benchmark.CometPlanOutputBenchmark
```
It runs on `local[1]` unless `COMET_BENCHMARK_MASTER` says otherwise.
<details>
<summary>Full benchmark output</summary>
```text
1 bigint output columns, 1 column read, cacheOutputSchema=false: 1025
batches in 1 tasks, the schema sent with 1025 of them
1 bigint output columns, 1 column read, cacheOutputSchema=true: 1025 batches
in 1 tasks, the schema sent with 1 of them
OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Mac OS X 26.6.2
Apple M3 Max
1 bigint output columns, 1 column read: Best Time(ms) Avg Time(ms)
Stdev(ms) Rate(M/s) Per Row(ns) Relative
------------------------------------------------------------------------------------------------------------------------
cacheOutputSchema OFF 56 61
3 150.7 6.6 1.0X
cacheOutputSchema ON 54 59
3 156.8 6.4 1.0X
cacheOutputSchema OFF (repeat) 55 61
4 151.2 6.6 1.0X
1 bigint output columns, 1 column read: process CPU per batch 79.1 us with
the schema in every batch, 78.5 us with it cached, 0.6 us (0.8%) less
10 bigint output columns, 1 column read, cacheOutputSchema=false: 1025
batches in 1 tasks, the schema sent with 1025 of them
10 bigint output columns, 1 column read, cacheOutputSchema=true: 1025
batches in 1 tasks, the schema sent with 1 of them
OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Mac OS X 26.6.2
Apple M3 Max
10 bigint output columns, 1 column read: Best Time(ms) Avg Time(ms)
Stdev(ms) Rate(M/s) Per Row(ns) Relative
------------------------------------------------------------------------------------------------------------------------
cacheOutputSchema OFF 52 58
4 160.5 6.2 1.0X
cacheOutputSchema ON 51 57
5 164.6 6.1 1.0X
cacheOutputSchema OFF (repeat) 51 57
5 164.9 6.1 1.0X
10 bigint output columns, 1 column read: process CPU per batch 90.3 us with
the schema in every batch, 85.5 us with it cached, 4.9 us (5.4%) less
100 bigint output columns, 1 column read, cacheOutputSchema=false: 1025
batches in 1 tasks, the schema sent with 1025 of them
100 bigint output columns, 1 column read, cacheOutputSchema=true: 1025
batches in 1 tasks, the schema sent with 1 of them
OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Mac OS X 26.6.2
Apple M3 Max
100 bigint output columns, 1 column read: Best Time(ms) Avg Time(ms)
Stdev(ms) Rate(M/s) Per Row(ns) Relative
------------------------------------------------------------------------------------------------------------------------
cacheOutputSchema OFF 214 218
4 39.3 25.5 1.0X
cacheOutputSchema ON 178 203
18 47.1 21.2 1.2X
cacheOutputSchema OFF (repeat) 232 254
17 36.2 27.6 0.9X
100 bigint output columns, 1 column read: process CPU per batch 314.8 us
with the schema in every batch, 264.7 us with it cached, 50.1 us (15.9%) less
1 string output columns, 1 column read, cacheOutputSchema=false: 1025
batches in 1 tasks, the schema sent with 1025 of them
1 string output columns, 1 column read, cacheOutputSchema=true: 1025 batches
in 1 tasks, the schema sent with 1 of them
OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Mac OS X 26.6.2
Apple M3 Max
1 string output columns, 1 column read: Best Time(ms) Avg Time(ms)
Stdev(ms) Rate(M/s) Per Row(ns) Relative
------------------------------------------------------------------------------------------------------------------------
cacheOutputSchema OFF 142 163
12 59.2 16.9 1.0X
cacheOutputSchema ON 113 150
20 74.6 13.4 1.3X
cacheOutputSchema OFF (repeat) 121 131
7 69.5 14.4 1.2X
1 string output columns, 1 column read: process CPU per batch 158.9 us with
the schema in every batch, 190.2 us with it cached, -31.3 us (-19.7%) less
10 string output columns, 1 column read, cacheOutputSchema=false: 1025
batches in 1 tasks, the schema sent with 1025 of them
10 string output columns, 1 column read, cacheOutputSchema=true: 1025
batches in 1 tasks, the schema sent with 1 of them
OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Mac OS X 26.6.2
Apple M3 Max
10 string output columns, 1 column read: Best Time(ms) Avg Time(ms)
Stdev(ms) Rate(M/s) Per Row(ns) Relative
------------------------------------------------------------------------------------------------------------------------
cacheOutputSchema OFF 132 165
20 63.5 15.7 1.0X
cacheOutputSchema ON 132 156
15 63.5 15.8 1.0X
cacheOutputSchema OFF (repeat) 138 171
16 60.8 16.5 1.0X
10 string output columns, 1 column read: process CPU per batch 210.0 us with
the schema in every batch, 207.0 us with it cached, 3.0 us (1.4%) less
100 string output columns, 1 column read, cacheOutputSchema=false: 1025
batches in 1 tasks, the schema sent with 1025 of them
100 string output columns, 1 column read, cacheOutputSchema=true: 1025
batches in 1 tasks, the schema sent with 1 of them
OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Mac OS X 26.6.2
Apple M3 Max
100 string output columns, 1 column read: Best Time(ms) Avg Time(ms)
Stdev(ms) Rate(M/s) Per Row(ns) Relative
------------------------------------------------------------------------------------------------------------------------
cacheOutputSchema OFF 269 278
9 31.2 32.0 1.0X
cacheOutputSchema ON 211 229
8 39.8 25.1 1.3X
cacheOutputSchema OFF (repeat) 271 279
9 31.0 32.3 1.0X
100 string output columns, 1 column read: process CPU per batch 395.1 us
with the schema in every batch, 340.5 us with it cached, 54.6 us (13.8%) less
1 struct<bigint,string> output columns, 1 column read,
cacheOutputSchema=false: 1025 batches in 1 tasks, the schema sent with 1025 of
them
1 struct<bigint,string> output columns, 1 column read,
cacheOutputSchema=true: 1025 batches in 1 tasks, the schema sent with 1 of them
OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Mac OS X 26.6.2
Apple M3 Max
1 struct<bigint,string> output columns, 1 column read: Best Time(ms) Avg
Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
-------------------------------------------------------------------------------------------------------------------------------------
cacheOutputSchema OFF 155
202 26 53.9 18.5 1.0X
cacheOutputSchema ON 164
188 17 51.1 19.6 0.9X
cacheOutputSchema OFF (repeat) 154
188 24 54.4 18.4 1.0X
1 struct<bigint,string> output columns, 1 column read: process CPU per batch
235.7 us with the schema in every batch, 232.2 us with it cached, 3.4 us (1.5%)
less
10 struct<bigint,string> output columns, 1 column read,
cacheOutputSchema=false: 1025 batches in 1 tasks, the schema sent with 1025 of
them
10 struct<bigint,string> output columns, 1 column read,
cacheOutputSchema=true: 1025 batches in 1 tasks, the schema sent with 1 of them
OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Mac OS X 26.6.2
Apple M3 Max
10 struct<bigint,string> output columns, 1 column read: Best Time(ms) Avg
Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
--------------------------------------------------------------------------------------------------------------------------------------
cacheOutputSchema OFF 177
215 26 47.4 21.1 1.0X
cacheOutputSchema ON 152
212 25 55.2 18.1 1.2X
cacheOutputSchema OFF (repeat) 186
211 18 45.1 22.2 1.0X
10 struct<bigint,string> output columns, 1 column read: process CPU per
batch 274.6 us with the schema in every batch, 265.8 us with it cached, 8.8 us
(3.2%) less
100 struct<bigint,string> output columns, 1 column read,
cacheOutputSchema=false: 1025 batches in 1 tasks, the schema sent with 1025 of
them
100 struct<bigint,string> output columns, 1 column read,
cacheOutputSchema=true: 1025 batches in 1 tasks, the schema sent with 1 of them
OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Mac OS X 26.6.2
Apple M3 Max
100 struct<bigint,string> output columns, 1 column read: Best Time(ms)
Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
---------------------------------------------------------------------------------------------------------------------------------------
cacheOutputSchema OFF 627
651 23 13.4 74.7 1.0X
cacheOutputSchema ON 514
529 19 16.3 61.3 1.2X
cacheOutputSchema OFF (repeat) 622
636 18 13.5 74.1 1.0X
100 struct<bigint,string> output columns, 1 column read: process CPU per
batch 801.7 us with the schema in every batch, 679.5 us with it cached, 122.2
us (15.2%) less
100 bigint output columns, 100 columns read, cacheOutputSchema=false: 129
batches in 3 tasks, the schema sent with 129 of them
100 bigint output columns, 100 columns read, cacheOutputSchema=true: 129
batches in 3 tasks, the schema sent with 3 of them
OpenJDK 64-Bit Server VM 17.0.10+7-LTS on Mac OS X 26.6.2
Apple M3 Max
100 bigint output columns, 100 columns read: Best Time(ms) Avg Time(ms)
Stdev(ms) Rate(M/s) Per Row(ns) Relative
---------------------------------------------------------------------------------------------------------------------------
cacheOutputSchema OFF 479 488
7 2.2 456.9 1.0X
cacheOutputSchema ON 482 488
6 2.2 459.7 1.0X
cacheOutputSchema OFF (repeat) 477 488
7 2.2 455.3 1.0X
100 bigint output columns, 100 columns read: process CPU per batch 4354.2 us
with the schema in every batch, 4320.5 us with it cached, 33.6 us (0.8%) less
```
</details>
--
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]
