comphead commented on code in PR #26074:
URL: https://github.com/apache/datafusion/pull/26074#discussion_r4198007007
##########
datafusion/physical-plan/src/aggregates/order/partial.rs:
##########
@@ -267,9 +266,10 @@ impl GroupOrderingPartial {
Ok(())
}
- /// Return the size of memory allocated by this structure
- pub(crate) fn size(&self) -> usize {
- size_of::<Self>() + self.order_indices.allocated_size() +
self.state.size()
+ /// Returns retained heap allocations, excluding the inline descriptor
+ /// already counted by [`super::GroupOrdering::size`].
+ pub(crate) fn heap_size(&self) -> usize {
+ self.order_indices.allocated_size() + self.state.size()
Review Comment:
Not introduced here, so a follow-up under #23393 is fine. `order_indices` is
charged by capacity, but `State::size` charges `sort_key` by length.
`get_row_at_idx` collects into `Result<Vec<_>>`, and a std-only stand-in with a
64-byte element ends at `len = 1`, `capacity = 4`. So a one-column key likely
holds 3 spare `ScalarValue` slots (192 bytes) that `heap_size` does not report,
and the `in_progress` expectation in the new test pins that length-based value.
`Vec<ScalarValue>` already implements `DFHeapSize` (capacity based), which
`CountGroupsAccumulator::size` uses the same way. `ScalarValue::size_of_vec`
minus the `Vec` header, as in `nth_value.rs`, is the other existing option.
##########
datafusion/physical-plan/src/aggregates/order/mod.rs:
##########
@@ -164,7 +166,75 @@ mod tests {
use std::sync::Arc;
- use arrow::array::Int32Array;
+ use arrow::array::{Int32Array, StringArray};
+ use datafusion_common::ScalarValue;
+
+ #[test]
+ fn test_size_none() {
+ assert_eq!(GroupOrdering::None.size(), size_of::<GroupOrdering>());
+ }
+
+ #[test]
+ fn test_size_full() -> Result<()> {
+ let mut ordering = GroupOrdering::try_new(&InputOrderMode::Sorted)?;
+ let expected = size_of::<GroupOrdering>();
+ assert_eq!(ordering.size(), expected);
+
+ ordering.new_groups(&[], &[0, 1, 2], 3)?;
+ assert_eq!(ordering.size(), expected);
+ ordering.remove_groups(2);
+ assert_eq!(ordering.size(), expected);
+ ordering.input_done();
+ assert_eq!(ordering.size(), expected);
+ ordering.reset();
+ assert_eq!(ordering.size(), expected);
+ Ok(())
+ }
Review Comment:
One test is enough here: assert that `None` and a `Sorted` ordering both
report exactly `size_of::<GroupOrdering>()`. `None | Full(_)` never reads the
variant, so the `new_groups`, `remove_groups`, `input_done` and `reset` steps
cannot change the result. The first `Full` assertion already catches the old
code, and `test_size_none` passed before this PR too.
##########
datafusion/physical-plan/src/aggregates/order/mod.rs:
##########
@@ -164,7 +166,75 @@ mod tests {
use std::sync::Arc;
- use arrow::array::Int32Array;
+ use arrow::array::{Int32Array, StringArray};
+ use datafusion_common::ScalarValue;
+
+ #[test]
+ fn test_size_none() {
+ assert_eq!(GroupOrdering::None.size(), size_of::<GroupOrdering>());
+ }
+
+ #[test]
+ fn test_size_full() -> Result<()> {
+ let mut ordering = GroupOrdering::try_new(&InputOrderMode::Sorted)?;
+ let expected = size_of::<GroupOrdering>();
+ assert_eq!(ordering.size(), expected);
+
+ ordering.new_groups(&[], &[0, 1, 2], 3)?;
+ assert_eq!(ordering.size(), expected);
+ ordering.remove_groups(2);
+ assert_eq!(ordering.size(), expected);
+ ordering.input_done();
+ assert_eq!(ordering.size(), expected);
+ ordering.reset();
+ assert_eq!(ordering.size(), expected);
+ Ok(())
+ }
+
+ #[test]
+ fn test_size_partial_retained_allocations() -> Result<()> {
+ // Grow then truncate: unused order-index capacity must remain charged.
+ let mut order_indices = vec![0];
+ order_indices.extend(1..32);
+ order_indices.truncate(1);
+ assert!(order_indices.capacity() > order_indices.len());
+ let expected =
+ size_of::<GroupOrdering>() + order_indices.capacity() *
size_of::<usize>();
+ let mut ordering =
+
GroupOrdering::Partial(GroupOrderingPartial::try_new(order_indices)?);
+ assert_eq!(ordering.size(), expected);
+
+ // A variable-width key checks both the scalar descriptor and its
payload.
+ let key = "retained sort key";
+ let batch_group_values: Vec<ArrayRef> =
+ vec![Arc::new(StringArray::from(vec!["a", key]))];
+ ordering.new_groups(&batch_group_values, &[0, 1], 2)?;
+ let in_progress = expected +
ScalarValue::Utf8(Some(key.to_owned())).size();
+ assert_eq!(ordering.size(), in_progress);
+ ordering.remove_groups(1);
+ assert_eq!(ordering.size(), in_progress);
+
+ // Replacing the key charges only the new payload, not the previous
key.
+ let replacement_key = "updated sort key with a longer payload";
+ let replacement_group_values: Vec<ArrayRef> =
+ vec![Arc::new(StringArray::from(vec![replacement_key]))];
+ ordering.new_groups(&replacement_group_values, &[1], 2)?;
+ let replaced =
+ expected +
ScalarValue::Utf8(Some(replacement_key.to_owned())).size();
+ assert!(replaced > in_progress);
+ assert_eq!(ordering.size(), replaced);
+
+ // Completing or resetting drops the key, but retains order-index
capacity.
+ ordering.input_done();
+ assert_eq!(ordering.size(), expected);
+ ordering.reset();
+ assert_eq!(ordering.size(), expected);
+ ordering.new_groups(&batch_group_values, &[0, 1], 2)?;
+ assert_eq!(ordering.size(), in_progress);
+ ordering.reset();
+ assert_eq!(ordering.size(), expected);
Review Comment:
`heap_size` is a pure function of the current state, so these steps cannot
fail independently of the earlier ones:
- `remove_groups` leaves `sort_key` untouched.
- The key replacement block hits the same `InProgress` arm as the first
`new_groups` call. Replacement itself is already covered by
`test_group_ordering_partial` in `partial.rs`.
- `reset` and the re-populate steps repeat the `Start` and `InProgress`
assertions made above.
The spare-capacity assert, one `new_groups` call with the `Utf8` key and the
`input_done` assert cover every accounting term this PR touches at about half
the length.
--
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]
