james-willis opened a new pull request, #1225: URL: https://github.com/apache/sedona-db/pull/1225
## What this PR does `ST_Collect_Agg`'s `Accumulator::size()` (`rust/sedona-functions/src/st_collect_agg.rs`) reported `size_of::<CollectionAccumulator>() + item.capacity()`, omitting the heap tables of `unique_geometry_types: HashSet<GeometryTypeId>` and `unique_dimensions: HashSet<Dimensions>`, which allocate 24 and 76 bytes respectively on the first insert into every non-empty group (measured, hashbrown 0.15.5). Because a UDAF accumulator is instantiated per group by `GroupsAccumulatorAdapter`, a high-cardinality `GROUP BY` under-reports ~100 bytes per group to the memory pool. Observed in practice: SpatialBench Q5 at sf=10 (10M+ groups, 12 GiB container, 7.73 GB fair memory pool) peaked at 9.55 GB anon with **zero spills and no `ResourcesExhausted`** — the accounted total never reached the pool's spillable budget, so whether the container survives is luck rather than the pool's decision. The minimal fix would be to add the missing terms to `size()` (4 LoC: `+ self.unique_geometry_types.allocation_size() + self.unique_dimensions.allocation_size()`). This PR instead removes the allocations that needed accounting: both HashSets are replaced with the existing `GeometryTypeAndDimensionsSet` u32 bitset (`rust/sedona-geometry/src/types.rs`), which strictly supersedes the minimal fix: - `size()` becomes exact with no extra terms; - `size_of::<CollectionAccumulator>()` shrinks, and two heap allocations per group disappear (~170 bytes/group real at typical capacities — multiple GB at sf=10 group counts); - the per-row `convert_to_state` path (used when the skip-partial-aggregation probe engages at high cardinality) sheds two malloc/free pairs per input row. If the swap is contentious, the 4-LoC accounting fix above is the fallback. ## Compatibility - **State wire format unchanged**: `GeometryTypeAndDimensionsSet` already serializes to the same JSON as the HashSet form (kept deliberately compatible in `sedona-geometry`), and is already used by `ST_Analyze_Agg`. Both existing `state()` columns are derived from the single set; tests assert the encoding is byte-identical. - **Semantics preserved**: uses `insert()`, not `insert_or_ignore()` — the latter silently drops `Dimensions::Unknown(_)`, which would let a mixed/unknown-dimension group emit a header claiming XY instead of erroring. A mixed-dimensions test is included. ## Testing - Full `sedona-functions` suite green locally; new tests cover the mixed-dimensions error path and the state-encoding compatibility. - End-to-end memory-accounting validation on our deployment is in progress (with exact accounting, a ~500 MB collect under a 512 MiB pool must now spill or fail with `ResourcesExhausted` instead of silently overrunning RSS); will report results here. ## Related follow-ups (intentionally not in this PR) - DataFusion `GroupsAccumulatorAdapter`: per-group `indices` Vec capacity never enters `allocation_bytes`, plus a small inline double-count (will be proposed upstream separately). - A native `GroupsAccumulator` for `ST_Collect_Agg` (shared buffer + per-group offsets), which would remove most of the remaining per-group adapter overhead and re-enable partial aggregation at high cardinality. - `ItemCrsAccumulator` has the same class of gap (per-group `crs: Option<String>` heap uncounted). -- 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]
