[SPARK-58918][SQL] Push column pruning into the Arrow cache read path - #58177
Closed
viirya wants to merge 1 commit into
Closed
[SPARK-58918][SQL] Push column pruning into the Arrow cache read path#58177viirya wants to merge 1 commit into
viirya wants to merge 1 commit into
Conversation
When reading a projection of an Arrow-cached relation, the reader deserialized every cached column -- decompressing and loading all of them off-heap -- and then discarded the unselected ones. This wastes work proportional to the pruned columns, which dominates wide-relation scans, especially with compression. Read the encapsulated IPC RecordBatch message's metadata (a small flatbuffer that lists every buffer's offset and length within the body) and copy only the byte ranges belonging to the selected columns straight out of the in-memory cached bytes, so the unselected columns are never copied off-heap, loaded, or decompressed. The selected buffers become windows into a single off-heap allocation, matching how the standard IPC reader slices one body buffer, so ownership stays a single allocation freed once. Both the columnar and row read paths use it; a projection whose selected attribute is absent from the cache schema falls back to the full read. Co-authored-by: Claude Code
uros-b
approved these changes
Aug 20, 2026
Member
|
LGTM |
viirya
added a commit
that referenced
this pull request
Aug 21, 2026
### What changes were proposed in this pull request? When reading a projection of an Arrow-cached relation, `ArrowCachedBatchSerializer`'s read path deserialized **every** cached column -- decompressing and loading all of them off-heap -- and then discarded the unselected ones. This change reads only the selected columns out of the cached bytes. A new helper `ArrowCachedBatchSerializer.readProjectedRecordBatch` reads the encapsulated IPC RecordBatch message's metadata (a small flatbuffer that lists every buffer's offset and length within the body) and copies just the byte ranges belonging to the selected columns straight out of the in-memory cached `Array[Byte]`, so the unselected columns are never copied off-heap, loaded, or decompressed. The selected buffers become windows into a single off-heap allocation, mirroring how the standard IPC reader slices one body buffer, so ownership stays a single allocation freed once. It is wired into both read paths -- `convertCachedBatchToColumnarBatch` (columnar) and `convertCachedBatchToInternalRow` (row). A projection whose selected attribute is absent from the cache schema (index `-1`) falls back to reading the full batch. ### Why are the changes needed? The wasted work is proportional to the pruned columns and dominates wide-relation scans, especially with compression. An e2e SQL benchmark (`sum(col0)` over a cached relation, cache materialized outside the timed region, 1M rows x 50 long columns) shows: | | before | after | |---|---|---| | Arrow cache -- uncompressed | 31 ms | 16 ms (~1.9x) | | Arrow cache -- zstd level 1 | 326 ms | 22 ms (~15x) | The compressed case gains most, since the 49 pruned columns are no longer decompressed. ### Does this PR introduce _any_ user-facing change? No. The Arrow cache serializer (SPARK-57268) is unreleased, and this is a read-path performance improvement with identical results. ### How was this patch tested? - A new `ArrowCachedBatchSerializerSuite` test covers projection shapes (reordering, single column at each position, complex-after-var-width, duplicate selection, full projection) on both read paths, and is negative-validated (an off-by-one in the buffer-span arithmetic fails it). - Full `ArrowCachedBatchSerializerSuite` and `ArrowCachedBatchKryoRegistrationSuite` pass (77 tests). - A new `columnPruningWideTable` case in `ArrowCacheBenchmark` measures the read path with the cache built outside the timed region; the committed benchmark result files are regenerated by the benchmark CI job. ### Was this patch authored or co-authored using generative AI tooling? Yes, this pull request and its description were written by Claude Code. Closes #58177 from viirya/arrow-cache-projection-pushdown. Authored-by: Liang-Chi Hsieh <viirya@gmail.com> Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com> (cherry picked from commit ca6a869) Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com>
Member
Author
Member
Author
|
Thanks @uros-b |
viirya
added a commit
that referenced
this pull request
Aug 22, 2026
…rrow cache column-pruning read ### What changes were proposed in this pull request? Add a test to `ArrowCachedBatchSerializerSuite` covering deeply nested types in the Arrow cache column-pruning read path added by #58177 (SPARK-58918): `array<struct<...>>`, `struct<struct<struct<...>>>`, and `map<int, array<int>>`, each selected alone, mixed with primitives, reordered, and with other columns pruned in between, under both the row and vectorized read paths. ### Why are the changes needed? The pruning read copies each selected top-level column's entire buffer subtree, so `readProjectedRecordBatch`'s buffer-span arithmetic (`fieldBufferCount` / `fieldNodeCount`) recurses through nested types to arbitrary depth. The original test only exercised one level of nesting (`array<int>`, `struct` of primitives), so an off-by-one in the recursive span over a deeper subtree -- which would shift the buffers of a following selected column and produce wrong values -- was not covered. The new test fails if that recursion is broken (verified locally by making the buffer count ignore children). ### Does this PR introduce _any_ user-facing change? No, test-only. ### How was this patch tested? The new test; full `ArrowCachedBatchSerializerSuite` passes. ### Was this patch authored or co-authored using generative AI tooling? Yes, this pull request and its description were written by Claude Code. Closes #58215 from viirya/arrow-cache-nested-projection-test. Authored-by: Liang-Chi Hsieh <viirya@gmail.com> Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com>
viirya
added a commit
that referenced
this pull request
Aug 22, 2026
…rrow cache column-pruning read ### What changes were proposed in this pull request? Add a test to `ArrowCachedBatchSerializerSuite` covering deeply nested types in the Arrow cache column-pruning read path added by #58177 (SPARK-58918): `array<struct<...>>`, `struct<struct<struct<...>>>`, and `map<int, array<int>>`, each selected alone, mixed with primitives, reordered, and with other columns pruned in between, under both the row and vectorized read paths. ### Why are the changes needed? The pruning read copies each selected top-level column's entire buffer subtree, so `readProjectedRecordBatch`'s buffer-span arithmetic (`fieldBufferCount` / `fieldNodeCount`) recurses through nested types to arbitrary depth. The original test only exercised one level of nesting (`array<int>`, `struct` of primitives), so an off-by-one in the recursive span over a deeper subtree -- which would shift the buffers of a following selected column and produce wrong values -- was not covered. The new test fails if that recursion is broken (verified locally by making the buffer count ignore children). ### Does this PR introduce _any_ user-facing change? No, test-only. ### How was this patch tested? The new test; full `ArrowCachedBatchSerializerSuite` passes. ### Was this patch authored or co-authored using generative AI tooling? Yes, this pull request and its description were written by Claude Code. Closes #58215 from viirya/arrow-cache-nested-projection-test. Authored-by: Liang-Chi Hsieh <viirya@gmail.com> Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com> (cherry picked from commit 09ebec3) Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com>
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
2 participants
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.
What changes were proposed in this pull request?
When reading a projection of an Arrow-cached relation,
ArrowCachedBatchSerializer's read path deserialized every cached column -- decompressing and loading all of them off-heap -- and then discarded the unselected ones. This change reads only the selected columns out of the cached bytes.A new helper
ArrowCachedBatchSerializer.readProjectedRecordBatchreads the encapsulated IPC RecordBatch message's metadata (a small flatbuffer that lists every buffer's offset and length within the body) and copies just the byte ranges belonging to the selected columns straight out of the in-memory cachedArray[Byte], so the unselected columns are never copied off-heap, loaded, or decompressed. The selected buffers become windows into a single off-heap allocation, mirroring how the standard IPC reader slices one body buffer, so ownership stays a single allocation freed once.It is wired into both read paths --
convertCachedBatchToColumnarBatch(columnar) andconvertCachedBatchToInternalRow(row). A projection whose selected attribute is absent from the cache schema (index-1) falls back to reading the full batch.Why are the changes needed?
The wasted work is proportional to the pruned columns and dominates wide-relation scans, especially with compression. An e2e SQL benchmark (
sum(col0)over a cached relation, cache materialized outside the timed region, 1M rows x 50 long columns) shows:The compressed case gains most, since the 49 pruned columns are no longer decompressed.
Does this PR introduce any user-facing change?
No. The Arrow cache serializer (SPARK-57268) is unreleased, and this is a read-path performance improvement with identical results.
How was this patch tested?
ArrowCachedBatchSerializerSuitetest covers projection shapes (reordering, single column at each position, complex-after-var-width, duplicate selection, full projection) on both read paths, and is negative-validated (an off-by-one in the buffer-span arithmetic fails it).ArrowCachedBatchSerializerSuiteandArrowCachedBatchKryoRegistrationSuitepass (77 tests).columnPruningWideTablecase inArrowCacheBenchmarkmeasures the read path with the cache built outside the timed region; the committed benchmark result files are regenerated by the benchmark CI job.Was this patch authored or co-authored using generative AI tooling?
Yes, this pull request and its description were written by Claude Code.