perf(rust/sedona-functions): replace ST_Collect_Agg's per-group hash sets with the type/dimension bitset - #1225
Open
james-willis wants to merge 4 commits into
Open
Conversation
james-willis
force-pushed
the
jw/collect-agg-accounting
branch
from
September 2, 2026 22:41
03f4e8c to
e920d4f
Compare
…sets with the type/dimension bitset CollectionAccumulator kept two per-group hashbrown sets (unique_geometry_types, unique_dimensions). Each allocates its table on the first insert - i.e. in every non-empty group - costing ~100 heap bytes per group that Accumulator::size() reported as ZERO, plus two mallocs (and matching frees) per accumulated row. On a traced 12 GiB container running a collect aggregation over 60M rows into ~10-16M groups, the aggregate state was invisible to the memory pool: anon heap peaked 1.8 GB PAST the pool limit with zero spill the entire run. Replace both sets with the existing GeometryTypeAndDimensionsSet inline u32 bitset (already used by st_analyze_agg): size_of the accumulator drops 184 -> 112 bytes, the per-group heap allocations disappear (~2.8 GB real at 16.5M groups), size() becomes exact, and the per-row update path loses both mallocs. Semantics: - update_batch uses insert() - NOT insert_or_ignore() - and propagates the error, so geometries with unknown (non-XY/XYZ/XYM/XYZM) dimensions fail loudly instead of being silently dropped from the dimension check; the hash sets used to retain them. The existing "Can't ST_Collect_Agg() mixed dimension geometries" behavior in make_wkb_result is preserved via the dimensions marginal. - The serialized state wire format is unchanged: column 0 stays a JSON list of geometry types and column 1 a JSON list of (Geometry, dims) pairs, both now derived from the bitset, so states merge with previous versions. merge_batch reconstructs the pair bitset from the two marginals via their cross product; only the marginals are ever consumed, and a state produced by update_batch never has one marginal empty while the other is not. Tests: state wire format pinned byte-for-byte for the Point/XY case; size() exactness (inserts allocate nothing); the existing mixed dimension and CRS tests cover the error paths.
james-willis
force-pushed
the
jw/collect-agg-accounting
branch
from
September 2, 2026 22:49
e920d4f to
c96a525
Compare
james-willis
marked this pull request as ready for review
September 2, 2026 22:54
james-willis
requested review from
Kontinuation,
jiayuasu and
paleolimbot
and removed request for
zhangfengcdt
September 2, 2026 22:54
paleolimbot
reviewed
Sep 3, 2026
…tset Per review: the aggregate intermediate state is ephemeral within a query execution and is never persisted or merged across engine versions, so it need not stay byte-compatible with the pre-bitset HashSet JSON. Store the GeometryTypeAndDimensionsSet's u32 directly as a single Int64 state column (dropping the two JSON marginal columns), so merge_batch becomes a bitwise OR instead of parsing JSON and rebuilding the pair set via a cross product. Swaps the wire-format-invariance test for a state round-trip test.
Per review: geometry_types()/dimensions() deduped via Vec::contains in a loop. Since the set is a u32 with geometry types at fixed bit offsets per dimension byte, the distinct types are the OR of the four dimension bytes and the distinct dimensions are the non-empty bytes -- pure bit ops, no per-element scan.
james-willis
force-pushed
the
jw/collect-agg-accounting
branch
from
September 3, 2026 22:19
fa85c55 to
5144c42
Compare
paleolimbot
approved these changes
Sep 4, 2026
paleolimbot
left a comment
paleolimbot
left a comment
Member
There was a problem hiding this comment.
Thank you!
Just the one test and then this is good to go.
Comment on lines
+445
to
+474
| /// The distinct geometry types in this set, in ascending WKB-id order. | ||
| pub fn geometry_types(&self) -> Vec<GeometryTypeId> { | ||
| // Collapse the four dimension bytes onto one: bit i is set iff a | ||
| // geometry with WKB id i is present under any dimension. | ||
| let merged = | ||
| (self.types | (self.types >> 8) | (self.types >> 16) | (self.types >> 24)) & 0xFF; | ||
| (0..8) | ||
| .filter(|i| merged & (1 << i) != 0) | ||
| .map(|i| { | ||
| GeometryTypeId::try_from_wkb_id(i) | ||
| .expect("Invalid geometry type wkb_id in GeometryTypeAndDimensionsSet") | ||
| }) | ||
| .collect() | ||
| } | ||
|
|
||
| /// The distinct dimensions in this set, in XY, XYZ, XYM, XYZM order. | ||
| pub fn dimensions(&self) -> Vec<Dimensions> { | ||
| // Each dimension occupies one byte; the dimension is present iff its | ||
| // byte has any type bit set. | ||
| [ | ||
| (0x0000_00FF, Dimensions::Xy), | ||
| (0x0000_FF00, Dimensions::Xyz), | ||
| (0x00FF_0000, Dimensions::Xym), | ||
| (0xFF00_0000_u32, Dimensions::Xyzm), | ||
| ] | ||
| .into_iter() | ||
| .filter(|(mask, _)| self.types & mask != 0) | ||
| .map(|(_, dim)| dim) | ||
| .collect() | ||
| } |
Member
There was a problem hiding this comment.
If you are going to add these two functions you should check the matrix of geometry type by dimensions that each pair correctly returns the correct list-of-one GeometryType and Dimensions (I don't think there's anything making sure the bits are correct at the moment)
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 this PR does
ST_Collect_Agg'sAccumulator::size()under-reported memory: it countedsize_of::<CollectionAccumulator>() + item.capacity()but omitted the heap ofunique_geometry_types: HashSet<GeometryTypeId>andunique_dimensions: HashSet<Dimensions>(~100 bytes/group on first insert, measured). Since a UDAF accumulator is instantiated per group, a high-cardinalityGROUP BYhides ~100 bytes/group from the memory pool — on SpatialBench Q5 at sf=10 (10M+ groups, 12 GiB container, 7.73 GB fair pool) the query peaked at 9.55 GB anon with zero spills and noResourcesExhausted, so survival was luck rather than the pool's decision.This replaces both HashSets with the existing
GeometryTypeAndDimensionsSetu32 bitset (rust/sedona-geometry/src/types.rs):size()becomes exact with no extra terms;convert_to_statepath (used when skip-partial-aggregation engages at high cardinality) sheds two malloc/free pairs per row.Notes
u32(a singleInt64column) andmerge_batchis a bitwise OR. (Earlier revisions kept the pre-bitset JSON wire format; dropped per review, since aggregate intermediate state is ephemeral within a query execution and never persisted or merged across versions.)insert(), notinsert_or_ignore(), so a mixed/unknown-dimension group errors rather than silently emitting an XY header. A mixed-dimensions test is included.Testing
Full
sedona-functionsandsedona-geometrysuites green. End-to-end memory validation on our deployment is in progress (with exact accounting, a ~500 MB collect under a 512 MiB pool must now spill orResourcesExhaustedinstead of silently overrunning RSS); will report here.Follow-ups (not in this PR)
GroupsAccumulatorAdapter: per-groupindicesVec capacity never entersallocation_bytes, plus a small inline double-count — to be proposed upstream separately.GroupsAccumulatorforST_Collect_Agg(shared buffer + per-group offsets) to remove the remaining per-group adapter overhead and re-enable partial aggregation.ItemCrsAccumulatorhas the same class of gap (per-groupcrs: Option<String>uncounted).