Skip to content

Commit

Permalink
destination offset
Browse files Browse the repository at this point in the history
  • Loading branch information
paleolimbot committed May 1, 2024
1 parent 846a189 commit 81c76d3
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 3 deletions.
12 changes: 9 additions & 3 deletions python/src/nanoarrow/_lib.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -1815,7 +1815,7 @@ cdef class CBufferView:
else:
return self._iter_dispatch(offset, length)

def unpack_bits_into(self, dest, offset=0, length=None):
def unpack_bits_into(self, dest, offset=0, length=None, dest_offset=0):
if self._data_type != NANOARROW_TYPE_BOOL:
raise ValueError("Can't unpack non-boolean buffer")

Expand All @@ -1834,14 +1834,20 @@ cdef class CBufferView:
PyBuffer_Release(&buffer)
raise ValueError("Destination buffer has itemsize != 1")

if buffer.len < length:
if buffer.len < (dest_offset + length):
buffer_len = buffer.len
PyBuffer_Release(&buffer)
raise IndexError(
f"Can't unpack {length} elements into buffer of size {buffer_len}"
)

ArrowBitsUnpackInt8(self._ptr.data.as_uint8, offset, length, <int8_t*>buffer.buf)
ArrowBitsUnpackInt8(
self._ptr.data.as_uint8,
offset,
length,
&(<int8_t*>buffer.buf)[dest_offset]
)

PyBuffer_Release(&buffer)

def unpack_bits(self, offset=0, length=None):
Expand Down
5 changes: 5 additions & 0 deletions python/tests/test_c_buffer_view.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,11 @@ def test_buffer_view_bool_unpack():
assert len(unpacked_some) == 4
assert list(unpacked_some) == [0, 0, 1, 0]

# Check with non-zero destination offset
out = bytearray([255] * 10)
view.unpack_bits_into(out, dest_offset=2)
assert list(out) == [255, 255, 1, 0, 0, 1, 0, 0, 0, 0]

# Check errors from requesting out-of-bounds slices
msg = "do not describe a valid slice"
with pytest.raises(IndexError, match=msg):
Expand Down

0 comments on commit 81c76d3

Please sign in to comment.