Reduce the number of device allocations in the ORC writer - #23433
Reduce the number of device allocations in the ORC writer #23433vuule wants to merge 18 commits into
Conversation
`encode_columns` and `gather_stripes` previously allocated one `device_uvector<uint8_t>` per (stripe, stream) pair, producing ~1,000 driver allocations each per write_orc call. Replace both patterns with a single arena per phase plus non-owning `device_span<uint8_t>` views, holding the encoded and gathered arenas on the `encoded_data` struct for lifetime management. Each per-(stripe, stream) region inside an arena starts on a 256-byte boundary (`rmm::CUDA_ALLOCATION_ALIGNMENT`) to match the alignment that RMM previously gave each independent allocation: the in-loop alignment fix-up only honors `compress_required_chunk_alignment(compression)`, which is 1 for `compression == NONE` and as low as 8 for nvcomp paths, so the encoder/codec kernels rely on the underlying allocation being naturally aligned. Effect on `orc_write_encode` (INTEGRAL_SIGNED, 512 MB): * `cudaMallocFromPoolAsync` per iter: ~2,135 -> ~88 (-96%). * pool GPU time: 53.95 ms -> 48.23 ms (-10.6%). * async GPU time: 104.65 ms -> 48.52 ms (-53.6%). * async-vs-pool delta: +94% -> +0.6% (within noise). * peak memory: +1 MiB on a 1.97 GiB run (alignment padding). * encoded_file_size unchanged. Made-with: Cursor
Splitting the encoder output into one arena rather than a device_uvector per (stripe, stream) kept the upper-bound-sized allocation alive through compression, which raised peak memory by up to 53% on the ORC writer benchmarks. Partition the output into two arenas instead, by whether a region is expected to be compacted into the gather buffer: a region is, when its stripe spans several rowgroups and its size is a strict upper bound rather than exact. gather_stripes compacts every region in the transient arena as well as any region it measures a gap in, so the transient arena can be released as soon as gathering completes. Predicting which regions carry slack only decides arena placement; whether a region is compacted is still measured, so a misprediction cannot leave non-contiguous chunks to be read as a contiguous stream. It only costs the allocation that would have been released. Peak memory is now within 0.08% of main on all 109 benchmark configs, and timings are unchanged against the previous arena layout under the pool resource.
Throughput vs
|
| benchmark | configs | vs main |
|---|---|---|
orc_write_encode INTEGRAL_SIGNED |
4 | -2.6% to -7.9% |
orc_write_encode FLOAT |
4 | -2.9% to -12.9% |
orc_write_encode DECIMAL |
4 | -4.3% to -14.2% |
orc_write_encode TIMESTAMP |
4 | -4.9% to -14.6% |
orc_write_encode STRING |
4 | -9.1% to -14.2% |
orc_write_encode LIST |
4 | -0.7% to -1.4% |
orc_write_encode STRUCT |
4 | -22.4% to -34.8% |
orc_write |
2 | -1.1% to -1.4% |
orc_chunked_write |
4 | -0.2% to -0.3% |
The saving is roughly fixed per write (3-5 ms on the encode configurations), so it reads large where the write is short and small where other work dominates. STRUCT leads because it has the most streams, and so the most allocations removed.
Per-configuration numbers
| benchmark | axes | main |
this branch | change |
|---|---|---|---|---|
| orc_write_encode | data_type=INTEGRAL_SIGNED cardinality=0 run_length=1 | 49.2 ms | 46.3 ms | -5.8% |
| orc_write_encode | data_type=INTEGRAL_SIGNED cardinality=0 run_length=32 | 40.2 ms | 37.0 ms | -7.9% |
| orc_write_encode | data_type=INTEGRAL_SIGNED cardinality=1000 run_length=1 | 120.6 ms | 117.5 ms | -2.6% |
| orc_write_encode | data_type=INTEGRAL_SIGNED cardinality=1000 run_length=32 | 42.9 ms | 39.9 ms | -6.9% |
| orc_write_encode | data_type=FLOAT cardinality=0 run_length=1 | 37.4 ms | 34.2 ms | -8.6% |
| orc_write_encode | data_type=FLOAT cardinality=0 run_length=32 | 25.2 ms | 22.0 ms | -12.9% |
| orc_write_encode | data_type=FLOAT cardinality=1000 run_length=1 | 122.6 ms | 119.0 ms | -2.9% |
| orc_write_encode | data_type=FLOAT cardinality=1000 run_length=32 | 25.4 ms | 22.2 ms | -12.6% |
| orc_write_encode | data_type=DECIMAL cardinality=0 run_length=1 | 62.5 ms | 57.3 ms | -8.3% |
| orc_write_encode | data_type=DECIMAL cardinality=0 run_length=32 | 39.3 ms | 33.7 ms | -14.2% |
| orc_write_encode | data_type=DECIMAL cardinality=1000 run_length=1 | 132.4 ms | 126.6 ms | -4.3% |
| orc_write_encode | data_type=DECIMAL cardinality=1000 run_length=32 | 39.4 ms | 33.8 ms | -14.2% |
| orc_write_encode | data_type=TIMESTAMP cardinality=0 run_length=1 | 35.0 ms | 29.9 ms | -14.6% |
| orc_write_encode | data_type=TIMESTAMP cardinality=0 run_length=32 | 34.1 ms | 29.4 ms | -14.0% |
| orc_write_encode | data_type=TIMESTAMP cardinality=1000 run_length=1 | 109.3 ms | 103.9 ms | -4.9% |
| orc_write_encode | data_type=TIMESTAMP cardinality=1000 run_length=32 | 35.2 ms | 30.5 ms | -13.5% |
| orc_write_encode | data_type=STRING cardinality=0 run_length=1 | 45.5 ms | 41.3 ms | -9.1% |
| orc_write_encode | data_type=STRING cardinality=0 run_length=32 | 45.5 ms | 41.3 ms | -9.1% |
| orc_write_encode | data_type=STRING cardinality=1000 run_length=1 | 38.6 ms | 33.1 ms | -14.2% |
| orc_write_encode | data_type=STRING cardinality=1000 run_length=32 | 38.7 ms | 33.2 ms | -14.2% |
| orc_write_encode | data_type=LIST cardinality=0 run_length=1 | 72.6 ms | 71.6 ms | -1.4% |
| orc_write_encode | data_type=LIST cardinality=0 run_length=32 | 98.4 ms | 97.7 ms | -0.8% |
| orc_write_encode | data_type=LIST cardinality=1000 run_length=1 | 158.4 ms | 157.3 ms | -0.7% |
| orc_write_encode | data_type=LIST cardinality=1000 run_length=32 | 118.3 ms | 117.5 ms | -0.7% |
| orc_write_encode | data_type=STRUCT cardinality=0 run_length=1 | 67.4 ms | 47.5 ms | -29.4% |
| orc_write_encode | data_type=STRUCT cardinality=0 run_length=32 | 69.5 ms | 46.7 ms | -32.8% |
| orc_write_encode | data_type=STRUCT cardinality=1000 run_length=1 | 111.9 ms | 86.9 ms | -22.4% |
| orc_write_encode | data_type=STRUCT cardinality=1000 run_length=32 | 75.9 ms | 49.5 ms | -34.8% |
| orc_write | num_columns=64 | 509.3 ms | 502.2 ms | -1.4% |
| orc_chunked_write | num_columns=64 num_chunks=64 | 7164.7 ms | 7145.8 ms | -0.3% |
| orc_chunked_write | num_columns=64 num_chunks=8 | 3526.0 ms | 3519.6 ms | -0.2% |
| orc_write | num_columns=8 | 93.5 ms | 92.5 ms | -1.1% |
| orc_chunked_write | num_columns=8 num_chunks=64 | 2037.7 ms | 2033.6 ms | -0.2% |
| orc_chunked_write | num_columns=8 num_chunks=8 | 331.1 ms | 330.3 ms | -0.2% |
Faster on 34 of 34 configs; best -34.8% (orc_write_encode data_type=STRUCT cardinality=1000 run_length=32).
Throughput vs
|
| benchmark | configs | vs main |
|---|---|---|
orc_write_encode INTEGRAL_SIGNED |
4 | -24.0% to -52.8% |
orc_write_encode FLOAT |
4 | -19.1% to -56.0% |
orc_write_encode DECIMAL |
4 | -21.0% to -57.4% |
orc_write_encode TIMESTAMP |
4 | -34.8% to -70.2% |
orc_write_encode STRING |
4 | -38.7% to -69.7% |
orc_write_encode LIST |
4 | -8.4% to -17.1% |
orc_write_encode STRUCT |
4 | -67.2% to -88.7% |
orc_write |
2 | -13.1% to -14.6% |
orc_chunked_write |
4 | -0.8% to -3.9% |
Releasing transient_buffer before compression is what costs 5-28% here versus keeping one arena. Isolated with a third build in which only the release is disabled: splitting into two arenas is itself free (-0.2% to -2.6%, slightly faster), and the whole cost is the single cudaFreeAsync. It tracks the size of the buffer returned, not the number of calls, so it appears exactly where the peak-memory benefit is largest (INTEGRAL_SIGNED run_length=1: +6.3 ms; the same type at run_length=32, where RLE shrinks the arena: +0.16 ms). Skipping the release for small arenas would therefore buy nothing.
Per-configuration numbers
| benchmark | axes | main |
this branch | change |
|---|---|---|---|---|
| orc_write_encode | data_type=INTEGRAL_SIGNED cardinality=0 run_length=1 | 97.0 ms | 52.0 ms | -46.4% |
| orc_write_encode | data_type=INTEGRAL_SIGNED cardinality=0 run_length=32 | 77.5 ms | 36.6 ms | -52.8% |
| orc_write_encode | data_type=INTEGRAL_SIGNED cardinality=1000 run_length=1 | 162.2 ms | 123.3 ms | -24.0% |
| orc_write_encode | data_type=INTEGRAL_SIGNED cardinality=1000 run_length=32 | 76.1 ms | 39.5 ms | -48.1% |
| orc_write_encode | data_type=FLOAT cardinality=0 run_length=1 | 86.7 ms | 38.2 ms | -56.0% |
| orc_write_encode | data_type=FLOAT cardinality=0 run_length=32 | 53.8 ms | 25.0 ms | -53.5% |
| orc_write_encode | data_type=FLOAT cardinality=1000 run_length=1 | 155.4 ms | 125.7 ms | -19.1% |
| orc_write_encode | data_type=FLOAT cardinality=1000 run_length=32 | 52.6 ms | 28.0 ms | -46.7% |
| orc_write_encode | data_type=DECIMAL cardinality=0 run_length=1 | 120.3 ms | 59.5 ms | -50.5% |
| orc_write_encode | data_type=DECIMAL cardinality=0 run_length=32 | 78.8 ms | 33.5 ms | -57.4% |
| orc_write_encode | data_type=DECIMAL cardinality=1000 run_length=1 | 160.4 ms | 126.7 ms | -21.0% |
| orc_write_encode | data_type=DECIMAL cardinality=1000 run_length=32 | 58.3 ms | 33.6 ms | -42.3% |
| orc_write_encode | data_type=TIMESTAMP cardinality=0 run_length=1 | 99.8 ms | 29.7 ms | -70.2% |
| orc_write_encode | data_type=TIMESTAMP cardinality=0 run_length=32 | 75.1 ms | 29.2 ms | -61.1% |
| orc_write_encode | data_type=TIMESTAMP cardinality=1000 run_length=1 | 159.6 ms | 104.1 ms | -34.8% |
| orc_write_encode | data_type=TIMESTAMP cardinality=1000 run_length=32 | 75.9 ms | 30.3 ms | -60.1% |
| orc_write_encode | data_type=STRING cardinality=0 run_length=1 | 92.0 ms | 41.3 ms | -55.1% |
| orc_write_encode | data_type=STRING cardinality=0 run_length=32 | 136.7 ms | 41.5 ms | -69.7% |
| orc_write_encode | data_type=STRING cardinality=1000 run_length=1 | 94.1 ms | 38.0 ms | -59.6% |
| orc_write_encode | data_type=STRING cardinality=1000 run_length=32 | 59.2 ms | 36.3 ms | -38.7% |
| orc_write_encode | data_type=LIST cardinality=0 run_length=1 | 87.1 ms | 72.2 ms | -17.1% |
| orc_write_encode | data_type=LIST cardinality=0 run_length=32 | 113.7 ms | 99.1 ms | -12.9% |
| orc_write_encode | data_type=LIST cardinality=1000 run_length=1 | 173.5 ms | 158.8 ms | -8.4% |
| orc_write_encode | data_type=LIST cardinality=1000 run_length=32 | 133.6 ms | 119.1 ms | -10.9% |
| orc_write_encode | data_type=STRUCT cardinality=0 run_length=1 | 472.4 ms | 53.7 ms | -88.6% |
| orc_write_encode | data_type=STRUCT cardinality=0 run_length=32 | 432.8 ms | 48.8 ms | -88.7% |
| orc_write_encode | data_type=STRUCT cardinality=1000 run_length=1 | 273.9 ms | 89.8 ms | -67.2% |
| orc_write_encode | data_type=STRUCT cardinality=1000 run_length=32 | 261.0 ms | 52.5 ms | -79.9% |
| orc_write | num_columns=64 | 590.9 ms | 504.6 ms | -14.6% |
| orc_chunked_write | num_columns=64 num_chunks=64 | 7452.8 ms | 7163.3 ms | -3.9% |
| orc_chunked_write | num_columns=64 num_chunks=8 | 3554.1 ms | 3521.8 ms | -0.9% |
| orc_write | num_columns=8 | 105.8 ms | 92.0 ms | -13.1% |
| orc_chunked_write | num_columns=8 num_chunks=64 | 2077.8 ms | 2050.2 ms | -1.3% |
| orc_chunked_write | num_columns=8 num_chunks=8 | 334.6 ms | 331.9 ms | -0.8% |
Faster on 34 of 34 configs; best -88.7% (orc_write_encode data_type=STRUCT cardinality=0 run_length=32).
Peak memory vs
|
| benchmark | configs | mean change | worst change | worst delta |
|---|---|---|---|---|
| orc_write_encode | 28 | +0.02% | +0.07% | +0.8 MiB |
| orc_write_io_compression | 60 | +0.03% | +0.08% | +0.5 MiB |
| orc_write_statistics | 15 | +0.03% | +0.06% | +0.4 MiB |
| orc_write | 2 | +0.02% | +0.03% | +0.5 MiB |
| orc_chunked_write | 4 | +0.03% | +0.06% | +0.0 MiB |
109 configs total; largest increase +0.08%, largest decrease +0.00%.
Peak memory is unchanged. This is the metric that motivated the two-arena split described in the PR body: with a single arena holding the encoded output through compression, orc_write_encode regressed by a mean of +21.9% and a worst case of +53.0% (+854.8 MiB), and the worst configurations overall were the wide chunked writes at +46.8% (+506.4 MiB).
Treat the pass-through float/double data stream as an exact size estimate, so those regions stay in the arena that is read in place instead of being copied out for no reason. Hand the transient-region set from `encode_columns` to `gather_stripes` as its own argument rather than parking it in `encoded_data`, where it outlived its only use, and make the release of `transient_buffer` unconditional with the invariant it relies on checked explicitly. Flatten the per-(stripe, stream) bookkeeping into single allocations, and cover stripes that span several rowgroups with compression disabled and with pass-through floats.
The alignment is not there because the encoder needs naturally aligned streams; `encode_columns` already realigns every per-rowgroup pointer by absolute address, so the encoded arenas would tolerate any base. It is the gather that needs it: compaction drops that per-rowgroup alignment, leaving the region base as the pointer the compressor receives, so it has to meet the codec's requirement. Check that rather than assuming 256 bytes is always enough.
The string column in the test is dictionary encoded, whose per-stripe stream is a single region no matter how many rowgroups it spans. Add the types whose per-rowgroup sizes are known exactly and chained end to end -- direct encoded strings and decimals -- since those are what actually exercise reading a multi-rowgroup stripe in place.
The decimal parameterization already covers stripes of one and of several rowgroups, and decimal data streams are the exactly sized kind, so adding a compression axis to it covers writing a multi-rowgroup stripe straight from the encoder output. Drops the standalone test that was doing the same thing.
OrcChunkedReaderInputLimitTest already writes null-free doubles with compression disabled at 20k-row stripes, which is the same multi-rowgroup pass-through layout at a much larger scale.
"Region" was a second name for something ORC already names. An extent is the aligned byte range within an arena that holds one (stripe, stream) pair, which is what replaces the allocation each pair used to get.
vuule
changed the title
vuule
changed the title
vuule
added
Performance
The bespoke class was a 2D array of bools with a hand-rolled index; cudf already has host_2dspan for passing [row][column] host data between the writer's helpers.
|
/ok to test 0b47b6b |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
📝 WalkthroughSummary by CodeRabbit
WalkthroughChangesORC extent management
Estimated code review effort: 4 (Complex) | ~45 minutes Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (5)
cpp/tests/io/orc_test.cpp (1)
1287-1292: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftAdd empty-input and sliced-column coverage.
The matrix covers nulls, boundary sizes, and multi-block sizes, but all cases use non-empty, unsliced columns. Add an empty case and a dedicated sliced-column case.
As per coding guidelines, C++ tests must cover empty inputs, nulls, sliced columns, boundary and multi-block sizes.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/io/orc_test.cpp` around lines 1287 - 1292, Add empty-input and sliced-column scenarios to the parameterization for OrcWriterTestDecimal alongside the existing size, scale, and compression combinations. Ensure the new tests exercise a truly empty column and a column created from a non-zero slice while preserving the current coverage.Source: Coding guidelines
cpp/src/io/orc/writer_impl.hpp (1)
108-111: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the span/arena invariant on the members themselves.
dataholds raw views into the three arenas, so clearing or reassigning any arena silently invalidates the corresponding spans (relied upon ingather_stripesand in the post-compression cleanup). A///<doxygen note per member stating which arena each span may point into, and thatdatamust be repointed or cleared before an arena is released, would make the coupling explicit.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/io/orc/writer_impl.hpp` around lines 108 - 111, Update the member comments for persistent_buffer, transient_buffer, gathered_buffer, and data to document which arena the data spans may reference. Add ///< Doxygen notes stating that clearing or reassigning any referenced arena invalidates those spans and that data must be repointed or cleared before the arena is released.Source: Coding guidelines
cpp/src/io/orc/writer_impl.cu (3)
856-867: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider a doxygen block for the new dual return.
The return contract (
encoded_dataplus a flattened[stripe][strm_id]transient flag array indexed withstreams.size()stride) is non-obvious and is consumed bygather_stripesas ahost_2dspan. A@brief/@param/@returnblock would match the style used forgather_stripesdirectly below.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/io/orc/writer_impl.cu` around lines 856 - 867, Add a Doxygen documentation block immediately before encode_columns describing its inputs, the paired encoded_data and flattened transient-buffer flag array return value, and the array’s streams.size()-based [stripe][strm_id] indexing consumed by gather_stripes. Match the existing gather_stripes documentation style.Source: Coding guidelines
2582-2586: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueClear
databefore releasing the arenas it points into.Purely defensive — nothing dereferences the spans in between — but inverting the order removes the transient window where
enc_data.dataholds views into freed memory.♻️ Reorder
+ enc_data.data.clear(); enc_data.persistent_buffer = rmm::device_uvector<uint8_t>{0, stream}; enc_data.gathered_buffer = rmm::device_uvector<uint8_t>{0, stream}; - enc_data.data.clear();🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/io/orc/writer_impl.cu` around lines 2582 - 2586, In the encoded-data cleanup block, update the deallocation order so enc_data.data.clear() runs before resetting enc_data.persistent_buffer and enc_data.gathered_buffer. Keep the existing cleanup operations unchanged otherwise, ensuring the spans are cleared before their backing arenas are released.
1119-1124: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider designated initializers here.
Positional aggregate init of five members (three of which are same-typed
device_uvector<uint8_t>) will silently mis-bind ifencoded_data's member order ever changes. C++20 designated initializers (.persistent_buffer = ...) make it order-safe and self-documenting.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/io/orc/writer_impl.cu` around lines 1119 - 1124, Update the encoded_data aggregate construction in the surrounding return statement to use C++20 designated initializers for all five members, including the three device_uvector<uint8_t> values, while preserving their current values and move semantics.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cpp/src/io/orc/writer_impl.cu`:
- Around line 1024-1056: Make the alignment invariant explicit in the validation
near the existing extent_alignment and uncomp_block_align checks: require that
extent_alignment is divisible by uncomp_block_align, not merely greater than or
equal to it, and document that the RMM arena base is aligned to
extent_alignment. Preserve the single-rowgroup optimization while ensuring its
skipped alignment shift remains valid.
---
Nitpick comments:
In `@cpp/src/io/orc/writer_impl.cu`:
- Around line 856-867: Add a Doxygen documentation block immediately before
encode_columns describing its inputs, the paired encoded_data and flattened
transient-buffer flag array return value, and the array’s streams.size()-based
[stripe][strm_id] indexing consumed by gather_stripes. Match the existing
gather_stripes documentation style.
- Around line 2582-2586: In the encoded-data cleanup block, update the
deallocation order so enc_data.data.clear() runs before resetting
enc_data.persistent_buffer and enc_data.gathered_buffer. Keep the existing
cleanup operations unchanged otherwise, ensuring the spans are cleared before
their backing arenas are released.
- Around line 1119-1124: Update the encoded_data aggregate construction in the
surrounding return statement to use C++20 designated initializers for all five
members, including the three device_uvector<uint8_t> values, while preserving
their current values and move semantics.
In `@cpp/src/io/orc/writer_impl.hpp`:
- Around line 108-111: Update the member comments for persistent_buffer,
transient_buffer, gathered_buffer, and data to document which arena the data
spans may reference. Add ///< Doxygen notes stating that clearing or reassigning
any referenced arena invalidates those spans and that data must be repointed or
cleared before the arena is released.
In `@cpp/tests/io/orc_test.cpp`:
- Around line 1287-1292: Add empty-input and sliced-column scenarios to the
parameterization for OrcWriterTestDecimal alongside the existing size, scale,
and compression combinations. Ensure the new tests exercise a truly empty column
and a column created from a non-zero slice while preserving the current
coverage.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 73a4f3e3-0bc5-43f5-a7ea-93eb44812e6f
📒 Files selected for processing (3)
cpp/src/io/orc/writer_impl.cucpp/src/io/orc/writer_impl.hppcpp/tests/io/orc_test.cpp
bdice
left a comment
bdice
left a comment
There was a problem hiding this comment.
I will have to read the PR more closely, but I wanted to check a couple of things. A different way to reduce the number of (small) allocations from the current memory resource is to create and suballocate from a temporary pool. Could we create a temporary RMM pool that uses the current device resource as its base allocator, and use that pool as the allocator for internal temporary allocations? This doesn’t need to return the final allocations, as far as I can tell. We should avoid doing extensive padding math and manual alignment if we can just suballocate from a pool object that handles that for us.
Using a temporary pool is doable, but I would argue that using arenas is the right approach, since we know in advance how much memory we'll need. Using the arenas should be faster and have lower peak memory than a pool. Also, the Parquet writer already does something very similar. The additional complexity here is to have two arenas to reduce the peak memory use, which is justified IMO. |
Extent bases are offset from the arena base by multiples of extent_alignment, so the codec's chunk alignment has to divide it. The previous >= check only implied that because every alignment in play today is a power of two.
|
/ok to test a86b485 |
2 participants
Description
Contributes to #23464
encode_columnsallocates onedevice_uvectorper (stripe, stream) pair, so a wide table with many stripes produces thousands of device allocations.With this PR, we compute the size of every stream up front and suballocate an extent (an aligned byte range within an arena) for each. Each extent is aligned exactly as a fresh RMM allocation would be.
Extent sizes are upper bounds, because the encoded length of an RLE stream is not known until it has been encoded.
gather_stripescompacts a stripe's per-rowgroup chunks into a contiguous, tightly sized stream. The encoder output is therefore split across two arenas, by whether an extent might be compacted:transient_buffer- extents whose stripe spans several rowgroups and whose size is a strict upper bound rather than exact.gather_stripescompacts every extent in this arena, which lets the whole arena be released as soon as gathering completes, before compression allocates.persistent_buffer- everything else: single-rowgroup stripes, and multi-rowgroup extents whose size estimate is exact. This arena lives until the encoded data is no longer needed.gathered_buffer- the compaction destination.Checklist