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..88ee2af5f26da 100644 --- a/datafusion/physical-plan/src/aggregates/order/mod.rs +++ b/datafusion/physical-plan/src/aggregates/order/mod.rs @@ -148,12 +148,15 @@ 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(), + // These variants hold only inline state, already counted above. + GroupOrdering::None | GroupOrdering::Full(_) => 0, + GroupOrdering::Partial(partial) => partial.heap_size(), } } } @@ -164,7 +167,42 @@ mod tests { use std::sync::Arc; - use arrow::array::Int32Array; + use arrow::array::{Int32Array, StringArray}; + use datafusion_common::ScalarValue; + + #[test] + fn test_size_inline_only() -> Result<()> { + let expected = size_of::(); + assert_eq!(GroupOrdering::None.size(), expected); + let ordering = GroupOrdering::try_new(&InputOrderMode::Sorted)?; + assert_eq!(ordering.size(), expected); + Ok(()) + } + + #[test] + fn test_size_partial_retained_allocations() -> Result<()> { + // 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 = + 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); + + // Completing drops the key, but retains order-index capacity. + ordering.input_done(); + 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..2488c24913c8a 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; @@ -102,7 +101,7 @@ enum State { } impl State { - fn size(&self) -> usize { + fn heap_size(&self) -> usize { match self { State::Taken => 0, State::Start => 0, @@ -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.heap_size() } }