viirya commented on code in PR #1050:
URL: https://github.com/apache/datafusion-comet/pull/1050#discussion_r1828439388
##########
native/core/src/common/buffer.rs:
##########
@@ -163,11 +168,28 @@ impl CometBuffer {
/// because of the iterator-style pattern, the content of the original
mutable buffer will only
/// be updated once upstream operators fully consumed the previous output
batch. For breaking
/// operators, they are responsible for copying content out of the buffers.
- pub unsafe fn to_arrow(&self) -> ArrowBuffer {
+ pub unsafe fn to_arrow(&self) -> Result<ArrowBuffer, ExecutionError> {
let ptr = NonNull::new_unchecked(self.data.as_ptr());
- // Uses a dummy `Arc::new(0)` as `Allocation` to ensure the memory
region pointed by
- // `ptr` won't be freed when the returned `ArrowBuffer` goes out of
scope.
- ArrowBuffer::from_custom_allocation(ptr, self.len, Arc::new(0))
+ self.check_reference()?;
+ Ok(ArrowBuffer::from_custom_allocation(
+ ptr,
+ self.len,
+ Arc::<CometBufferAllocation>::clone(&self.allocation),
+ ))
+ }
+
+ /// Checks if this buffer is exclusively owned by Comet. If not, an error
is returned.
+ /// We run this check when we want to update the buffer. If the buffer is
also shared by
+ /// other components, e.g. one DataFusion operator stores the buffer,
Comet cannot safely
+ /// modify the buffer.
+ pub fn check_reference(&self) -> Result<(), ExecutionError> {
Review Comment:
This is called after filling new values into the buffer. The reference count
should not be changed at the point. We also don't access scan iterator in
multiple threads.
--
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-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
