andygrove opened a new issue, #7613: URL: https://github.com/apache/datafusion-comet/issues/7613
### What is the problem the feature request solves? Rust 1.100 stabilized the core of the Allocator API: the `Allocator` trait, `Global`/`System`, and allocator-aware `Vec` and `Box` (`new_in`, `with_capacity_in`, `allocator()`, raw-parts variants). See the [release announcement](https://blog.rust-lang.org/2026/11/12/Rust-1.100.0/). Comet charges memory-pool reservations for buffers it owns by estimating sizes manually, and the tests observe allocations by installing a crate-wide `#[global_allocator]`. A custom `Allocator` could make that accounting exact and let buffer growth fail gracefully under memory pressure. This issue is to research whether that is worth doing. ### Describe the potential solution Ideas to evaluate, in priority order: 1. **Pool-charging allocator for Comet-owned byte buffers.** An `Allocator` that grows and shrinks a `MemoryReservation` (or calls the `MemoryPool` directly) on allocate and deallocate. Use it for the shuffle path's own `Vec<u8>`/`Box<[u8]>` buffers, e.g. `shuffle/src/ipc.rs`, the scratch buffers in `writers/buf_batch_writer.rs`, and the `Cursor<Vec<u8>>` in `writers/rss/rss_partition_writer.rs`. - Accounting would be exact instead of estimated (see the note near `shuffle_writer.rs:1082`). - `Vec::try_reserve` works with any allocator, so pool exhaustion could become an `Err` that triggers a spill instead of over-allocating. - Open questions: `Allocator` methods take `&self` but `MemoryReservation::try_grow` needs `&mut self`. Infallible calls (`push`, `extend`) panic on `AllocError`, so hot paths would need `try_reserve`. 2. **Scoped counting allocators in tests.** Replace the crate-global `ObservedAllocator` in `shuffle/src/writers/rss/mod.rs` and the `SERIAL` lock in `core/src/alloc_accounting.rs` with per-test allocators. This only works once the code under test is generic over the allocator. 3. **Keep `AccountingAllocator`** (`core/src/alloc_accounting.rs`) as a `GlobalAlloc`. It provides the process-wide view, which the new API does not replace. ### Additional context - **Arrow buffers are out of scope.** arrow-rs `MutableBuffer` allocates directly and does not take an `Allocator`, so most of Comet's memory is not reachable through this API. Allocator support in arrow-rs would be a separate upstream discussion. - **`String` and `HashMap` are not allocator-aware in std yet.** - **MSRV:** `native/Cargo.toml` sets `rust-version = "1.95.0"`. Using the std API means raising it to 1.100, in line with DataFusion's MSRV policy. - **Prototype on stable now:** the [`allocator-api2`](https://crates.io/crates/allocator-api2) crate provides the same API on current stable Rust, so the research can start before any MSRV bump, then switch to `std::alloc::Allocator` later. Suggested first step: prototype (1) on the shuffle write path with `allocator-api2`. Measure how closely reservations track actual buffer bytes, and whether `try_reserve`-driven spilling behaves correctly under pressure. -- 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]
