[PoC] Add API for tracking distinct buffers in MemoryPool by reference count

[Dandandan]

Which issue does this PR close?

• Closes #.

Rationale for this change

Currently, memory is tracked based on self reporting of bytes by each individual consumer. While that often works, it will overreport the amount of memory

• Whenever the child operator produces data based on calling slice on the RecordBatch / arrays. (this over-reports a lot for very large RecordBatches, i.e. those coming out of aggregates).
• Whenever child arrays are re-used within the same RecordBatch (for example, two same arrays with a different name).

This PR proposes an API extension to MemoryPool and MemoryReservation that tracks Arc<dyn Array> based on their memory address.

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

[2010YOUY01]

Thank you, this solves the memory overcounting issue across batches. I got some suggestions/questions.

This new pool implementation might cause some issue for MemoryReservation API:
If we first count the memory usage of several arrays by grow_with_arrays() interface, and then use the following functions to release (like reservation.resize(0)), memory used value will be subtracted, but the inner hash table won't be cleared

datafusion/datafusion/execution/src/memory_pool/mod.rs

Lines 308 to 314 in f35416e

	pub fn free(&mut self) -> usize {
	let size = self.size;
	if size != 0 {
	self.shrink(size)
	}
	size
	}

datafusion/datafusion/execution/src/memory_pool/mod.rs

Lines 345 to 351 in f35416e

	pub fn resize(&mut self, capacity: usize) {
	match capacity.cmp(&self.size) {
	Ordering::Greater => self.grow(capacity - self.size),
	Ordering::Less => self.shrink(self.size - capacity),
	_ => {}
	}
	}

This causes a potential inconsistent state, maybe we can use another semantics for the MemoryReservation API:

• Memory for arrays can only be managed with grow_with_arrays()/shrink_with_arrays()
• Other interfaces like shrink()/grow() will used for other memory usage

And the implementation will be like

#[derive(Debug)]
pub struct GreedyMemoryPoolWithTracking {
    pool_size: usize,
    used_others: AtomicUsize, // managed with non-array APIs
    
    used_array: AtomicUsize, // ref-counted with buffer addr stored in the hash table
    references: Mutex<HashMap<usize, usize>>,
}

[2010YOUY01]

Can we make the API take RecordBatch instead of arrays? Since inside df it's more common to passing batches around, and we can use a utility function to do batch -> [Array] for array usages.

[Dandandan]

Thats possible, let me have a look at whether we can use recordbatch always. My corcern was we might not always have a RecordBatch, but might have an Array. In that case conversion to recordbatch would be strange.

[Dandandan]

Thank you, this solves the memory overcounting issue across batches. I got some suggestions/questions.

This new pool implementation might cause some issue for MemoryReservation API: If we first count the memory usage of several arrays by grow_with_arrays() interface, and then use the following functions to release (like reservation.resize(0)), memory used value will be subtracted, but the inner hash table won't be cleared

datafusion/datafusion/execution/src/memory_pool/mod.rs

Lines 308 to 314 in f35416e

	pub fn free(&mut self) -> usize {
	let size = self.size;
	if size != 0 {
	self.shrink(size)
	}
	size
	}

datafusion/datafusion/execution/src/memory_pool/mod.rs

Lines 345 to 351 in f35416e

	pub fn resize(&mut self, capacity: usize) {
	match capacity.cmp(&self.size) {
	Ordering::Greater => self.grow(capacity - self.size),
	Ordering::Less => self.shrink(self.size - capacity),
	_ => {}
	}
	}

This causes a potential inconsistent state, maybe we can use another semantics for the MemoryReservation API:

• Memory for arrays can only be managed with grow_with_arrays()/shrink_with_arrays()
• Other interfaces like shrink()/grow() will used for other memory usage

And the implementation will be like

#[derive(Debug)]
pub struct GreedyMemoryPoolWithTracking {
    pool_size: usize,
    used_others: AtomicUsize, // managed with non-array APIs
    
    used_array: AtomicUsize, // ref-counted with buffer addr stored in the hash table
    references: Mutex<HashMap<usize, usize>>,
}

Hi

Thank you, this solves the memory overcounting issue across batches. I got some suggestions/questions.

This new pool implementation might cause some issue for MemoryReservation API: If we first count the memory usage of several arrays by grow_with_arrays() interface, and then use the following functions to release (like reservation.resize(0)), memory used value will be subtracted, but the inner hash table won't be cleared

datafusion/datafusion/execution/src/memory_pool/mod.rs

Lines 308 to 314 in f35416e

	pub fn free(&mut self) -> usize {
	let size = self.size;
	if size != 0 {
	self.shrink(size)
	}
	size
	}

datafusion/datafusion/execution/src/memory_pool/mod.rs

Lines 345 to 351 in f35416e

	pub fn resize(&mut self, capacity: usize) {
	match capacity.cmp(&self.size) {
	Ordering::Greater => self.grow(capacity - self.size),
	Ordering::Less => self.shrink(self.size - capacity),
	_ => {}
	}
	}

This causes a potential inconsistent state, maybe we can use another semantics for the MemoryReservation API:

• Memory for arrays can only be managed with grow_with_arrays()/shrink_with_arrays()
• Other interfaces like shrink()/grow() will used for other memory usage

And the implementation will be like

#[derive(Debug)]
pub struct GreedyMemoryPoolWithTracking {
    pool_size: usize,
    used_others: AtomicUsize, // managed with non-array APIs
    
    used_array: AtomicUsize, // ref-counted with buffer addr stored in the hash table
    references: Mutex<HashMap<usize, usize>>,
}

Thanks for the feedback!

Yes, tracking the size seperately makes sense, I'll change that!