Repository navigation
Conversation
andygrove
added
enhancement
executePlan exported an ArrowSchema for every column of every output batch, and the JVM imported each one, although a plan's output schema rarely changes between batches. Native now records the data type and field metadata of each output column it last exported, and while a batch matches them it fills only the ArrowArray structs. It sets bit 32 of the row count it returns when it filled the schemas too. The JVM passes the same empty ArrowSchema structs with every batch, and imports a batch that arrives without a schema with the fields of the last one. A batch with a dictionary-encoded column always carries its schema, because the JVM gives each dictionary it imports a new id. The internal config spark.comet.exec.cacheOutputSchema.enabled, on by default, makes native send the schema with every batch when it is off.
CometPlanOutputBenchmark compares spark.comet.exec.cacheOutputSchema.enabled on and off for native plans with 1, 10 and 100 output columns of bigint, string and struct<bigint,string>, and for one that reads 100 distinct columns. Besides wall time, which leaves planning out, it reports each arm's process CPU time per batch, since the JVM's import overlaps the native plan when that runs on a Tokio worker.
andygrove
force-pushed
the
cache-native-output-schema
branch
from
November 2, 2026 18:30
5fcdd4c to
566fc87
Compare
Member
Author
|
This matches the M3 Max results in the description. With 100 projected columns, wall time is 1.36x faster for Full results file |
andygrove
marked this pull request as ready for review
November 2, 2026 18:58
comphead
reviewed
Nov 3, 2026
comphead
left a comment
comphead
left a comment
Contributor
There was a problem hiding this comment.
Thanks, this reads cleanly. I checked it against arrow-rs 59.3 and Arrow Java 18.3. The cache compares exactly what move_to_spark builds the FFI_ArrowSchema from, and arrow-rs Field equality includes child names, nullability and metadata, so a nested change sends the schema again. Always sending it for dictionaries matches ArrayImporter.doImport, which clears the provider's dictionary vector on every import. I found no leak or double release on the error or end-of-output paths. I have not run the tests or the benchmark.
#6883 makes prepare_output pub(crate) and calls it from its new native/core/src/local.rs, and its LocalQueryIterator reads through getNextBatch. Whichever lands second has to pass None there, since getNextBatch does not handle OUTPUT_SCHEMA_EXPORTED.
Comment on lines
+473
to
+474
| private def hasDictionary(field: Field): Boolean = | ||
| field.getDictionary != null || field.getChildren.asScala.exists(hasDictionary) |
Contributor
There was a problem hiding this comment.
CometArrowPythonRunnerBase (under spark-4.x) has the same private hasDictionary(field: Field). Could this one be private[apache], or move to CometVectorUtils, so the Python runner reuses it instead of keeping a second copy?
Comment on lines
+341
to
+344
| val batches = Seq( | ||
| nextBatch(Seq(1, 2), withSchema = true), | ||
| nextBatch(Seq(3), withSchema = false), | ||
| nextBatch(Seq(4, 5, 6), withSchema = false)) |
Contributor
There was a problem hiding this comment.
Could this sequence also send a batch of another type with its schema (a BigIntVector, for example) and then one more without a schema, and check that the last one reads back as longs? With the current tests, a getNextPlanBatch that kept only the first schema's fields would still pass: the schema arrives once here and in the dictionary test, and with the cache off the fields it keeps are never used. The native side of that sequence is covered by output_schema_cache_writes_the_schema_only_when_it_changes, but nothing on the JVM side reads a batch without a schema after a second schema.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
executePlanhands each output batch to the JVM as oneArrowArrayand oneArrowSchemaper 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 aField, 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 to 1.4 µs perstruct<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
ArrowArrayStreamReaderloads every batch throughData.importIntoVectorSchemaRoot, which builds a temporaryStructVectorand 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 intoIOExceptiontext, the problem #6243 worked around for input, and it would reuse one root for every batch, which changes who owns the batchesCometExecIteratorreturns. KeepingexecutePlanand dropping only the repeated schema leaves errors, metrics and batch ownership as they are.What changes are included in this PR?
jni_api.rs): anOutputSchemaCacheheld by theExecutionContextrecords each output column's data type and field metadata asexecutePlanlast exported them. These two are whatmove_to_sparkbuilds theFFI_ArrowSchemafrom. While a batch matches them,prepare_outputwrites only the arrays, with the newmove_array_to_spark, which still applieszero_offsets, and doesn't read the schema addresses. Otherwise it writes the schemas too, records the new ones, and setsOUTPUT_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.CometExecIteratorfetches batches throughNativeUtil.getNextPlanBatch, which allocates theArrowSchemastructs once and passes the same, still empty, structs with every batch. When native sets the flag, it imports the batch asgetNextBatchdoes, which consumes the structs, and keeps the batch's fields unless one of them has a dictionary. Otherwise it imports the arrays with those fields through the newArrowImporter.importVector(array, field, provider), which the schema overload now delegates to. Both imports share one loop and its cleanup.close()frees the schema structs that are left. Decoded shuffle blocks still carry a schema each, throughgetNextBatch.spark.comet.exec.cacheOutputSchema.enabled, on by default, makes native send the schema with every batch when it is off. Only native reads it; the JVM follows the flag. It is there for the benchmark and as a fallback.ffi.mdandmemory_management.mddescribe the protocol.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 with every batch with the cache off and once per task with it on. Wall time leaves planning out. The benchmark also 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
NativeUtilSuitetests check that:close();New Rust tests cover the cache's sequence of exports (the first batch writes the schema, a matching one doesn't, a new type writes it again, a dictionary writes it every time, and no addresses write nothing), a disabled cache writing the schema every time,
move_array_to_sparkread back through a schema exported for an earlier array, the column comparison (type and metadata count, name and nullability don't), andhas_dictionaryon nested types.Locally,
NativeUtilSuite,CometExecIteratorLifecycleSuite,CometNativeSuite,SparkErrorConverterSuite,CometArrowStreamSuite,CometNativeShuffleSuiteandCometExecSuitepass (367 tests), and so doescargo clippy --all-targets --workspace -- -D warnings.CometPlanOutputBenchmarkon 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.bigint, 1 column readstring, 1 column readstruct<bigint,string>, 1 column readbigint, 100 columns readEarlier runs gave 1.20x to 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 14 µs on batches that cost 55 to 245 µs of CPU. In the last row the saving, 29 µs per batch, is under 1% of decoding 100 columns.
To run it:
It runs on
local[1]unlessCOMET_BENCHMARK_MASTERsays otherwise.Full benchmark output