[arrow] improve test coverage and error handling

This commit is contained in:
Jan Petykiewicz 2026-04-21 21:50:18 -07:00
commit 151a7f846f
3 changed files with 160 additions and 9 deletions

View file

@ -37,6 +37,7 @@ import sys
import tempfile
from pprint import pformat
from klamath.basic import KlamathError
import numpy
from numpy.typing import ArrayLike, NDArray
import pyarrow
@ -54,9 +55,10 @@ logger = logging.getLogger(__name__)
ffi.cdef(
"""
void read_path(char* path, struct ArrowArray* array, struct ArrowSchema* schema);
void scan_bytes(uint8_t* data, size_t size, struct ArrowArray* array, struct ArrowSchema* schema);
void read_cells_bytes(
const char* last_error_message(void);
int read_path(const char* path, struct ArrowArray* array, struct ArrowSchema* schema);
int scan_bytes(uint8_t* data, size_t size, struct ArrowArray* array, struct ArrowSchema* schema);
int read_cells_bytes(
uint8_t* data,
size_t size,
uint64_t* ranges,
@ -203,11 +205,11 @@ def _read_to_arrow(
tmp_stream.write(data)
tmp_name = tmp_stream.name
try:
_get_clib().read_path(tmp_name.encode(), ptr_array, ptr_schema)
_call_native(_get_clib().read_path(tmp_name.encode(), ptr_array, ptr_schema), 'read_path')
finally:
pathlib.Path(tmp_name).unlink(missing_ok=True)
else:
_get_clib().read_path(str(path).encode(), ptr_array, ptr_schema)
_call_native(_get_clib().read_path(str(path).encode(), ptr_array, ptr_schema), 'read_path')
return _import_arrow_array(ptr_array, ptr_schema)
@ -217,12 +219,24 @@ def _import_arrow_array(ptr_array: Any, ptr_schema: Any) -> pyarrow.Array:
return pyarrow.Array._import_from_c(iptr_array, iptr_schema)
def _call_native(status: int, action: str) -> None:
if status == 0:
return
err_ptr = _get_clib().last_error_message()
if err_ptr == ffi.NULL:
raise KlamathError(f'{action} failed')
message = ffi.string(err_ptr).decode(errors='replace')
raise KlamathError(message)
def _scan_buffer_to_arrow(buffer: bytes | mmap.mmap | memoryview) -> pyarrow.Array:
ptr_array = ffi.new('struct ArrowArray[]', 1)
ptr_schema = ffi.new('struct ArrowSchema[]', 1)
buf_view = memoryview(buffer)
cbuf = ffi.from_buffer('uint8_t[]', buf_view)
_get_clib().scan_bytes(cbuf, len(buf_view), ptr_array, ptr_schema)
_call_native(_get_clib().scan_bytes(cbuf, len(buf_view), ptr_array, ptr_schema), 'scan_bytes')
return _import_arrow_array(ptr_array, ptr_schema)
@ -236,7 +250,10 @@ def _read_selected_cells_to_arrow(
cbuf = ffi.from_buffer('uint8_t[]', buf_view)
flat_ranges = numpy.require(ranges, dtype=numpy.uint64, requirements=('C_CONTIGUOUS', 'ALIGNED'))
cranges = ffi.from_buffer('uint64_t[]', flat_ranges)
_get_clib().read_cells_bytes(cbuf, len(buf_view), cranges, int(flat_ranges.shape[0]), ptr_array, ptr_schema)
_call_native(
_get_clib().read_cells_bytes(cbuf, len(buf_view), cranges, int(flat_ranges.shape[0]), ptr_array, ptr_schema),
'read_cells_bytes',
)
return _import_arrow_array(ptr_array, ptr_schema)
@ -713,7 +730,9 @@ def _gpaths_to_mpaths(
layer = layer_tups[int(elem_layer_inds[ee])]
vertices = xy_val[xy_offs[ee]:xy_offs[ee + 1]]
width = elem_widths[ee]
cap_int = elem_path_types[ee]
cap_int = int(elem_path_types[ee])
if cap_int not in path_cap_map:
raise PatternError(f'Unrecognized path type: {cap_int}')
cap = path_cap_map[cap_int]
if cap_int == 4:
cap_extensions = elem_extensions[ee]