jorisvandenbossche commented on code in PR #378:
URL: https://github.com/apache/arrow-nanoarrow/pull/378#discussion_r1481452416
##########
python/src/nanoarrow/_lib.pyx:
##########
@@ -176,6 +178,188 @@ cdef object alloc_c_array_shallow_copy(object base, const
ArrowArray* c_array) n
return array_capsule
+cdef void pycapsule_buffer_deleter(object stream_capsule) noexcept:
+ cdef ArrowBuffer* buffer = <ArrowBuffer*>PyCapsule_GetPointer(
+ stream_capsule, 'nanoarrow_buffer'
+ )
+
+ ArrowBufferReset(buffer)
+ ArrowFree(buffer)
+
+
+cdef object alloc_c_buffer(ArrowBuffer** c_buffer) noexcept:
+ c_buffer[0] = <ArrowBuffer*> ArrowMalloc(sizeof(ArrowBuffer))
+ ArrowBufferInit(c_buffer[0])
+ return PyCapsule_New(c_buffer[0], 'nanoarrow_buffer',
&pycapsule_buffer_deleter)
+
+cdef void c_deallocate_pybuffer(ArrowBufferAllocator* allocator, uint8_t* ptr,
int64_t size) noexcept with gil:
+ cdef Py_buffer* buffer = <Py_buffer*>allocator.private_data
+ PyBuffer_Release(buffer)
+ ArrowFree(buffer)
+
+
+cdef ArrowBufferAllocator c_pybuffer_deallocator(Py_buffer* buffer):
+ # This should probably be changed in nanoarrow C; however, currently, the
deallocator
+ # won't get called if buffer.buf is NULL.
+ if buffer.buf == NULL:
+ PyBuffer_Release(buffer)
+ return ArrowBufferAllocatorDefault()
+
+ cdef Py_buffer* allocator_private =
<Py_buffer*>ArrowMalloc(sizeof(Py_buffer))
+ if allocator_private == NULL:
+ PyBuffer_Release(buffer)
+ raise MemoryError()
+
+ memcpy(allocator_private, buffer, sizeof(Py_buffer))
+ return
ArrowBufferDeallocator(<ArrowBufferDeallocatorCallback>&c_deallocate_pybuffer,
allocator_private)
+
+
+cdef c_arrow_type_from_format(format):
+ # PyBuffer_SizeFromFormat() was added in Python 3.9 (potentially faster)
+ item_size = calcsize(format)
+
+ # Don't allow non-native endian values
+ if sys_byteorder == "little" and (">" in format or "!" in format):
+ raise ValueError(f"Can't convert format '{format}' to Arrow type")
+ elif sys_byteorder == "big" and "<" in format:
+ raise ValueError(f"Can't convert format '{format}' to Arrow type")
+
+ # Strip system endian specifiers
+ format = format.strip("=@")
+
+ if format == "c":
+ return 0, NANOARROW_TYPE_STRING
Review Comment:
Would binary be more correct here?
##########
python/tests/test_c_lib.py:
##########
@@ -0,0 +1,489 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import struct
+import sys
+
+import pytest
+from nanoarrow._lib import NanoarrowException
+from nanoarrow.c_lib import (
+ CArrayBuilder,
+ CBuffer,
+ CBufferBuilder,
+ c_array_empty,
+ c_array_from_buffers,
+ c_array_from_pybuffer,
+ c_buffer,
+ c_buffer_from_iterable,
+)
+
+import nanoarrow as na
+
+
+def test_buffer_invalid():
+ invalid = CBuffer()
+
+ with pytest.raises(RuntimeError, match="CBuffer is not valid"):
+ invalid._addr()
+ with pytest.raises(RuntimeError, match="CBuffer is not valid"):
+ invalid.size_bytes
+ with pytest.raises(RuntimeError, match="CBuffer is not valid"):
+ invalid.capacity_bytes
+ with pytest.raises(RuntimeError, match="CBuffer is not valid"):
+ invalid.data
+
+ assert repr(invalid) == "CBuffer(<invalid>)"
+
+
+def test_c_buffer_constructor():
+ invalid = CBuffer()
+ assert c_buffer(invalid) is invalid
+
+ buffer = c_buffer(b"1234")
+ assert isinstance(buffer, CBuffer)
+ assert bytes(buffer.data) == b"1234"
Review Comment:
The CBuffer itself could also support the buffer protocol?
##########
python/src/nanoarrow/c_lib.py:
##########
@@ -257,7 +465,74 @@ def c_array_view(obj, requested_schema=None) -> CArrayView:
return CArrayView.from_cpu_array(c_array(obj, requested_schema))
-def allocate_c_schema():
+def c_buffer(obj) -> CBuffer:
+ """Owning, read-only ArrowBuffer wrapper
+
+ Wraps obj in nanoarrow's owning buffer structure, the ArrowBuffer,
Review Comment:
It's "owning" in the sense that it will take care of releasing the reference
when it goes out of scope. But it's still the original python objects that owns
the actual data?
##########
python/src/nanoarrow/_lib.pyx:
##########
@@ -176,6 +178,188 @@ cdef object alloc_c_array_shallow_copy(object base, const
ArrowArray* c_array) n
return array_capsule
+cdef void pycapsule_buffer_deleter(object stream_capsule) noexcept:
+ cdef ArrowBuffer* buffer = <ArrowBuffer*>PyCapsule_GetPointer(
+ stream_capsule, 'nanoarrow_buffer'
+ )
+
+ ArrowBufferReset(buffer)
+ ArrowFree(buffer)
+
+
+cdef object alloc_c_buffer(ArrowBuffer** c_buffer) noexcept:
+ c_buffer[0] = <ArrowBuffer*> ArrowMalloc(sizeof(ArrowBuffer))
+ ArrowBufferInit(c_buffer[0])
+ return PyCapsule_New(c_buffer[0], 'nanoarrow_buffer',
&pycapsule_buffer_deleter)
+
+cdef void c_deallocate_pybuffer(ArrowBufferAllocator* allocator, uint8_t* ptr,
int64_t size) noexcept with gil:
+ cdef Py_buffer* buffer = <Py_buffer*>allocator.private_data
+ PyBuffer_Release(buffer)
+ ArrowFree(buffer)
+
+
+cdef ArrowBufferAllocator c_pybuffer_deallocator(Py_buffer* buffer):
+ # This should probably be changed in nanoarrow C; however, currently, the
deallocator
+ # won't get called if buffer.buf is NULL.
+ if buffer.buf == NULL:
+ PyBuffer_Release(buffer)
+ return ArrowBufferAllocatorDefault()
+
+ cdef Py_buffer* allocator_private =
<Py_buffer*>ArrowMalloc(sizeof(Py_buffer))
+ if allocator_private == NULL:
+ PyBuffer_Release(buffer)
+ raise MemoryError()
+
+ memcpy(allocator_private, buffer, sizeof(Py_buffer))
+ return
ArrowBufferDeallocator(<ArrowBufferDeallocatorCallback>&c_deallocate_pybuffer,
allocator_private)
+
+
+cdef c_arrow_type_from_format(format):
+ # PyBuffer_SizeFromFormat() was added in Python 3.9 (potentially faster)
+ item_size = calcsize(format)
+
+ # Don't allow non-native endian values
+ if sys_byteorder == "little" and (">" in format or "!" in format):
+ raise ValueError(f"Can't convert format '{format}' to Arrow type")
+ elif sys_byteorder == "big" and "<" in format:
+ raise ValueError(f"Can't convert format '{format}' to Arrow type")
+
+ # Strip system endian specifiers
+ format = format.strip("=@")
+
+ if format == "c":
+ return 0, NANOARROW_TYPE_STRING
+ elif format == "e":
+ return item_size, NANOARROW_TYPE_HALF_FLOAT
+ elif format == "f":
+ return item_size, NANOARROW_TYPE_FLOAT
+ elif format == "d":
+ return item_size, NANOARROW_TYPE_DOUBLE
+
+ # Check for signed integers
+ if format in ("b", "?", "h", "i", "l", "q", "n"):
+ if item_size == 1:
+ return item_size, NANOARROW_TYPE_INT8
+ elif item_size == 2:
+ return item_size, NANOARROW_TYPE_INT16
+ elif item_size == 4:
+ return item_size, NANOARROW_TYPE_INT32
+ elif item_size == 8:
+ return item_size, NANOARROW_TYPE_INT64
+
+ # Check for unsinged integers
+ if format in ("B", "H", "I", "L", "Q", "N"):
+ if item_size == 1:
+ return item_size, NANOARROW_TYPE_UINT8
+ elif item_size == 2:
+ return item_size, NANOARROW_TYPE_UINT16
+ elif item_size == 4:
+ return item_size, NANOARROW_TYPE_UINT32
+ elif item_size == 8:
+ return item_size, NANOARROW_TYPE_UINT64
+
+ # If all else fails, return opaque fixed-size binary
+ return item_size, NANOARROW_TYPE_BINARY
+
+
+cdef int c_format_from_arrow_type(ArrowType type_id, int element_size_bits,
size_t out_size, char* out):
+ if type_id in (NANOARROW_TYPE_BINARY, NANOARROW_TYPE_FIXED_SIZE_BINARY)
and element_size_bits > 0:
+ snprintf(out, out_size, "%ds", <int>(element_size_bits // 8))
+ return element_size_bits
+
+ cdef const char* format_const = ""
+ cdef int element_size_bits_calc = 0
+ if type_id == NANOARROW_TYPE_STRING:
+ format_const = "c"
+ element_size_bits_calc = 0
+ elif type_id == NANOARROW_TYPE_BINARY:
+ format_const = "B"
+ element_size_bits_calc = 0
+ elif type_id == NANOARROW_TYPE_BOOL:
+ # Bitmaps export as unspecified binary
+ format_const = "B"
+ element_size_bits_calc = 1
+ elif type_id == NANOARROW_TYPE_INT8:
+ format_const = "b"
+ element_size_bits_calc = 8
+ elif type_id == NANOARROW_TYPE_UINT8:
+ format_const = "B"
+ element_size_bits_calc = 8
+ elif type_id == NANOARROW_TYPE_INT16:
+ format_const = "h"
+ element_size_bits_calc = 16
+ elif type_id == NANOARROW_TYPE_UINT16:
+ format_const = "H"
+ element_size_bits_calc = 16
+ elif type_id in (NANOARROW_TYPE_INT32, NANOARROW_TYPE_INTERVAL_MONTHS):
+ format_const = "i"
+ element_size_bits_calc = 32
+ elif type_id == NANOARROW_TYPE_UINT32:
+ format_const = "I"
+ element_size_bits_calc = 32
+ elif type_id == NANOARROW_TYPE_INT64:
+ format_const = "q"
+ element_size_bits_calc = 64
+ elif type_id == NANOARROW_TYPE_UINT64:
+ format_const = "Q"
+ element_size_bits_calc = 64
+ elif type_id == NANOARROW_TYPE_HALF_FLOAT:
+ format_const = "e"
+ element_size_bits_calc = 16
+ elif type_id == NANOARROW_TYPE_FLOAT:
+ format_const = "f"
+ element_size_bits_calc = 32
+ elif type_id == NANOARROW_TYPE_DOUBLE:
+ format_const = "d"
+ element_size_bits_calc = 64
+ elif type_id == NANOARROW_TYPE_INTERVAL_DAY_TIME:
+ format_const = "ii"
+ element_size_bits_calc = 64
+ elif type_id == NANOARROW_TYPE_INTERVAL_MONTH_DAY_NANO:
+ format_const = "iiq"
+ element_size_bits_calc = 128
+ elif type_id == NANOARROW_TYPE_DECIMAL128:
+ format_const = "16s"
+ element_size_bits_calc = 128
+ elif type_id == NANOARROW_TYPE_DECIMAL256:
+ format_const = "32s"
+ element_size_bits_calc = 256
+ else:
+ raise ValueError(f"Unsupported Arrow type_id for format conversion:
{type_id}")
+
+ snprintf(out, out_size, "%s", format_const)
+ return element_size_bits_calc
+
+
+cdef object c_buffer_set_pybuffer(object obj, ArrowBuffer** c_buffer):
+ ArrowBufferReset(c_buffer[0])
+
+ cdef Py_buffer buffer
+ cdef int rc = PyObject_GetBuffer(obj, &buffer, PyBUF_FORMAT |
PyBUF_ANY_CONTIGUOUS)
+ if rc != 0:
+ raise BufferError()
+
+ # Parse the buffer's format string to get the ArrowType and element size
+ try:
+ if buffer.format == NULL:
+ format = "B"
+ else:
+ format = buffer.format.decode("UTF-8")
+ except Exception as e:
+ PyBuffer_Release(&buffer)
+ raise e
+
+ # Transfers ownership of buffer to c_buffer, whose finalizer will be
called by
+ # the capsule when the capsule is deleted or garbage collected
+ c_buffer[0].data = <uint8_t*>buffer.buf
+ c_buffer[0].size_bytes = <int64_t>buffer.len
+ c_buffer[0].capacity_bytes = 0
+ c_buffer[0].allocator = c_pybuffer_deallocator(&buffer)
Review Comment:
Should this function somehow increase the reference of `buffer`, or how does
this ensure the Python buffer is kept alive?
##########
python/src/nanoarrow/c_lib.py:
##########
@@ -125,10 +138,205 @@ def c_array(obj=None, requested_schema=None) -> CArray:
out = CArray.allocate(CSchema.allocate())
obj._export_to_c(out._addr(), out.schema._addr())
return out
- else:
+
+ # Try buffer protocol (e.g., numpy arrays)
+ try:
+ return c_array_from_pybuffer(obj)
+ except Exception as e:
raise TypeError(
f"Can't convert object of type {type(obj).__name__} to
nanoarrow.c_array"
+ ) from e
+
+
+def c_array_from_pybuffer(obj) -> CArray:
+ """Create an ArrowArray wrapper from the Python buffer protocol
+
+ Invokes the Python buffer protocol to wrap the buffer represented by obj
+ if possible.
+
+ Examples
+ --------
+
+ >>> import nanoarrow as na
+ >>> from nanoarrow.c_lib import c_array_from_pybuffer
+ >>> na.c_array_view(c_array_from_pybuffer(b"1234"))
+ <nanoarrow.c_lib.CArrayView>
+ - storage_type: 'uint8'
+ - length: 4
+ - offset: 0
+ - null_count: 0
+ - buffers[2]:
+ - validity <bool[0 b] >
+ - data <uint8[4 b] 49 50 51 52>
+ - dictionary: NULL
+ - children[0]:
+ """
+
+ buffer = CBuffer().set_pybuffer(obj)
+ view = buffer.data
Review Comment:
The `.data` is a strange name for the view? (because it's actually _not_ the
data but a view on it?)
--
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]
