kosiew commented on code in PR #26074:
URL: https://github.com/apache/datafusion/pull/26074#discussion_r4203938158
##########
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:
I’ll keep only the spare-capacity setup, initial assertion, one `Utf8` key
update, and `input_done` assertion.
--
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]
