You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "LiaCastaneda (via GitHub)" <gi...@apache.org> on 2025/12/26 16:53:55 UTC
[PR] Use arrow pool to fix memory over accounting aggregations [datafusion]
LiaCastaneda opened a new pull request, #19501:
URL: https://github.com/apache/datafusion/pull/19501
## Which issue does this PR close?
Closes #
Related to [the https://github.com/apache/datafusion/issues/16841#issuecomment-3563643947](https://github.com/apache/datafusion/issues/16841#issue-3248316384)
## Rationale for this change
Aggregation accumulators that store Arrow Arrays can have memory over accounting when array buffers are shared between multiple ScalarValues or directly Arrow Arrays. This occurs because doing ScalarValue::try_from_array() create slices that reference the same underlying buffers and each ScalarValue reports the full buffer size when calculating memory usage and the same physical buffer gets counted multiple times, leading to over accounting ([this](https://github.com/apache/datafusion/issues/16841#issuecomment-3107483475) comment explains very well why we are seeing this)
There have been several attempts to fix this before which included compacting data to not keep the whole array alive, or using `get_slice_memory_size` instead of `get_array_memory_size`. However, we observed that both had downsides:
- `compact()` was CPU inefficient (copies data which can be expensive)
- `get_slice_memory_size()` only accounts for logical memory, but it is not the actual physical buffer capacity, therefore the amount returned is not accurate.
## What changes are included in this PR?
This approach avoids double-counting memory by using Arrow's TrackingMemoryPool, which automatically deduplicates shared buffers when accounting them. This means we don't need to compact() or call `get_slice_memory_size()` just to solve the accounting problem. Note that compact() might still be useful when we want to release memory pressure.
- Updated Accumulator::size() and GroupsAccumulator::size() signatures to accept Option<&dyn MemoryPool>:
- When pool is `None` Returns total memory size including Arrow buffers using either get_slice_memory_size or just ScalarValue-> size () (same as before, so its backward compatible)
- When pool is Some iteturns structural size only and claims buffers with the pool for deduplication tracking. Callers using the pool must add pool.used() to get total memory.
Updated accumulators that use the pool parameter:
- DistinctCountAccumulator
- ArrayAggAccumulator
- For OrderSensitiveArrayAggAccumulator and DistinctArrayAggAccumulator I removed the compacting since it was introduced specifically to solve the over accounting, and its not needed anymore.
- FirstValueAccumulator / LastValueAccumulator
All other accumulator implementations had to be updated to match new signature
## Are these changes tested?
Added `distinct_count_does_not_over_account_memory() `test to test memory pool deduplication for COUNT(DISTINCT) with array types. Also updated the existing accumulator tests to use memory pool, it verifies the accounted memory is still less than when not using the memory pool (in some cases even less than when we compacted).
## Are there any user-facing changes?
yes, the API size for Accumulators and GroupAccumulators changed from` fn size(&self) -> usize;` to `fn size(&self, pool: Option<&dyn MemoryPool>) -> usize;
`
Not sure if this is the best API design... I'm open to suggestions. In any case, if None is passed the behavior will remain the same as before. Also IIUC this function is mainly used to keep the DF memory pool within its bounds during aggregations.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use arrow pool to fix memory over accounting in aggregations [datafusion]
Posted by "LiaCastaneda (via GitHub)" <gi...@apache.org>.
LiaCastaneda commented on code in PR #19501:
URL: https://github.com/apache/datafusion/pull/19501#discussion_r2655376608
##########
datafusion/functions-aggregate/src/array_agg.rs:
##########
@@ -518,13 +518,29 @@ impl Accumulator for DistinctArrayAggAccumulator {
Ok(ScalarValue::List(arr))
}
- fn size(&self) -> usize {
- size_of_val(self) + ScalarValue::size_of_hashset(&self.values)
- - size_of_val(&self.values)
+ fn size(&self, pool: Option<&dyn MemoryPool>) -> usize {
+ let mut total = size_of_val(self)
+ + size_of_val(&self.values)
+ + (size_of::<ScalarValue>() * self.values.capacity())
+ self.datatype.size()
- size_of_val(&self.datatype)
- size_of_val(&self.sort_options)
- + size_of::<Option<SortOptions>>()
+ + size_of::<Option<SortOptions>>();
+
+ for scalar in &self.values {
+ if let Some(array) = scalar.get_array_ref() {
+ total += size_of::<Arc<dyn Array>>();
+ if let Some(pool) = pool {
+ claim_buffers_recursive(&array.to_data(), pool);
Review Comment:
Yes, that's why after calling size() we add arrow_pool.used() to account for the buffer memory (I forgot to add it). I also considered calling arrow_pool.used() directly inside each accumulator's size() function so the caller doesn't have to remember to do it. However, that would still cause over-accounting in scenarios like `update_memory_reservation()` where we sum size() across multiple accumulators since they can share buffers as well:
```
let total_size = self.group_values.size()
+ self.group_ordering.size()
+ self.current_group_indices.capacity() * size_of::<usize>()
+ self
.accumulators
.iter()
.map(|x| x.size(Some(&self.arrow_pool)))
.sum::<usize>() <--- If each size() returned arrow_pool.used() we would still be over counting the pool
```
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use arrow pool to fix memory over accounting in aggregations [datafusion]
Posted by "gabotechs (via GitHub)" <gi...@apache.org>.
gabotechs commented on code in PR #19501:
URL: https://github.com/apache/datafusion/pull/19501#discussion_r2655417840
##########
datafusion/functions-aggregate/src/array_agg.rs:
##########
@@ -518,13 +518,29 @@ impl Accumulator for DistinctArrayAggAccumulator {
Ok(ScalarValue::List(arr))
}
- fn size(&self) -> usize {
- size_of_val(self) + ScalarValue::size_of_hashset(&self.values)
- - size_of_val(&self.values)
+ fn size(&self, pool: Option<&dyn MemoryPool>) -> usize {
+ let mut total = size_of_val(self)
+ + size_of_val(&self.values)
+ + (size_of::<ScalarValue>() * self.values.capacity())
+ self.datatype.size()
- size_of_val(&self.datatype)
- size_of_val(&self.sort_options)
- + size_of::<Option<SortOptions>>()
+ + size_of::<Option<SortOptions>>();
+
+ for scalar in &self.values {
+ if let Some(array) = scalar.get_array_ref() {
+ total += size_of::<Arc<dyn Array>>();
+ if let Some(pool) = pool {
+ claim_buffers_recursive(&array.to_data(), pool);
Review Comment:
🤔 Ok, and that works because we have 1 `TrackingArrowPool` per `GroupedHashAggregateStream` 👍
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use arrow pool to fix memory over accounting in aggregations [datafusion]
Posted by "LiaCastaneda (via GitHub)" <gi...@apache.org>.
LiaCastaneda commented on code in PR #19501:
URL: https://github.com/apache/datafusion/pull/19501#discussion_r2652852551
##########
datafusion/functions-aggregate-common/src/min_max.rs:
##########
@@ -501,7 +502,7 @@ impl Accumulator for MinAccumulator {
Ok(self.min.clone())
}
- fn size(&self) -> usize {
+ fn size(&self, _pool: Option<&dyn MemoryPool>) -> usize {
size_of_val(self) - size_of_val(&self.min) + self.min.size()
}
Review Comment:
`MinAccumulator` and `MaxAccumulator` could also use the pool since they hold a `ScalarValue`, which could hold an Arrow Array --> have shared buffers. I didn’t do it to avoid adding more changes in this PR.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use arrow pool to fix memory over accounting in aggregations [datafusion]
Posted by "gabotechs (via GitHub)" <gi...@apache.org>.
gabotechs commented on PR #19501:
URL: https://github.com/apache/datafusion/pull/19501#issuecomment-3709742565
Even if this PR looks great, I'd love to explore what would it take to get the memory tracking fundamentals very right from the beginning, and see if there's an opportunity to be in a situation where memory tracking is more automatic, and we do not need to deal with two different `MemoryPool` traits with overlapping intentions.
Wrote a small comment about that here https://github.com/apache/arrow-rs/issues/8938#issuecomment-3709632585. What do you think?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use arrow pool to fix memory over accounting in aggregations [datafusion]
Posted by "LiaCastaneda (via GitHub)" <gi...@apache.org>.
LiaCastaneda commented on code in PR #19501:
URL: https://github.com/apache/datafusion/pull/19501#discussion_r2655323607
##########
datafusion/functions-aggregate/src/array_agg.rs:
##########
@@ -390,27 +393,25 @@ impl Accumulator for ArrayAggAccumulator {
Ok(SingleRowListArrayBuilder::new(concated_array).build_list_scalar())
}
- fn size(&self) -> usize {
- size_of_val(self)
+ fn size(&self, pool: Option<&dyn MemoryPool>) -> usize {
+ let mut total = size_of_val(self)
+ (size_of::<ArrayRef>() * self.values.capacity())
- + self
+ + self.datatype.size()
+ - size_of_val(&self.datatype);
+
+ if let Some(pool) = pool {
+ for arr in &self.values {
+ claim_buffers_recursive(&arr.to_data(), pool);
Review Comment:
it will call claim() multiple times on the same `Buffer`/`Bytes`, but is safe because each call replaces the previous reservation (and does not add to it), so the pool will only track each buffer once regardless of how many times size() is called
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use arrow pool to fix memory over accounting in aggregations [datafusion]
Posted by "LiaCastaneda (via GitHub)" <gi...@apache.org>.
LiaCastaneda commented on code in PR #19501:
URL: https://github.com/apache/datafusion/pull/19501#discussion_r2655404934
##########
datafusion/functions-aggregate/src/array_agg.rs:
##########
@@ -518,13 +518,29 @@ impl Accumulator for DistinctArrayAggAccumulator {
Ok(ScalarValue::List(arr))
}
- fn size(&self) -> usize {
- size_of_val(self) + ScalarValue::size_of_hashset(&self.values)
- - size_of_val(&self.values)
+ fn size(&self, pool: Option<&dyn MemoryPool>) -> usize {
+ let mut total = size_of_val(self)
+ + size_of_val(&self.values)
+ + (size_of::<ScalarValue>() * self.values.capacity())
+ self.datatype.size()
- size_of_val(&self.datatype)
- size_of_val(&self.sort_options)
- + size_of::<Option<SortOptions>>()
+ + size_of::<Option<SortOptions>>();
+
+ for scalar in &self.values {
+ if let Some(array) = scalar.get_array_ref() {
+ total += size_of::<Arc<dyn Array>>();
+ if let Some(pool) = pool {
+ claim_buffers_recursive(&array.to_data(), pool);
Review Comment:
So this is how we would use the pool properly ([link](https://github.com/LiaCastaneda/datafusion/blob/5a0eecbdf0b0a53c9645764d7be37ea8a2e26d4e/datafusion/physical-plan/src/aggregates/row_hash.rs#L1065)) to calculate the total size without counting shared buffers multiple times
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use arrow pool to fix memory over accounting in aggregations [datafusion]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #19501:
URL: https://github.com/apache/datafusion/pull/19501#issuecomment-4340381183
Thank you for your contribution. Unfortunately, this pull request is stale because it has been open 60 days with no activity. Please remove the stale label or comment or this will be closed in 7 days.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use arrow pool to fix memory over accounting aggregations [datafusion]
Posted by "gabotechs (via GitHub)" <gi...@apache.org>.
gabotechs commented on PR #19501:
URL: https://github.com/apache/datafusion/pull/19501#issuecomment-3698954062
Nice! I'll review this one soon
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use arrow pool to fix memory over accounting in aggregations [datafusion]
Posted by "LiaCastaneda (via GitHub)" <gi...@apache.org>.
LiaCastaneda commented on code in PR #19501:
URL: https://github.com/apache/datafusion/pull/19501#discussion_r2652852551
##########
datafusion/functions-aggregate-common/src/min_max.rs:
##########
@@ -501,7 +502,7 @@ impl Accumulator for MinAccumulator {
Ok(self.min.clone())
}
- fn size(&self) -> usize {
+ fn size(&self, _pool: Option<&dyn MemoryPool>) -> usize {
size_of_val(self) - size_of_val(&self.min) + self.min.size()
}
Review Comment:
`MinAccumulator` and `MaxAccumulator` could also use the pool since they hold a `ScalarValue`, which could hold an Arrow Array --> have shared buffers. I didn’t do it to avoid adding more changes in this PR.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use arrow pool to fix memory over accounting in aggregations [datafusion]
Posted by "LiaCastaneda (via GitHub)" <gi...@apache.org>.
LiaCastaneda commented on code in PR #19501:
URL: https://github.com/apache/datafusion/pull/19501#discussion_r2652864208
##########
datafusion/functions-aggregate/src/min_max.rs:
##########
@@ -437,7 +438,7 @@ impl Accumulator for SlidingMaxAccumulator {
true
}
- fn size(&self) -> usize {
+ fn size(&self, _pool: Option<&dyn MemoryPool>) -> usize {
size_of_val(self) - size_of_val(&self.max) + self.max.size()
}
}
Review Comment:
MinAccumulator and MaxAccumulator could also use the pool since they hold a ScalarValue, which could hold an Arrow Array -> shared buffers. I didn’t do it to avoid adding more changes in this PR.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use arrow pool to fix memory over accounting in aggregations [datafusion]
Posted by "gabotechs (via GitHub)" <gi...@apache.org>.
gabotechs commented on code in PR #19501:
URL: https://github.com/apache/datafusion/pull/19501#discussion_r2655060697
##########
datafusion/functions-aggregate/src/array_agg.rs:
##########
@@ -518,13 +518,29 @@ impl Accumulator for DistinctArrayAggAccumulator {
Ok(ScalarValue::List(arr))
}
- fn size(&self) -> usize {
- size_of_val(self) + ScalarValue::size_of_hashset(&self.values)
- - size_of_val(&self.values)
+ fn size(&self, pool: Option<&dyn MemoryPool>) -> usize {
+ let mut total = size_of_val(self)
+ + size_of_val(&self.values)
+ + (size_of::<ScalarValue>() * self.values.capacity())
+ self.datatype.size()
- size_of_val(&self.datatype)
- size_of_val(&self.sort_options)
- + size_of::<Option<SortOptions>>()
+ + size_of::<Option<SortOptions>>();
+
+ for scalar in &self.values {
+ if let Some(array) = scalar.get_array_ref() {
+ total += size_of::<Arc<dyn Array>>();
+ if let Some(pool) = pool {
+ claim_buffers_recursive(&array.to_data(), pool);
Review Comment:
I see that if a `&dyn MemoryPool` is passed, then the `array` size does not compute towards the `total` size, and it's instead claimed in the `&dyn MemoryPool` instead.
Imagine this scenario:
- The underlying `array` is huge
- A `dyn MemoryPool` is passed, so the `array` size does not compute towards the `total_size`, it's just claimed in the Arrow Buffer memory pool
- In [GroupedHashAggregateStream::update_memory_reservation](https://github.com/apache/datafusion/blob/main/datafusion/physical-plan/src/aggregates/row_hash.rs#L1064-L1064), the `total_size` is very small, as the `array` size did not compute towards it.
- When calling `reservation.try_resize()` with the small `total_size`, the reservation succeeds
Isn't this a problematic scenario?
##########
datafusion/functions-aggregate/src/first_last.rs:
##########
@@ -1524,11 +1576,36 @@ impl Accumulator for LastValueAccumulator {
Ok(self.last.clone())
}
- fn size(&self) -> usize {
- size_of_val(self) - size_of_val(&self.last)
- + self.last.size()
- + ScalarValue::size_of_vec(&self.orderings)
- - size_of_val(&self.orderings)
+ fn size(&self, pool: Option<&dyn MemoryPool>) -> usize {
+ let mut total =
+ size_of_val(self) - size_of_val(&self.last) - size_of_val(&self.orderings)
+ + size_of::<ScalarValue>() * self.orderings.capacity();
+
+ if let Some(array) = self.last.get_array_ref() {
+ total += size_of::<Arc<dyn Array>>();
+ if let Some(pool) = pool {
+ claim_buffers_recursive(&array.to_data(), pool);
+ } else {
+ total += self.last.size() - size_of_val(&self.last);
+ }
+ } else {
+ total += self.last.size();
+ }
+
+ for scalar in &self.orderings {
+ if let Some(array) = scalar.get_array_ref() {
+ total += size_of::<Arc<dyn Array>>();
+ if let Some(pool) = pool {
+ claim_buffers_recursive(&array.to_data(), pool);
+ } else {
+ total += scalar.size() - size_of_val(scalar);
+ }
+ } else {
+ total += scalar.size() - size_of_val(scalar);
+ }
Review Comment:
This pattern seems to be repeated several times across the project. Maybe a helper could be useful?
##########
datafusion/functions-aggregate/src/array_agg.rs:
##########
@@ -390,27 +393,25 @@ impl Accumulator for ArrayAggAccumulator {
Ok(SingleRowListArrayBuilder::new(concated_array).build_list_scalar())
}
- fn size(&self) -> usize {
- size_of_val(self)
+ fn size(&self, pool: Option<&dyn MemoryPool>) -> usize {
+ let mut total = size_of_val(self)
+ (size_of::<ArrayRef>() * self.values.capacity())
- + self
+ + self.datatype.size()
+ - size_of_val(&self.datatype);
+
+ if let Some(pool) = pool {
+ for arr in &self.values {
+ claim_buffers_recursive(&arr.to_data(), pool);
Review Comment:
I imagine that `.size()` here can be called an arbitrary amount of times. What would happen with `claim_buffers_recursive` if this is called a lot of times?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use arrow pool to fix memory over accounting aggregations [datafusion]
Posted by "LiaCastaneda (via GitHub)" <gi...@apache.org>.
LiaCastaneda commented on PR #19501:
URL: https://github.com/apache/datafusion/pull/19501#issuecomment-3698920288
This is probably not the perfect solution -- but maybe a starting point? We've seen over accounting problems in topk as well, so maybe the pool could also be integrated there? I'm open to suggestions :)
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use arrow pool to fix memory over accounting in aggregations [datafusion]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] closed pull request #19501: Use arrow pool to fix memory over accounting in aggregations
URL: https://github.com/apache/datafusion/pull/19501
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use arrow pool to fix memory over accounting in aggregations [datafusion]
Posted by "LiaCastaneda (via GitHub)" <gi...@apache.org>.
LiaCastaneda commented on code in PR #19501:
URL: https://github.com/apache/datafusion/pull/19501#discussion_r2655395063
##########
datafusion/functions-aggregate/src/first_last.rs:
##########
@@ -1524,11 +1576,36 @@ impl Accumulator for LastValueAccumulator {
Ok(self.last.clone())
}
- fn size(&self) -> usize {
- size_of_val(self) - size_of_val(&self.last)
- + self.last.size()
- + ScalarValue::size_of_vec(&self.orderings)
- - size_of_val(&self.orderings)
+ fn size(&self, pool: Option<&dyn MemoryPool>) -> usize {
+ let mut total =
+ size_of_val(self) - size_of_val(&self.last) - size_of_val(&self.orderings)
+ + size_of::<ScalarValue>() * self.orderings.capacity();
+
+ if let Some(array) = self.last.get_array_ref() {
+ total += size_of::<Arc<dyn Array>>();
+ if let Some(pool) = pool {
+ claim_buffers_recursive(&array.to_data(), pool);
+ } else {
+ total += self.last.size() - size_of_val(&self.last);
+ }
+ } else {
+ total += self.last.size();
+ }
+
+ for scalar in &self.orderings {
+ if let Some(array) = scalar.get_array_ref() {
+ total += size_of::<Arc<dyn Array>>();
+ if let Some(pool) = pool {
+ claim_buffers_recursive(&array.to_data(), pool);
+ } else {
+ total += scalar.size() - size_of_val(scalar);
+ }
+ } else {
+ total += scalar.size() - size_of_val(scalar);
+ }
Review Comment:
yep, good idea
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use arrow pool to fix memory over accounting in aggregations [datafusion]
Posted by "LiaCastaneda (via GitHub)" <gi...@apache.org>.
LiaCastaneda commented on code in PR #19501:
URL: https://github.com/apache/datafusion/pull/19501#discussion_r2655376608
##########
datafusion/functions-aggregate/src/array_agg.rs:
##########
@@ -518,13 +518,29 @@ impl Accumulator for DistinctArrayAggAccumulator {
Ok(ScalarValue::List(arr))
}
- fn size(&self) -> usize {
- size_of_val(self) + ScalarValue::size_of_hashset(&self.values)
- - size_of_val(&self.values)
+ fn size(&self, pool: Option<&dyn MemoryPool>) -> usize {
+ let mut total = size_of_val(self)
+ + size_of_val(&self.values)
+ + (size_of::<ScalarValue>() * self.values.capacity())
+ self.datatype.size()
- size_of_val(&self.datatype)
- size_of_val(&self.sort_options)
- + size_of::<Option<SortOptions>>()
+ + size_of::<Option<SortOptions>>();
+
+ for scalar in &self.values {
+ if let Some(array) = scalar.get_array_ref() {
+ total += size_of::<Arc<dyn Array>>();
+ if let Some(pool) = pool {
+ claim_buffers_recursive(&array.to_data(), pool);
Review Comment:
Yes, that's why after calling size() we add arrow_pool.used() to account for the buffer memory (I forgot to add it). I also considered calling arrow_pool.used() directly inside each accumulator's size() function so the caller doesn't have to remember to do it. However, that would cause over-accounting in scenarios like `update_memory_reservation()` where we sum size() across multiple accumulators:
```
let total_size = self.group_values.size()
+ self.group_ordering.size()
+ self.current_group_indices.capacity() * size_of::<usize>()
+ self
.accumulators
.iter()
.map(|x| x.size(Some(&self.arrow_pool)))
.sum::<usize>() <--- If each size() returned arrow_pool.used() we would still be over counting the pool
```
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use arrow pool to fix memory over accounting in aggregations [datafusion]
Posted by "LiaCastaneda (via GitHub)" <gi...@apache.org>.
LiaCastaneda commented on PR #19501:
URL: https://github.com/apache/datafusion/pull/19501#issuecomment-3702136469
Yes, I completely agree with you. I created https://github.com/apache/arrow-rs/issues/8938 to raise this issue on the Arrow side. In most other Arrow implementations, they carry around a context object that contains the memory pool, so newly created arrays are immediately accounted for in the pool. In that case, DataFusion wouldn't have to do anything other than use the Arrow memory pool instead of the DataFusion pool. However, I'm aware this would require a considerable amount of effort on the Arrow side, and I'm not sure what the community there thinks of this idea.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org