jorisvandenbossche commented on code in PR #378:
URL: https://github.com/apache/arrow-nanoarrow/pull/378#discussion_r1492211648
##########
python/src/nanoarrow/_lib.pyx:
##########
@@ -1177,11 +1462,435 @@ cdef class CBufferView:
buffer.strides = &self._strides
buffer.suboffsets = NULL
- def __releasebuffer__(self, Py_buffer *buffer):
+ cdef _do_releasebuffer(self, Py_buffer* buffer):
pass
def __repr__(self):
- return f"<nanoarrow.c_lib.CBufferView>\n
{_lib_utils.buffer_view_repr(self)[1:]}"
+ return f"CBufferView({_lib_utils.buffer_view_repr(self)})"
+
+
+cdef class CBuffer:
+ """Wrapper around readable owned buffer content
+
+ Like the CBufferView, the CBuffer represents readable buffer content;
however,
+ unlike the CBufferView, the CBuffer always represents a valid ArrowBuffer
C object.
+ """
+ cdef object _base
+ cdef ArrowBuffer* _ptr
+ cdef ArrowType _data_type
+ cdef int _element_size_bits
+ cdef char _format[32]
+ cdef CDevice _device
+ cdef CBufferView _view
+ cdef int _get_buffer_count
+
+ def __cinit__(self):
+ self._base = None
+ self._ptr = NULL
+ self._data_type = NANOARROW_TYPE_BINARY
+ self._element_size_bits = 0
+ self._device = CDEVICE_CPU
+ # Set initial format to "B" (Cython makes this hard)
+ self._format[0] = 66
+ self._format[1] = 0
+ self._get_buffer_count = 0
+ self._view = CBufferView(None, 0, 0, NANOARROW_TYPE_BINARY, 0,
self._device)
+
+ cdef _assert_valid(self):
+ if self._ptr == NULL:
+ raise RuntimeError("CBuffer is not valid")
+
+ cdef _assert_buffer_count_zero(self):
+ if self._get_buffer_count != 0:
+ raise RuntimeError(
+ f"CBuffer already open ({self._get_buffer_count} ",
+ f"references, {self._writable_get_buffer_count} writable)")
+
+ cdef _populate_view(self):
+ self._assert_valid()
+ self._assert_buffer_count_zero()
+ self._view = CBufferView(
+ self._base, <uintptr_t>self._ptr.data,
+ self._ptr.size_bytes, self._data_type, self._element_size_bits,
+ self._device
+ )
+
+ @staticmethod
+ def empty():
+ cdef CBuffer out = CBuffer()
+ out._base = alloc_c_buffer(&out._ptr)
+ return out
+
+ @staticmethod
+ def from_pybuffer(obj):
+ cdef CBuffer out = CBuffer()
+ out._base = alloc_c_buffer(&out._ptr)
+ out._set_format(c_buffer_set_pybuffer(obj, &out._ptr))
+ out._device = CDEVICE_CPU
+ out._populate_view()
+ return out
+
+ def _set_format(self, str format):
+ self._assert_buffer_count_zero()
+
+ element_size_bytes, data_type = c_arrow_type_from_format(format)
+ self._data_type = data_type
+ self._element_size_bits = element_size_bytes * 8
+ format_bytes = format.encode("UTF-8")
+ snprintf(self._format, sizeof(self._format), "%s", <const
char*>format_bytes)
+ self._populate_view()
+ return self
+
+ def _set_data_type(self, ArrowType type_id, int element_size_bits=0):
+ self._assert_buffer_count_zero()
+
+ self._element_size_bits = c_format_from_arrow_type(
+ type_id,
+ element_size_bits,
+ sizeof(self._format),
+ self._format
+ )
+ self._data_type = type_id
+
+ self._populate_view()
+ return self
+
+ def _addr(self):
+ self._assert_valid()
+ return <uintptr_t>self._ptr.data
+
+ @property
+ def _get_buffer_count(self):
+ return self._get_buffer_count
+
+ @property
+ def size_bytes(self):
+ self._assert_valid()
+ return self._ptr.size_bytes
+
+ @property
+ def data_type(self):
+ return ArrowTypeString(self._data_type).decode("UTF-8")
+
+ @property
+ def data_type_id(self):
+ return self._data_type
+
+ @property
+ def element_size_bits(self):
+ return self._element_size_bits
+
+ @property
+ def item_size(self):
+ self._assert_valid()
+ return self._view.item_size
+
+ @property
+ def format(self):
+ return self._format.decode("UTF-8")
+
+ def __len__(self):
+ self._assert_valid()
+ return len(self._view)
+
+ def __getitem__(self, k):
+ self._assert_valid()
+ return self._view[k]
+
+ def __iter__(self):
+ self._assert_valid()
+ return iter(self._view)
+
+ @property
+ def n_elements(self):
+ self._assert_valid()
+ return self._view.n_elements
+
+ def element(self, i):
+ self._assert_valid()
+ return self._view.element(i)
+
+ @property
+ def elements(self):
+ self._assert_valid()
+ return self._view.elements
+
+ def __getbuffer__(self, Py_buffer* buffer, int flags):
+ self._assert_valid()
+ self._view._do_getbuffer(buffer, flags)
+ self._get_buffer_count += 1
+
+ def __releasebuffer__(self, Py_buffer* buffer):
+ if self._get_buffer_count <= 0:
+ raise RuntimeError("CBuffer buffer reference count underflow
(releasebuffer)")
+
+ self._view._do_releasebuffer(buffer)
+ self._get_buffer_count -= 1
+
+ def __repr__(self):
+ if self._ptr == NULL:
+ return "CBuffer(<invalid>)"
+
+ return f"CBuffer({_lib_utils.buffer_view_repr(self._view)})"
+
+
+cdef class CBufferBuilder:
+ """Wrapper around writable owned buffer CPU content"""
+ cdef CBuffer _buffer
+
+ def __cinit__(self):
+ self._buffer = CBuffer.empty()
+
+ @property
+ def size_bytes(self):
+ return self._buffer.size_bytes
+
+ @property
+ def capacity_bytes(self):
+ return self._buffer._ptr.capacity_bytes
+
+ def set_data_type(self, ArrowType type_id, int element_size_bits=0):
+ self._buffer._set_data_type(type_id, element_size_bits)
+ return self
+
+ def reserve_bytes(self, int64_t additional_bytes):
+ cdef int code = ArrowBufferReserve(self._buffer._ptr, additional_bytes)
+ Error.raise_error_not_ok("ArrowBufferReserve()", code)
+ return self
+
+ def write(self, content):
+ cdef Py_buffer buffer
+ cdef int64_t out
+ PyObject_GetBuffer(content, &buffer, PyBUF_ANY_CONTIGUOUS)
+
+ cdef int code = ArrowBufferReserve(self._buffer._ptr, buffer.len)
+ if code != NANOARROW_OK:
+ PyBuffer_Release(&buffer)
+ Error.raise_error("ArrowBufferReserve()", code)
+
+ code = PyBuffer_ToContiguous(
+ self._buffer._ptr.data + self._buffer._ptr.size_bytes,
+ &buffer,
+ buffer.len,
+ # 'C' (not sure how to pass a character literal here)
+ 43
+ )
+ out = buffer.len
+ PyBuffer_Release(&buffer)
+ Error.raise_error_not_ok("PyBuffer_ToContiguous()", code)
+
+ self._buffer._ptr.size_bytes += out
+ return out
+
+ def write_elements(self, obj):
+ if self._buffer._data_type == NANOARROW_TYPE_BOOL:
+ return self._write_bits(obj)
+
+ cdef int64_t n_values = 0
+ struct_obj = Struct(self._buffer._format)
+ pack = struct_obj.pack
+ write = self.write
+
+ if self._buffer._data_type in (NANOARROW_TYPE_INTERVAL_DAY_TIME,
+ NANOARROW_TYPE_INTERVAL_MONTH_DAY_NANO):
+ for item in obj:
+ n_values += 1
+ write(pack(*item))
+ else:
+ for item in obj:
+ n_values += 1
+ write(pack(item))
+
+ return n_values
+
+ cdef _write_bits(self, obj):
+ if self._buffer._ptr.size_bytes != 0:
+ raise NotImplementedError("Append to bitmap that has already been
appended to")
+
+ cdef char buffer_item = 0
+ cdef int buffer_item_i = 0
+ cdef int code
+ cdef int64_t n_values = 0
+ for item in obj:
+ n_values += 1
+ if item:
+ buffer_item |= (<char>1 << buffer_item_i)
+
+ buffer_item_i += 1
+ if buffer_item_i == 8:
+ code = ArrowBufferAppendInt8(self._buffer._ptr, buffer_item)
+ Error.raise_error_not_ok("ArrowBufferAppendInt8()", code)
+ buffer_item = 0
+ buffer_item_i = 0
+
+ if buffer_item_i != 0:
+ code = ArrowBufferAppendInt8(self._buffer._ptr, buffer_item)
+ Error.raise_error_not_ok("ArrowBufferAppendInt8()", code)
+
+ return n_values
+
+ def finish(self):
+ cdef CBuffer out = self._buffer
+ out._populate_view()
+
+ self._buffer = CBuffer.empty()
+ return out
+
+
+cdef class CArrayBuilder:
+ cdef CArray c_array
+ cdef ArrowArray* _ptr
+ cdef bint _can_validate
+
+ def __cinit__(self, CArray array):
+ self.c_array = array
+ self._ptr = array._ptr
+ self._can_validate = True
+
+ @staticmethod
+ def allocate():
+ return CArrayBuilder(CArray.allocate(CSchema.allocate()))
+
+ def init_from_type(self, int type_id):
+ if self._ptr.release != NULL:
+ raise RuntimeError("CArrayBuilder is already initialized")
+
+ cdef int code = ArrowArrayInitFromType(self._ptr, <ArrowType>type_id)
+ Error.raise_error_not_ok("ArrowArrayInitFromType()", code)
+
+ code = ArrowSchemaInitFromType(self.c_array._schema._ptr,
<ArrowType>type_id)
+ Error.raise_error_not_ok("ArrowSchemaInitFromType()", code)
+
+ return self
+
+ def init_from_schema(self, CSchema schema):
+ if self._ptr.release != NULL:
+ raise RuntimeError("CArrayBuilder is already initialized")
+
+ cdef Error error = Error()
+ cdef int code = ArrowArrayInitFromSchema(self._ptr, schema._ptr,
&error.c_error)
+ error.raise_message_not_ok("ArrowArrayInitFromType()", code)
+
+ self.c_array._schema = schema
+ return self
+
+ def start_appending(self):
+ cdef int code = ArrowArrayStartAppending(self._ptr)
+ Error.raise_error_not_ok("ArrowArrayStartAppending()", code)
+ return self
+
+ def set_offset(self, int64_t offset):
+ self.c_array._assert_valid()
+ self._ptr.offset = offset
+ return self
+
+ def set_length(self, int64_t length):
+ self.c_array._assert_valid()
+ self._ptr.length = length
+ return self
+
+ def set_null_count(self, int64_t null_count):
+ self.c_array._assert_valid()
+ self._ptr.null_count = null_count
+ return self
+
+ def resolve_null_count(self):
+ self.c_array._assert_valid()
+
+ # This doesn't apply to unions. We currently don't have a schema view
+ # or array view we can use to query the type ID, so just use the format
+ # string for now.
+ format = self.c_array.schema.format
+ if format.startswith("+us:") or format.startswith("+ud:"):
+ return self
+
+ # Don't overwrite an explicit null count
+ if self._ptr.null_count != -1:
+ return self
+
+ cdef ArrowBuffer* validity_buffer = ArrowArrayBuffer(self._ptr, 0)
+ if validity_buffer.size_bytes == 0:
+ self._ptr.null_count = 0
+ return self
+
+ # From _ArrowBytesForBits(), which is not included in nanoarrow_c.pxd
+ # because it's an internal inline function.
+ cdef int64_t bits = self._ptr.offset + self._ptr.length
+ cdef int64_t bytes_required = (bits >> 3) + ((bits & 7) != 0)
+
+ if validity_buffer.size_bytes < bytes_required:
+ raise ValueError(
+ f"Expected validity bitmap >= {bytes_required} bytes "
+ f"but got validity bitmap with {validity_buffer.size_bytes}
bytes"
+ )
+
+ cdef int64_t count = ArrowBitCountSet(
+ validity_buffer.data,
+ self._ptr.offset,
+ self._ptr.length
+ )
+ self._ptr.null_count = self._ptr.length - count
+ return self
+
+ def set_buffer(self, int64_t i, CBuffer buffer, move=False):
Review Comment:
I see below that the docstring of the public function `c_array_from_buffers`
has an explanation: "Use ``True`` to move ownership of any input buffers or
children to the output array.".
I would repeat that here for when looking at the code here. And it's also
not super clear to me what the consequence is exactly of doing that? Or, why
would you want to set it to True?
--
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]
