parthchandra commented on code in PR #1050:
URL: https://github.com/apache/datafusion-comet/pull/1050#discussion_r1828407319
##########
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:
Does this need to be thread safe? Arc::strong_count is thread safe but the
reference count may change before the method returns
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
