Extend test for datafusion/physical-plan/src/aggregates/order/mod.rs

- Test coverage extended to cover longer sequences
- Replacement key charges only new allocation, not old key

Author: kosiew

datafusion/physical-plan/src/aggregates/order/mod.rs

@@ -214,6 +214,16 @@ mod tests {
         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);