Skip to content

Commit

Permalink
make sure get works
Browse files Browse the repository at this point in the history
  • Loading branch information
paleolimbot committed Sep 27, 2024
1 parent 4f582d9 commit 99e5173
Show file tree
Hide file tree
Showing 3 changed files with 52 additions and 9 deletions.
25 changes: 24 additions & 1 deletion python/src/nanoarrow/_array.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -22,10 +22,12 @@ from cpython.pycapsule cimport PyCapsule_GetPointer
from cpython.unicode cimport PyUnicode_AsUTF8AndSize
from cpython cimport (
Py_buffer,
PyObject_GetBuffer,
PyBuffer_Release,
PyBUF_ANY_CONTIGUOUS,
PyBUF_FORMAT,
PyBytes_FromStringAndSize,
PyObject_GetBuffer,
PyUnicode_FromStringAndSize,
)

from nanoarrow_c cimport (
Expand All @@ -43,6 +45,9 @@ from nanoarrow_c cimport (
ArrowArrayView,
ArrowArrayViewComputeNullCount,
ArrowArrayViewInitFromSchema,
ArrowArrayViewIsNull,
ArrowArrayViewGetStringUnsafe,
ArrowArrayViewGetBytesUnsafe,
ArrowArrayViewSetArray,
ArrowArrayViewSetArrayMinimal,
ArrowBitCountSet,
Expand Down Expand Up @@ -289,6 +294,24 @@ cdef class CArrayView:

return dictionary

def _iter_bytes(self, int64_t offset, int64_t length):
cdef ArrowBufferView item_view
for i in range(offset, length):
if ArrowArrayViewIsNull(self._ptr, i):
yield None
else:
item_view = ArrowArrayViewGetBytesUnsafe(self._ptr, i)
yield PyBytes_FromStringAndSize(item_view.data.as_char, item_view.size_bytes)

def _iter_str(self, int64_t offset, int64_t length):
cdef ArrowStringView item_view
for i in range(offset, length):
if ArrowArrayViewIsNull(self._ptr, i):
yield None
else:
item_view = ArrowArrayViewGetStringUnsafe(self._ptr, i)
yield PyUnicode_FromStringAndSize(item_view.data, item_view.size_bytes)

def __repr__(self):
return _repr_utils.array_view_repr(self)

Expand Down
8 changes: 8 additions & 0 deletions python/src/nanoarrow/iterator.py
Original file line number Diff line number Diff line change
Expand Up @@ -322,6 +322,12 @@ def _binary_iter(self, offset, length):
for start, end in zip(starts, ends):
yield bytes(data[start:end])

def _binary_view_iter(self, offset, length):
return self._array_view._iter_bytes(offset, length)

def _string_view_iter(self, offset, length):
return self._array_view._iter_str(offset, length)

def _decimal_iter(self, offset, length):
from decimal import Context, Decimal
from sys import byteorder
Expand Down Expand Up @@ -564,6 +570,8 @@ def _get_tzinfo(tz_string, strategy=None):
_types.DURATION: "_duration_iter",
_types.DECIMAL128: "_decimal_iter",
_types.DECIMAL256: "_decimal_iter",
_types.STRING_VIEW: "_string_view_iter",
_types.BINARY_VIEW: "_binary_view_iter",
}

_PRIMITIVE_TYPE_NAMES = [
Expand Down
28 changes: 20 additions & 8 deletions python/tests/test_iterator.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,35 +68,47 @@ def test_iterator_nullable_primitive():
assert list(iter_py(sliced)) == [2, 3, None]


def test_iterator_string():
array = na.c_array(["ab", "cde"], na.string())
@pytest.mark.parametrize(
"arrow_type", [na.string(), na.large_string(), na.string_view()]
)
def test_iterator_string(arrow_type):
array = na.c_array(["ab", "cde"], arrow_type)

assert list(iter_py(array)) == ["ab", "cde"]

sliced = array[1:]
assert list(iter_py(sliced)) == ["cde"]


def test_iterator_nullable_string():
array = na.c_array(["ab", "cde", None], na.string())
@pytest.mark.parametrize(
"arrow_type", [na.string(), na.large_string(), na.string_view()]
)
def test_iterator_nullable_string(arrow_type):
array = na.c_array(["ab", "cde", None], arrow_type)

assert list(iter_py(array)) == ["ab", "cde", None]

sliced = array[1:]
assert list(iter_py(sliced)) == ["cde", None]


def test_iterator_binary():
array = na.c_array([b"ab", b"cde"], na.binary())
@pytest.mark.parametrize(
"arrow_type", [na.binary(), na.large_binary(), na.binary_view()]
)
def test_iterator_binary(arrow_type):
array = na.c_array([b"ab", b"cde"], arrow_type)

assert list(iter_py(array)) == [b"ab", b"cde"]

sliced = array[1:]
assert list(iter_py(sliced)) == [b"cde"]


def test_iterator_nullable_binary():
array = na.c_array([b"ab", b"cde", None], na.binary())
@pytest.mark.parametrize(
"arrow_type", [na.binary(), na.large_binary(), na.binary_view()]
)
def test_iterator_nullable_binary(arrow_type):
array = na.c_array([b"ab", b"cde", None], arrow_type)

assert list(iter_py(array)) == [b"ab", b"cde", None]

Expand Down

0 comments on commit 99e5173

Please sign in to comment.