You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@sedona.apache.org by "james-willis (via GitHub)" <gi...@apache.org> on 2026/09/02 22:24:42 UTC
[PR] perf(rust/sedona-functions): replace ST_Collect_Agg's per-group hash sets with the type/dimension bitset [sedona-db]
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf(rust/sedona-functions): replace ST_Collect_Agg's per-group hash sets with the type/dimension bitset [sedona-db]
Posted by "james-willis (via GitHub)" <gi...@apache.org>.
james-willis commented on code in PR #1225:
URL: https://github.com/apache/sedona-db/pull/1225#discussion_r3929236585
##########
rust/sedona-functions/src/st_collect_agg.rs:
##########
@@ -275,10 +281,24 @@ impl Accumulator for CollectionAccumulator {
)
.map_err(|e| DataFusionError::External(Box::new(e)))?
.into_iter()
- .map(|item| item.dimensions());
-
- self.unique_geometry_types.extend(geometry_types);
- self.unique_dimensions.extend(dimensions);
+ .map(|item| item.dimensions())
+ .collect::<Vec<_>>();
+
+ // The state stores the two marginals; only the marginals
+ // are ever consumed, so inserting the cross product
+ // reconstructs them exactly in the pair bitset. (A state
+ // produced by update_batch never has one marginal empty
+ // while the other is not.)
+ for geometry_type in &geometry_types {
+ for dimensions in &dimensions {
+ self.types_and_dims
+ .insert(&GeometryTypeAndDimensions::new(
+ *geometry_type,
+ *dimensions,
+ ))
+ .map_err(|e| DataFusionError::External(Box::new(e)))?;
+ }
+ }
Review Comment:
Yeah I was overly cautious trying to maintain the wire format. I can't think of a reason that it'd be neccessary.
Serializing to use the u32 simplifies all the code here.
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf(rust/sedona-functions): replace ST_Collect_Agg's per-group hash sets with the type/dimension bitset [sedona-db]
Posted by "james-willis (via GitHub)" <gi...@apache.org>.
james-willis commented on code in PR #1225:
URL: https://github.com/apache/sedona-db/pull/1225#discussion_r3929237755
##########
rust/sedona-functions/src/st_collect_agg.rs:
##########
@@ -205,14 +208,17 @@ impl Accumulator for CollectionAccumulator {
}
fn state(&mut self) -> Result<Vec<ScalarValue>> {
- let geometry_types_value =
- serde_json::to_string(&self.unique_geometry_types.iter().collect::<Vec<_>>())
- .map_err(|e| DataFusionError::External(Box::new(e)))?;
+ // Both columns keep the exact pre-bitset wire format: a JSON list of
+ // geometry types, and a JSON list of dimensions wrapped as
+ // (Geometry, dimensions) pairs.
+ let geometry_types_value = serde_json::to_string(&self.types_and_dims.geometry_types())
Review Comment:
agreed. this really simplifies the code here. thanks
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf(rust/sedona-functions): replace ST_Collect_Agg's per-group hash sets with the type/dimension bitset [sedona-db]
Posted by "james-willis (via GitHub)" <gi...@apache.org>.
james-willis commented on code in PR #1225:
URL: https://github.com/apache/sedona-db/pull/1225#discussion_r3929228361
##########
rust/sedona-geometry/src/types.rs:
##########
@@ -429,6 +429,28 @@ impl GeometryTypeAndDimensionsSet {
current_bit: 0,
}
}
+
+ /// The distinct geometry types in this set, in bitset iteration order.
+ pub fn geometry_types(&self) -> Vec<GeometryTypeId> {
+ let mut out = Vec::new();
+ for item in self.iter() {
+ if !out.contains(&item.geometry_type()) {
+ out.push(item.geometry_type());
+ }
+ }
+ out
+ }
+
+ /// The distinct dimensions in this set, in bitset iteration order.
+ pub fn dimensions(&self) -> Vec<Dimensions> {
+ let mut out = Vec::new();
+ for item in self.iter() {
+ if !out.contains(&item.dimensions()) {
+ out.push(item.dimensions());
+ }
+ }
+ out
+ }
Review Comment:
done. probably doesnt really matter.
I added them to the class because I assumed itd be useful for future users and doesn't really cost anything. seems like a better home than in the UDF impl
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf(rust/sedona-functions): replace ST_Collect_Agg's per-group hash sets with the type/dimension bitset [sedona-db]
Posted by "james-willis (via GitHub)" <gi...@apache.org>.
james-willis commented on code in PR #1225:
URL: https://github.com/apache/sedona-db/pull/1225#discussion_r3962378570
##########
rust/sedona-geometry/src/types.rs:
##########
@@ -429,6 +441,37 @@ impl GeometryTypeAndDimensionsSet {
current_bit: 0,
}
}
+
+ /// 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()
+ }
Review Comment:
done
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf(rust/sedona-functions): replace ST_Collect_Agg's per-group hash sets with the type/dimension bitset [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #1225:
URL: https://github.com/apache/sedona-db/pull/1225#discussion_r3920719379
##########
rust/sedona-functions/src/st_collect_agg.rs:
##########
@@ -205,14 +208,17 @@ impl Accumulator for CollectionAccumulator {
}
fn state(&mut self) -> Result<Vec<ScalarValue>> {
- let geometry_types_value =
- serde_json::to_string(&self.unique_geometry_types.iter().collect::<Vec<_>>())
- .map_err(|e| DataFusionError::External(Box::new(e)))?;
+ // Both columns keep the exact pre-bitset wire format: a JSON list of
+ // geometry types, and a JSON list of dimensions wrapped as
+ // (Geometry, dimensions) pairs.
+ let geometry_types_value = serde_json::to_string(&self.types_and_dims.geometry_types())
Review Comment:
I don't think the wire format needs to be maintained?
##########
rust/sedona-geometry/src/types.rs:
##########
@@ -429,6 +429,28 @@ impl GeometryTypeAndDimensionsSet {
current_bit: 0,
}
}
+
+ /// The distinct geometry types in this set, in bitset iteration order.
+ pub fn geometry_types(&self) -> Vec<GeometryTypeId> {
+ let mut out = Vec::new();
+ for item in self.iter() {
+ if !out.contains(&item.geometry_type()) {
+ out.push(item.geometry_type());
+ }
+ }
+ out
+ }
+
+ /// The distinct dimensions in this set, in bitset iteration order.
+ pub fn dimensions(&self) -> Vec<Dimensions> {
+ let mut out = Vec::new();
+ for item in self.iter() {
+ if !out.contains(&item.dimensions()) {
+ out.push(item.dimensions());
+ }
+ }
+ out
+ }
Review Comment:
These seem like inefficient ways (particular vector contains) to do these...I think you can probably create bitmasks and & them against the current value if you truly need these.
##########
rust/sedona-functions/src/st_collect_agg.rs:
##########
@@ -525,4 +545,61 @@ mod test {
"CRS values not equal: ogc:crs84 vs epsg:3857"
);
}
+
+ /// The bitset swap must not change the serialized state wire format:
+ /// column 0 is a JSON list of geometry types and column 1 a JSON list of
+ /// dimensions wrapped as (Geometry, dims) pairs, exactly as the HashSet
+ /// implementation produced, so states merge across versions.
+ #[test]
+ fn state_wire_format_is_unchanged() {
Review Comment:
States merging across versions sounds fishy. Am I missing something about how this can be run?
##########
rust/sedona-functions/src/st_collect_agg.rs:
##########
@@ -275,10 +281,24 @@ impl Accumulator for CollectionAccumulator {
)
.map_err(|e| DataFusionError::External(Box::new(e)))?
.into_iter()
- .map(|item| item.dimensions());
-
- self.unique_geometry_types.extend(geometry_types);
- self.unique_dimensions.extend(dimensions);
+ .map(|item| item.dimensions())
+ .collect::<Vec<_>>();
+
+ // The state stores the two marginals; only the marginals
+ // are ever consumed, so inserting the cross product
+ // reconstructs them exactly in the pair bitset. (A state
+ // produced by update_batch never has one marginal empty
+ // while the other is not.)
+ for geometry_type in &geometry_types {
+ for dimensions in &dimensions {
+ self.types_and_dims
+ .insert(&GeometryTypeAndDimensions::new(
+ *geometry_type,
+ *dimensions,
+ ))
+ .map_err(|e| DataFusionError::External(Box::new(e)))?;
+ }
+ }
Review Comment:
How about just storing the 64-bit integer as the state? Otherwise, this comment needs to be rewritten for humans.
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf(rust/sedona-functions): replace ST_Collect_Agg's per-group hash sets with the type/dimension bitset [sedona-db]
Posted by "james-willis (via GitHub)" <gi...@apache.org>.
james-willis merged PR #1225:
URL: https://github.com/apache/sedona-db/pull/1225
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf(rust/sedona-functions): replace ST_Collect_Agg's per-group hash sets with the type/dimension bitset [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #1225:
URL: https://github.com/apache/sedona-db/pull/1225#discussion_r3937720950
##########
rust/sedona-geometry/src/types.rs:
##########
@@ -429,6 +441,37 @@ impl GeometryTypeAndDimensionsSet {
current_bit: 0,
}
}
+
+ /// 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()
+ }
Review 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 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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] perf(rust/sedona-functions): replace ST_Collect_Agg's per-group hash sets with the type/dimension bitset [sedona-db]
Posted by "james-willis (via GitHub)" <gi...@apache.org>.
james-willis commented on code in PR #1225:
URL: https://github.com/apache/sedona-db/pull/1225#discussion_r3929243960
##########
rust/sedona-functions/src/st_collect_agg.rs:
##########
@@ -525,4 +545,61 @@ mod test {
"CRS values not equal: ogc:crs84 vs epsg:3857"
);
}
+
+ /// The bitset swap must not change the serialized state wire format:
+ /// column 0 is a JSON list of geometry types and column 1 a JSON list of
+ /// dimensions wrapped as (Geometry, dims) pairs, exactly as the HashSet
+ /// implementation produced, so states merge across versions.
+ #[test]
+ fn state_wire_format_is_unchanged() {
Review Comment:
dropped this whole story by just changing the wire format. Agree that theres no reason to be consistent across version about it. AFAIK all uses of it are ephemeral
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org