From d3cef8c48b42f2e014be89431c9ad9214187a337 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Tue, 6 Oct 2026 16:12:35 +0800 Subject: [PATCH 1/5] Fix enum descriptor, preserve allocations, add lifecycle tests - Enum descriptor counted once. - Partial retained allocations preserved. - Added None/full/partial capacity and lifecycle tests. --- .../src/aggregates/order/full.rs | 5 -- .../physical-plan/src/aggregates/order/mod.rs | 68 +++++++++++++++++-- .../src/aggregates/order/partial.rs | 8 +-- 3 files changed, 68 insertions(+), 13 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/order/full.rs b/datafusion/physical-plan/src/aggregates/order/full.rs index ca818d6a2d598..5b17d527e52f9 100644 --- a/datafusion/physical-plan/src/aggregates/order/full.rs +++ b/datafusion/physical-plan/src/aggregates/order/full.rs @@ -16,7 +16,6 @@ // under the License. use datafusion_expr::EmitTo; -use std::mem::size_of; /// Tracks grouping state when the data is ordered entirely by its /// group keys @@ -143,10 +142,6 @@ impl GroupOrderingFull { } }; } - - pub(crate) fn size(&self) -> usize { - size_of::() - } } impl Default for GroupOrderingFull { diff --git a/datafusion/physical-plan/src/aggregates/order/mod.rs b/datafusion/physical-plan/src/aggregates/order/mod.rs index 147e4e0d6f18d..b540e6e90f710 100644 --- a/datafusion/physical-plan/src/aggregates/order/mod.rs +++ b/datafusion/physical-plan/src/aggregates/order/mod.rs @@ -148,12 +148,14 @@ impl GroupOrdering { } /// Returns the size of memory used by the ordering state, in bytes. + /// + /// Includes the enum descriptor once (covering the inline active variant) + /// and any heap allocations retained by that variant. pub fn size(&self) -> usize { size_of::() + match self { - GroupOrdering::None => 0, - GroupOrdering::Partial(partial) => partial.size(), - GroupOrdering::Full(full) => full.size(), + GroupOrdering::None | GroupOrdering::Full(_) => 0, + GroupOrdering::Partial(partial) => partial.heap_size(), } } } @@ -164,7 +166,65 @@ 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::()); + } + + #[test] + fn test_size_full() -> Result<()> { + let mut ordering = GroupOrdering::try_new(&InputOrderMode::Sorted)?; + let expected = size_of::(); + 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::() + order_indices.capacity() * size_of::(); + 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 = + 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); + + // 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); + Ok(()) + } #[test] fn test_oom_emit_to_none_ordering() { diff --git a/datafusion/physical-plan/src/aggregates/order/partial.rs b/datafusion/physical-plan/src/aggregates/order/partial.rs index 1603bb6d079be..86ff73d42fb31 100644 --- a/datafusion/physical-plan/src/aggregates/order/partial.rs +++ b/datafusion/physical-plan/src/aggregates/order/partial.rs @@ -16,7 +16,6 @@ // under the License. use std::cmp::Ordering; -use std::mem::size_of; use std::sync::Arc; use arrow::array::ArrayRef; @@ -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.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() } } From 4090508092fd4f13c2a261124eb2b60e8304f0d8 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Tue, 6 Oct 2026 16:19:09 +0800 Subject: [PATCH 2/5] 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 --- datafusion/physical-plan/src/aggregates/order/mod.rs | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/datafusion/physical-plan/src/aggregates/order/mod.rs b/datafusion/physical-plan/src/aggregates/order/mod.rs index b540e6e90f710..54f0f808365f4 100644 --- a/datafusion/physical-plan/src/aggregates/order/mod.rs +++ b/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 = + 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); From a2b65af1c2dccb3b6da07de2dc249ab9836887fa Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 7 Oct 2026 14:45:23 +0800 Subject: [PATCH 3/5] Rename State::size to heap_size MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Renamed State::size → heap_size. - Used with_capacity(32) + push(0); removed redundant assertion. - Documented inline-only variants. --- datafusion/physical-plan/src/aggregates/order/mod.rs | 9 ++++----- datafusion/physical-plan/src/aggregates/order/partial.rs | 4 ++-- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/order/mod.rs b/datafusion/physical-plan/src/aggregates/order/mod.rs index 54f0f808365f4..4c28c8b9b4a39 100644 --- a/datafusion/physical-plan/src/aggregates/order/mod.rs +++ b/datafusion/physical-plan/src/aggregates/order/mod.rs @@ -154,6 +154,7 @@ impl GroupOrdering { pub fn size(&self) -> usize { size_of::() + match self { + // These variants hold only inline state, already counted above. GroupOrdering::None | GroupOrdering::Full(_) => 0, GroupOrdering::Partial(partial) => partial.heap_size(), } @@ -193,11 +194,9 @@ mod tests { #[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()); + // The direct partial constructor retains spare order-index capacity. + let mut order_indices = Vec::with_capacity(32); + order_indices.push(0); let expected = size_of::() + order_indices.capacity() * size_of::(); let mut ordering = diff --git a/datafusion/physical-plan/src/aggregates/order/partial.rs b/datafusion/physical-plan/src/aggregates/order/partial.rs index 86ff73d42fb31..2488c24913c8a 100644 --- a/datafusion/physical-plan/src/aggregates/order/partial.rs +++ b/datafusion/physical-plan/src/aggregates/order/partial.rs @@ -101,7 +101,7 @@ enum State { } impl State { - fn size(&self) -> usize { + fn heap_size(&self) -> usize { match self { State::Taken => 0, State::Start => 0, @@ -269,7 +269,7 @@ impl GroupOrderingPartial { /// 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() + self.order_indices.allocated_size() + self.state.heap_size() } } From 99232b266ffcfeeb404e9d5524f3255132e00413 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 7 Oct 2026 14:50:42 +0800 Subject: [PATCH 4/5] fix: simplify test_size_inline_only - test_size_inline_only now checks None and Sorted only - lifecycle steps were removed --- .../physical-plan/src/aggregates/order/mod.rs | 19 +++---------------- 1 file changed, 3 insertions(+), 16 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/order/mod.rs b/datafusion/physical-plan/src/aggregates/order/mod.rs index 4c28c8b9b4a39..64c7ba6237715 100644 --- a/datafusion/physical-plan/src/aggregates/order/mod.rs +++ b/datafusion/physical-plan/src/aggregates/order/mod.rs @@ -171,23 +171,10 @@ mod tests { use datafusion_common::ScalarValue; #[test] - fn test_size_none() { - assert_eq!(GroupOrdering::None.size(), size_of::()); - } - - #[test] - fn test_size_full() -> Result<()> { - let mut ordering = GroupOrdering::try_new(&InputOrderMode::Sorted)?; + fn test_size_inline_only() -> Result<()> { let expected = size_of::(); - 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!(GroupOrdering::None.size(), expected); + let ordering = GroupOrdering::try_new(&InputOrderMode::Sorted)?; assert_eq!(ordering.size(), expected); Ok(()) } From df90baf93560d9a01c1fa84f9a9a5108e3a7382d Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 7 Oct 2026 14:54:34 +0800 Subject: [PATCH 5/5] Remove checks for removal, replacement, reset, and re-population - Removed removal, replacement, reset, and re-population checks. - Kept capacity, Utf8 update, and completion assertions. --- .../physical-plan/src/aggregates/order/mod.rs | 20 +------------------ 1 file changed, 1 insertion(+), 19 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/order/mod.rs b/datafusion/physical-plan/src/aggregates/order/mod.rs index 64c7ba6237715..88ee2af5f26da 100644 --- a/datafusion/physical-plan/src/aggregates/order/mod.rs +++ b/datafusion/physical-plan/src/aggregates/order/mod.rs @@ -197,28 +197,10 @@ mod tests { 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 = - 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. + // Completing 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); Ok(()) }