[arrow] improve test coverage and error handling
This commit is contained in:
parent
f3a60da30b
commit
151a7f846f
3 changed files with 160 additions and 9 deletions
|
|
@ -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]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue