From b932c11b2db04e11d3cd0e4e59d2f395ddaa3e18 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 28 Oct 2026 00:42:42 +0000 Subject: [PATCH 1/3] Initial plan From 400b4582109ea131ae26f223479e8be982ae87b2 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 28 Oct 2026 00:52:04 +0000 Subject: [PATCH 2/3] Add C Parquet-to-NumPy reader with optional PyArrow backend Co-authored-by: flexatone <395503+flexatone@users.noreply.github.com> --- README.rst | 35 +++++++ pyproject.toml | 3 + setup.py | 1 + src/__init__.py | 1 + src/__init__.pyi | 7 ++ src/_arraykit.c | 5 + src/parquet_to_arrays.c | 183 +++++++++++++++++++++++++++++++++ src/parquet_to_arrays.h | 4 + test/test_parquet_to_arrays.py | 174 +++++++++++++++++++++++++++++++ 9 files changed, 413 insertions(+) create mode 100644 src/parquet_to_arrays.c create mode 100644 src/parquet_to_arrays.h create mode 100644 test/test_parquet_to_arrays.py diff --git a/README.rst b/README.rst index 1ee56c47..962845ca 100644 --- a/README.rst +++ b/README.rst @@ -32,6 +32,41 @@ ArrayKit requires the following: - numpy>=1.24.3 +Reading Parquet +-------------- + +``parquet_to_arrays(file, *, columns_select=None)`` returns a list of immutable, +one-dimensional NumPy arrays, one per top-level Parquet column. Install the +optional dependency with ``pip install 'arraykit[parquet]'``. + +``file`` can be a string, bytes, or path-like filesystem path, or a seekable +binary file such as an open file or ``io.BytesIO``. Caller-owned files are not +closed; their position can change. Paths identify individual files, not datasets +or directories. + +``columns_select`` can be a list of unique column names. Only those columns are +decoded, in selection order. ``None`` loads all columns in file order; ``[]`` +returns an empty list. Missing or ambiguous names and repeated selections raise +``ValueError``. Names select whole top-level columns (including nested columns), +not dotted paths to nested fields. + +The C entry point calls PyArrow's compiled Parquet reader and NumPy conversion +methods directly, without Python's ``read_table`` / dataset layer, pandas, or +Python loops over columns or values. PyArrow is still required: this is not a +standalone C Parquet decoder. Decoding uses Arrow's codecs and column threading. +Compatible single-chunk numeric data can share Arrow-owned memory; multiple +chunks, booleans, nulls, and object-valued data can require copies. Performance +relative to Python PyArrow depends on the file, column types, and selection. + +Dtypes and null handling follow PyArrow's ``ChunkedArray.to_numpy()``: numeric +types are preserved where possible, nullable integers become floating-point +arrays with ``NaN``, nullable booleans and strings use object arrays with +``None``, and timestamps use ``datetime64`` with ``NaT``. Large nullable integers +can lose precision when converted to floating point. Dictionary columns are +decoded to values, and nested columns use object arrays. Unsupported conversions +and malformed files propagate PyArrow's exceptions. + + What is New in ArrayKit ------------------------- diff --git a/pyproject.toml b/pyproject.toml index a1367d5b..587a4df7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,9 @@ classifiers = [ [project.urls] Homepage = "https://github.com/static-frame/arraykit" +[project.optional-dependencies] +parquet = ["pyarrow>=26.0.0,<27"] + [tool.setuptools] package-dir = { "arraykit" = "src" } diff --git a/setup.py b/setup.py index 46e91660..d96e2d14 100644 --- a/setup.py +++ b/setup.py @@ -33,6 +33,7 @@ def get_ext_dir(*components: tp.Iterable[str]) -> tp.Sequence[str]: 'src/block_index.c', 'src/delimited_to_arrays.c', 'src/methods.c', + 'src/parquet_to_arrays.c', 'src/tri_map.c', 'src/auto_map.c', ], diff --git a/src/__init__.py b/src/__init__.py index 7ae728be..c3d5d13e 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -21,6 +21,7 @@ from ._arraykit import isna_element as isna_element from ._arraykit import dtype_from_element as dtype_from_element from ._arraykit import delimited_to_arrays as delimited_to_arrays +from ._arraykit import parquet_to_arrays as parquet_to_arrays from ._arraykit import iterable_str_to_array_1d as iterable_str_to_array_1d from ._arraykit import split_after_count as split_after_count from ._arraykit import get_new_indexers_and_screen as get_new_indexers_and_screen diff --git a/src/__init__.pyi b/src/__init__.pyi index aaeb3af3..85d1d0bc 100644 --- a/src/__init__.pyi +++ b/src/__init__.pyi @@ -1,5 +1,6 @@ import typing as tp import datetime +import os import numpy as np # type: ignore @@ -7,6 +8,12 @@ _T = tp.TypeVar('_T') __version__: str +def parquet_to_arrays( + file: tp.Union[str, bytes, os.PathLike, tp.BinaryIO], + *, + columns_select: tp.Optional[tp.List[str]] = None, +) -> tp.List[np.ndarray]: ... + _TLabel = tp.Union[ tp.Hashable, int, diff --git a/src/_arraykit.c b/src/_arraykit.c index 0d14b23f..12bef785 100644 --- a/src/_arraykit.c +++ b/src/_arraykit.c @@ -10,6 +10,7 @@ # include "block_index.h" # include "delimited_to_arrays.h" # include "methods.h" +# include "parquet_to_arrays.h" # include "tri_map.h" # include "auto_map.h" @@ -43,6 +44,10 @@ static PyMethodDef arraykit_methods[] = { (PyCFunction)delimited_to_arrays, METH_VARARGS | METH_KEYWORDS, NULL}, + {"parquet_to_arrays", + (PyCFunction)parquet_to_arrays, + METH_VARARGS | METH_KEYWORDS, + "Read a Parquet file into immutable 1D arrays; requires PyArrow."}, {"iterable_str_to_array_1d", (PyCFunction)iterable_str_to_array_1d, METH_VARARGS | METH_KEYWORDS, diff --git a/src/parquet_to_arrays.c b/src/parquet_to_arrays.c new file mode 100644 index 00000000..3b2dfe95 --- /dev/null +++ b/src/parquet_to_arrays.c @@ -0,0 +1,183 @@ +# include "Python.h" + +# define PY_ARRAY_UNIQUE_SYMBOL AK_ARRAY_API +# define NO_IMPORT_ARRAY +# define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION +# include "numpy/arrayobject.h" + +# include "parquet_to_arrays.h" + +PyObject * +parquet_to_arrays(PyObject *Py_UNUSED(self), PyObject *args, PyObject *kwargs) +{ + PyObject *source; + PyObject *columns_select = Py_None; + static char *kwlist[] = {"file", "columns_select", NULL}; + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|$O:parquet_to_arrays", + kwlist, &source, &columns_select)) { + return NULL; + } + + PyObject *selection = NULL; + PyObject *source_normalized = NULL; + PyObject *module = NULL; + PyObject *reader = NULL; + PyObject *paths = NULL; + PyObject *indices = NULL; + PyObject *schema = NULL; + PyObject *names = NULL; + PyObject *lookup = NULL; + PyObject *seen = NULL; + PyObject *table = NULL; + PyObject *result = NULL; + PyObject *value = NULL; + PyObject *column = NULL; + PyObject *array = NULL; + + if (columns_select != Py_None) { + if (!PyList_Check(columns_select)) { + PyErr_SetString(PyExc_TypeError, + "columns_select must be a list of strings or None"); + goto error; + } + selection = PyList_GetSlice(columns_select, 0, PyList_GET_SIZE(columns_select)); + if (selection == NULL) goto error; + seen = PySet_New(NULL); + if (seen == NULL) goto error; + for (Py_ssize_t i = 0; i < PyList_GET_SIZE(selection); ++i) { + PyObject *name = PyList_GET_ITEM(selection, i); + if (!PyUnicode_Check(name)) { + PyErr_SetString(PyExc_TypeError, + "columns_select must contain only strings"); + goto error; + } + int contains = PySet_Contains(seen, name); + if (contains < 0) goto error; + if (contains) { + PyErr_Format(PyExc_ValueError, "duplicate column selection: %R", name); + goto error; + } + if (PySet_Add(seen, name) < 0) goto error; + } + } + + source_normalized = PyOS_FSPath(source); + if (source_normalized == NULL) { + if (!PyErr_ExceptionMatches(PyExc_TypeError)) goto error; + PyErr_Clear(); + source_normalized = Py_NewRef(source); + } + if (PyBytes_Check(source_normalized)) { + value = PyUnicode_DecodeFSDefaultAndSize( + PyBytes_AS_STRING(source_normalized), PyBytes_GET_SIZE(source_normalized)); + if (value == NULL) goto error; + Py_SETREF(source_normalized, value); + value = NULL; + } + module = PyImport_ImportModule("pyarrow._parquet"); + if (module == NULL) goto error; + value = PyObject_GetAttrString(module, "ParquetReader"); + if (value == NULL) goto error; + reader = PyObject_CallNoArgs(value); + Py_CLEAR(value); + if (reader == NULL) goto error; + value = PyObject_CallMethod(reader, "open", "(O)", source_normalized); + if (value == NULL) goto error; + Py_CLEAR(value); + + if (selection != NULL) { + paths = PyObject_GetAttrString(reader, "column_paths"); + schema = PyObject_GetAttrString(reader, "schema_arrow"); + if (paths == NULL || schema == NULL) goto error; + names = PyObject_GetAttrString(schema, "names"); + if (names == NULL) goto error; + indices = PyList_New(0); + if (indices == NULL) goto error; + lookup = PyDict_New(); + if (lookup == NULL) goto error; + for (Py_ssize_t i = 0; i < PyList_Size(names); ++i) { + PyObject *name = PyList_GetItem(names, i); + int contains = PyDict_Contains(lookup, name); + if (contains < 0) goto error; + value = contains ? Py_NewRef(Py_None) : PyList_New(0); + if (value == NULL || PyDict_SetItem(lookup, name, value) < 0) goto error; + Py_CLEAR(value); + } + // Parquet indices refer to leaves, not top-level Arrow fields. + for (Py_ssize_t i = 0; i < PyList_Size(paths); ++i) { + PyObject *path = PyList_GetItem(paths, i); + PyObject *root = PyList_GetItem(path, 0); + if (root == NULL) goto error; + PyObject *leaves = PyDict_GetItemWithError(lookup, root); + if (leaves == NULL) { + if (!PyErr_Occurred()) { + PyErr_SetString(PyExc_ValueError, "inconsistent Parquet schema"); + } + goto error; + } + if (leaves != Py_None) { + value = PyLong_FromSsize_t(i); + if (value == NULL || PyList_Append(leaves, value) < 0) goto error; + Py_CLEAR(value); + } + } + for (Py_ssize_t i = 0; i < PyList_GET_SIZE(selection); ++i) { + PyObject *name = PyList_GET_ITEM(selection, i); + PyObject *leaves = PyDict_GetItemWithError(lookup, name); + if (leaves == NULL && PyErr_Occurred()) goto error; + if (leaves == NULL || leaves == Py_None) { + PyErr_Format(PyExc_ValueError, + "column name is missing or ambiguous: %R", name); + goto error; + } + Py_ssize_t size = PyList_GET_SIZE(indices); + if (PyList_SetSlice(indices, size, size, leaves) < 0) goto error; + } + } + + // Both reading and NumPy conversion are compiled PyArrow entry points. + table = PyObject_CallMethod(reader, "read_all", "(O)", + indices == NULL ? Py_None : indices); + if (table == NULL) goto error; + value = PyObject_GetAttrString(table, "num_columns"); + if (value == NULL) goto error; + Py_ssize_t count = PyLong_AsSsize_t(value); + Py_CLEAR(value); + if (count == -1 && PyErr_Occurred()) goto error; + result = PyList_New(count); + if (result == NULL) goto error; + for (Py_ssize_t i = 0; i < count; ++i) { + column = PyObject_CallMethod(table, "column", "n", i); + if (column == NULL) goto error; + array = PyObject_CallMethod(column, "to_numpy", NULL); + Py_CLEAR(column); + if (array == NULL) goto error; + if (!PyArray_Check(array) || PyArray_NDIM((PyArrayObject *)array) != 1) { + PyErr_SetString(PyExc_TypeError, "Parquet column did not produce a 1D array"); + goto error; + } + PyArray_CLEARFLAGS((PyArrayObject *)array, NPY_ARRAY_WRITEABLE); + PyList_SET_ITEM(result, i, array); + array = NULL; + } + goto done; + +error: + Py_CLEAR(result); +done: + Py_XDECREF(array); + Py_XDECREF(column); + Py_XDECREF(value); + Py_XDECREF(table); + Py_XDECREF(names); + Py_XDECREF(schema); + Py_XDECREF(lookup); + Py_XDECREF(seen); + Py_XDECREF(indices); + Py_XDECREF(paths); + Py_XDECREF(reader); + Py_XDECREF(module); + Py_XDECREF(selection); + Py_XDECREF(source_normalized); + return result; +} diff --git a/src/parquet_to_arrays.h b/src/parquet_to_arrays.h new file mode 100644 index 00000000..f748bfb5 --- /dev/null +++ b/src/parquet_to_arrays.h @@ -0,0 +1,4 @@ +# include "Python.h" + +PyObject * +parquet_to_arrays(PyObject *self, PyObject *args, PyObject *kwargs); diff --git a/test/test_parquet_to_arrays.py b/test/test_parquet_to_arrays.py new file mode 100644 index 00000000..2fc71ec8 --- /dev/null +++ b/test/test_parquet_to_arrays.py @@ -0,0 +1,174 @@ +import gc +import io +import os +from pathlib import Path +import tempfile +import unittest +from unittest import mock + +import numpy as np + +from arraykit import parquet_to_arrays + +try: + import pyarrow as pa + import pyarrow.parquet as pq +except ImportError: + pa = None + pq = None + + +@unittest.skipIf(pa is None, 'PyArrow is not installed') +class TestUnit(unittest.TestCase): + @staticmethod + def to_file(table, **kwargs): + file = io.BytesIO() + pq.write_table(table, file, **kwargs) + file.seek(0) + return file + + def test_parquet_to_arrays_types(self): + table = pa.table( + { + 'int': pa.array([1, 2, 3], type=pa.int32()), + 'uint': pa.array([0, 2**64 - 1, 3], type=pa.uint64()), + 'float': pa.array([1.5, 2.5, 3.5], type=pa.float32()), + 'bool': [True, False, True], + 'string': ['a', 'é', '水'], + 'binary': [b'a\x00', b'b', b'c'], + 'time': pa.array([1, 2, 3], type=pa.timestamp('us')), + } + ) + for compression in (None, 'snappy', 'gzip', 'zstd'): + with self.subTest(compression=compression): + arrays = parquet_to_arrays(self.to_file(table, compression=compression)) + self.assertIsInstance(arrays, list) + self.assertEqual(len(arrays), table.num_columns) + for array, expected in zip(arrays, table.columns): + self.assertEqual(array.ndim, 1) + self.assertFalse(array.flags.writeable) + np.testing.assert_array_equal(array, expected.to_numpy()) + self.assertEqual(array.dtype, expected.to_numpy().dtype) + + def test_parquet_to_arrays_paths(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / 'é.parquet' + pq.write_table(pa.table({'x': [1, 2]}), path) + for source in (path, str(path), os.fsencode(path)): + with self.subTest(source=source): + np.testing.assert_array_equal(parquet_to_arrays(source)[0], [1, 2]) + with path.open('rb') as file: + np.testing.assert_array_equal(parquet_to_arrays(file)[0], [1, 2]) + self.assertFalse(file.closed) + file.seek(0) + np.testing.assert_array_equal(parquet_to_arrays(file)[0], [1, 2]) + + def test_parquet_to_arrays_select(self): + file = self.to_file(pa.table({'a': [1], 'b': [2], 'c': [3]})) + arrays = parquet_to_arrays(file, columns_select=['c', 'a']) + self.assertEqual([a.tolist() for a in arrays], [[3], [1]]) + self.assertEqual(parquet_to_arrays(file, columns_select=[]), []) + self.assertEqual( + [a.tolist() for a in parquet_to_arrays(file, columns_select=None)], + [[1], [2], [3]], + ) + for selection in (['missing'], ['a', 'a']): + with self.subTest(selection=selection): + with self.assertRaises(ValueError): + parquet_to_arrays(file, columns_select=selection) + self.assertFalse(file.closed) + for selection in ('a', ('a',), [1], [None]): + with self.subTest(selection=selection): + with self.assertRaises(TypeError): + parquet_to_arrays(file, columns_select=selection) + with self.assertRaises(TypeError): + parquet_to_arrays(file, ['a']) + + def test_parquet_to_arrays_nulls_and_row_groups(self): + table = pa.table( + { + 'int': [1, None, 3, 4], + 'float': [1.5, None, float('nan'), 4.5], + 'bool': [True, None, False, True], + 'string': ['a', None, '', 'd'], + 'null': pa.nulls(4), + 'time': pa.array([1, None, 3, 4], type=pa.timestamp('ns')), + } + ) + arrays = parquet_to_arrays(self.to_file(table, row_group_size=1)) + gc.collect() + for array, column in zip(arrays, table.columns): + np.testing.assert_array_equal(array, column.to_numpy()) + self.assertFalse(array.flags.writeable) + + def test_parquet_to_arrays_empty(self): + table = pa.table({'int': pa.array([], type=pa.int64()), 'string': []}) + arrays = parquet_to_arrays(self.to_file(table)) + self.assertEqual([a.shape for a in arrays], [(0,), (0,)]) + self.assertEqual(arrays[0].dtype, np.dtype(np.int64)) + self.assertEqual(parquet_to_arrays(self.to_file(pa.table({}))), []) + + def test_parquet_to_arrays_nested(self): + table = pa.table( + { + 'struct': [{'a': 1, 'b': 'x'}, {'a': 2, 'b': 'y'}], + 'list': [[1, 2], [3]], + 'a.b': [5, 6], + } + ) + file = self.to_file(table) + arrays = parquet_to_arrays(file, columns_select=['a.b', 'struct', 'list']) + self.assertEqual(arrays[0].tolist(), [5, 6]) + self.assertEqual(arrays[1].tolist(), table.column('struct').to_pylist()) + self.assertEqual([a.tolist() for a in arrays[2]], [[1, 2], [3]]) + with self.assertRaises(ValueError): + parquet_to_arrays(file, columns_select=['struct.a']) + + def test_parquet_to_arrays_dictionary(self): + table = pa.table({'x': pa.array(['a', 'b', 'a']).dictionary_encode()}) + array = parquet_to_arrays(self.to_file(table))[0] + self.assertEqual(array.tolist(), ['a', 'b', 'a']) + self.assertEqual(array.dtype, np.dtype(object)) + + def test_parquet_to_arrays_wide_projection(self): + names = [f'c{i}' for i in range(128)] + table = pa.table({name: [i] for i, name in enumerate(names)}) + arrays = parquet_to_arrays(self.to_file(table), columns_select=names[::-1]) + self.assertEqual([a[0] for a in arrays], list(range(127, -1, -1))) + + def test_parquet_to_arrays_ambiguous_names(self): + table = pa.Table.from_arrays([pa.array([1]), pa.array([2])], names=['x', 'x']) + file = self.to_file(table) + self.assertEqual([a.tolist() for a in parquet_to_arrays(file)], [[1], [2]]) + with self.assertRaises(ValueError): + parquet_to_arrays(file, columns_select=['x']) + + def test_parquet_to_arrays_errors(self): + for source in (io.BytesIO(b'not parquet'), io.StringIO('not parquet'), 42): + with self.subTest(source=source): + with self.assertRaises((ValueError, TypeError, OSError)): + parquet_to_arrays(source) + file = self.to_file(pa.table({'a': [1]})) + file.close() + with self.assertRaises(ValueError): + parquet_to_arrays(file) + with tempfile.TemporaryDirectory() as directory: + with self.assertRaises(OSError): + parquet_to_arrays(Path(directory) / 'missing.parquet') + + def test_parquet_to_arrays_native_reader(self): + file = self.to_file(pa.table({'a': [1, 2]})) + with mock.patch.object(pq, 'read_table', side_effect=AssertionError): + with mock.patch.object(pq, 'ParquetFile', side_effect=AssertionError): + np.testing.assert_array_equal(parquet_to_arrays(file)[0], [1, 2]) + + +class TestOptionalDependency(unittest.TestCase): + def test_parquet_to_arrays_missing_pyarrow(self): + with mock.patch.dict('sys.modules', {'pyarrow._parquet': None}): + with self.assertRaises(ImportError): + parquet_to_arrays(io.BytesIO()) + + +if __name__ == '__main__': + unittest.main() From e2a2e18075d2aa661b25607efc34d452b736dfe1 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:16:58 +0000 Subject: [PATCH 3/3] Replace PyArrow wrapper with directly linked native Arrow/Parquet reader Co-authored-by: flexatone <395503+flexatone@users.noreply.github.com> --- MANIFEST.in | 1 + README.rst | 56 ++- pyproject.toml | 3 - setup.py | 55 +++ src/_arraykit.c | 2 +- src/parquet_reader.cpp | 664 +++++++++++++++++++++++++++++++++ src/parquet_to_arrays.c | 192 ++-------- test/test_parquet_build.py | 68 ++++ test/test_parquet_to_arrays.py | 129 ++++++- 9 files changed, 972 insertions(+), 198 deletions(-) create mode 100644 src/parquet_reader.cpp create mode 100644 test/test_parquet_build.py diff --git a/MANIFEST.in b/MANIFEST.in index cb6cad3a..b569572d 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,3 +1,4 @@ include src/__init__.pyi src/py.typed include src/*.h +include src/*.cpp include VERSION diff --git a/README.rst b/README.rst index 962845ca..9d7e1366 100644 --- a/README.rst +++ b/README.rst @@ -36,8 +36,29 @@ Reading Parquet -------------- ``parquet_to_arrays(file, *, columns_select=None)`` returns a list of immutable, -one-dimensional NumPy arrays, one per top-level Parquet column. Install the -optional dependency with ``pip install 'arraykit[parquet]'``. +one-dimensional NumPy arrays, one per top-level Parquet column. It uses the native +Arrow and Parquet C++ libraries, not Python PyArrow or ``arrow_python``. + +Parquet support is opt-in at build time. Install Arrow/Parquet >=24 development +libraries and a C++20 compiler, then build ArrayKit from source:: + + ARRAYKIT_WITH_PARQUET=1 pip install --no-binary=arraykit arraykit + +The build finds libraries using ``pkg-config``. Alternatively set ``ARROW_HOME`` +to a native installation prefix containing ``include`` and ``lib``, or set both +``ARROW_INCLUDE_DIR`` and ``ARROW_LIBRARY_DIR``. These settings are build-time +only; the resulting extension must be able to load the same native libraries at +runtime. Explicit library directories are embedded as runtime search paths on +Unix. On Windows, make the Arrow DLL directory available to the loader (for +example with ``os.add_dll_directory()`` before calling the reader). Keep native +library installations ABI-compatible with the version used to build the extension. +Locally built wheels with this backend are not portable unless their native +dependencies are also supplied. + +Normal builds still require only NumPy. Calling ``parquet_to_arrays`` without +the native backend raises ``ImportError`` with build instructions. PyArrow is +used only to generate fixtures and reference results in optional tests, never +for loading Parquet in ArrayKit. ``file`` can be a string, bytes, or path-like filesystem path, or a seekable binary file such as an open file or ``io.BytesIO``. Caller-owned files are not @@ -50,21 +71,24 @@ returns an empty list. Missing or ambiguous names and repeated selections raise ``ValueError``. Names select whole top-level columns (including nested columns), not dotted paths to nested fields. -The C entry point calls PyArrow's compiled Parquet reader and NumPy conversion -methods directly, without Python's ``read_table`` / dataset layer, pandas, or -Python loops over columns or values. PyArrow is still required: this is not a -standalone C Parquet decoder. Decoding uses Arrow's codecs and column threading. -Compatible single-chunk numeric data can share Arrow-owned memory; multiple -chunks, booleans, nulls, and object-valued data can require copies. Performance -relative to Python PyArrow depends on the file, column types, and selection. - -Dtypes and null handling follow PyArrow's ``ChunkedArray.to_numpy()``: numeric -types are preserved where possible, nullable integers become floating-point -arrays with ``NaN``, nullable booleans and strings use object arrays with -``None``, and timestamps use ``datetime64`` with ``NaT``. Large nullable integers +The native extension reads paths through Arrow's random-access file API. +Open Python files use binary ``read``, ``seek``, and ``tell`` callbacks; the +Parquet decoder and array conversion remain native. It releases the GIL during +reading and decoding. Arrow handles compression and page encodings; ArrayKit +copies decoded columns directly into NumPy-owned arrays through the C API, +without a Python Arrow table, pandas, or Python loops over columns or values. +Performance relative to Python PyArrow depends on the file, column types, and +selection. + +Dtypes and null handling preserve numeric types where possible; nullable integers +become floating-point arrays with ``NaN``, nullable booleans and strings use object +arrays with ``None``, and timestamps use ``datetime64`` with ``NaT``. Large nullable integers can lose precision when converted to floating point. Dictionary columns are -decoded to values, and nested columns use object arrays. Unsupported conversions -and malformed files propagate PyArrow's exceptions. +decoded to values, decimal columns use ``decimal.Decimal`` objects, and nested +struct/list columns use object arrays. Other Arrow types may not have a supported +NumPy conversion. Unsupported conversions raise ``NotImplementedError``; +malformed files raise ``ValueError`` or ``OSError``. Exceptions from Python file +methods are preserved. What is New in ArrayKit diff --git a/pyproject.toml b/pyproject.toml index 587a4df7..a1367d5b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,9 +44,6 @@ classifiers = [ [project.urls] Homepage = "https://github.com/static-frame/arraykit" -[project.optional-dependencies] -parquet = ["pyarrow>=26.0.0,<27"] - [tool.setuptools] package-dir = { "arraykit" = "src" } diff --git a/setup.py b/setup.py index d96e2d14..4d4db1a5 100644 --- a/setup.py +++ b/setup.py @@ -2,6 +2,8 @@ from setuptools import setup, Extension from setuptools.command.build_ext import build_ext import site, os +import shlex +import subprocess from pathlib import Path AK_VERSION = Path('VERSION').read_text(encoding='utf-8').strip() @@ -44,4 +46,57 @@ def get_ext_dir(*components: tp.Iterable[str]) -> tp.Sequence[str]: ) ] + +def get_parquet_extension() -> Extension: + prefix = os.environ.get('ARROW_HOME') + include_dir = os.environ.get('ARROW_INCLUDE_DIR') + library_dir = os.environ.get('ARROW_LIBRARY_DIR') + if prefix: + include_dir = include_dir or os.path.join(prefix, 'include') + library_dir = library_dir or os.path.join(prefix, 'lib') + compile_args = ['/std:c++20'] if os.name == 'nt' else ['-std=c++20'] + link_args = [] + library_dirs = [] + include_dirs = get_ext_dir('numpy', '_core', 'include') + if include_dir or library_dir: + if not (include_dir and library_dir): + raise RuntimeError('Set both ARROW_INCLUDE_DIR and ARROW_LIBRARY_DIR') + include_dirs.append(include_dir) + library_dirs.append(library_dir) + else: + try: + subprocess.check_call( + ['pkg-config', '--atleast-version=24', 'arrow', 'parquet'] + ) + compile_args.extend(shlex.split(subprocess.check_output( + ['pkg-config', '--cflags', 'arrow', 'parquet'], text=True + ))) + link_args = shlex.split(subprocess.check_output( + ['pkg-config', '--libs', 'arrow', 'parquet'], text=True + )) + except (OSError, subprocess.CalledProcessError) as error: + raise RuntimeError( + 'Native Parquet support requires Arrow/Parquet >=24 development ' + 'libraries. Install their pkg-config files or set ARROW_HOME ' + '(or ARROW_INCLUDE_DIR and ARROW_LIBRARY_DIR).' + ) from error + runtime_dirs = library_dirs + [ + flag[2:] for flag in link_args if flag.startswith('-L') + ] + return Extension( + name='arraykit._parquet', + sources=['src/parquet_reader.cpp'], + include_dirs=include_dirs, + library_dirs=library_dirs, + runtime_library_dirs=runtime_dirs if os.name != 'nt' else [], + libraries=['parquet', 'arrow'], + extra_compile_args=compile_args, + extra_link_args=link_args, + language='c++', + ) + + +if os.environ.get('ARRAYKIT_WITH_PARQUET') == '1': + ext_modules.append(get_parquet_extension()) + setup(ext_modules=ext_modules) diff --git a/src/_arraykit.c b/src/_arraykit.c index 12bef785..4e4f2550 100644 --- a/src/_arraykit.c +++ b/src/_arraykit.c @@ -47,7 +47,7 @@ static PyMethodDef arraykit_methods[] = { {"parquet_to_arrays", (PyCFunction)parquet_to_arrays, METH_VARARGS | METH_KEYWORDS, - "Read a Parquet file into immutable 1D arrays; requires PyArrow."}, + "Read a Parquet file into immutable 1D arrays; requires the native backend."}, {"iterable_str_to_array_1d", (PyCFunction)iterable_str_to_array_1d, METH_VARARGS | METH_KEYWORDS, diff --git a/src/parquet_reader.cpp b/src/parquet_reader.cpp new file mode 100644 index 00000000..e42bff00 --- /dev/null +++ b/src/parquet_reader.cpp @@ -0,0 +1,664 @@ +#define PY_SSIZE_T_CLEAN +#include +#define PY_ARRAY_UNIQUE_SYMBOL AK_PARQUET_ARRAY_API +#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +struct PythonError {}; + +class PyRef { +public: + explicit PyRef(PyObject* value = nullptr) : value_(value) {} + ~PyRef() { Py_XDECREF(value_); } + PyRef(const PyRef&) = delete; + PyRef& operator=(const PyRef&) = delete; + PyObject* get() const { return value_; } + PyObject* release() { PyObject* value = value_; value_ = nullptr; return value; } +private: + PyObject* value_; +}; + +class AcquireGIL { +public: + AcquireGIL() : state_(PyGILState_Ensure()) {} + ~AcquireGIL() { PyGILState_Release(state_); } +private: + PyGILState_STATE state_; +}; + +class ReleaseGIL { +public: + ReleaseGIL() : state_(PyEval_SaveThread()) {} + ~ReleaseGIL() { PyEval_RestoreThread(state_); } +private: + PyThreadState* state_; +}; + +class RecursionGuard { +public: + RecursionGuard() { + if (Py_EnterRecursiveCall(" while converting a Parquet value")) + throw PythonError(); + } + ~RecursionGuard() { Py_LeaveRecursiveCall(); } +}; + +PyObject* checked(PyObject* value) { + if (!value) throw PythonError(); + return value; +} + +void check_status(const arrow::Status& status) { + if (status.ok()) return; + PyObject* exception = PyExc_ValueError; + if (status.IsIOError()) exception = PyExc_OSError; + else if (status.IsNotImplemented()) exception = PyExc_NotImplementedError; + else if (status.IsOutOfMemory()) exception = PyExc_MemoryError; + else if (status.IsTypeError()) exception = PyExc_TypeError; + PyErr_SetString(exception, status.ToString().c_str()); + throw PythonError(); +} + +template +T unwrap(arrow::Result result) { + check_status(result.status()); + return std::move(result).ValueOrDie(); +} + +// Each operation holds the mutex across seek/read, and acquires the GIL only +// for Python callbacks. Arrow reading itself runs with the GIL released. +class PythonFile final : public arrow::io::RandomAccessFile { +public: + explicit PythonFile(PyObject* file) : file_(Py_NewRef(file)) {} + ~PythonFile() override { + AcquireGIL gil; + Py_DECREF(file_); + if (error_type_ && !PyErr_Occurred()) { + PyErr_Restore(error_type_, error_value_, error_traceback_); + error_type_ = error_value_ = error_traceback_ = nullptr; + } + Py_XDECREF(error_type_); + Py_XDECREF(error_value_); + Py_XDECREF(error_traceback_); + } + + arrow::Status Close() override { return arrow::Status::OK(); } + bool closed() const override { return false; } + + arrow::Result Tell() const override { + std::lock_guard lock(mutex_); + AcquireGIL gil; + return tell(); + } + + arrow::Status Seek(int64_t position) override { + std::lock_guard lock(mutex_); + AcquireGIL gil; + return seek(position, 0); + } + + arrow::Result GetSize() override { + std::lock_guard lock(mutex_); + AcquireGIL gil; + auto original = tell(); + if (!original.ok()) return original.status(); + auto status = seek(0, 2); + if (!status.ok()) return status; + auto size = tell(); + status = seek(*original, 0); + if (!status.ok()) return status; + return size; + } + + arrow::Result Read(int64_t count, void* out) override { + std::lock_guard lock(mutex_); + AcquireGIL gil; + return read(count, out); + } + + arrow::Result> Read(int64_t count) override { + auto buffer_result = arrow::AllocateResizableBuffer(count); + if (!buffer_result.ok()) return buffer_result.status(); + std::shared_ptr buffer = std::move(*buffer_result); + auto result = Read(count, buffer->mutable_data()); + if (!result.ok()) return result.status(); + auto status = buffer->Resize(*result, false); + if (!status.ok()) return status; + return buffer; + } + + arrow::Result ReadAt(int64_t position, int64_t count, void* out) override { + std::lock_guard lock(mutex_); + AcquireGIL gil; + auto status = seek(position, 0); + if (!status.ok()) return status; + return read(count, out); + } + + arrow::Result> ReadAt( + int64_t position, int64_t count) override { + auto buffer_result = arrow::AllocateResizableBuffer(count); + if (!buffer_result.ok()) return buffer_result.status(); + std::shared_ptr buffer = std::move(*buffer_result); + auto result = ReadAt(position, count, buffer->mutable_data()); + if (!result.ok()) return result.status(); + auto status = buffer->Resize(*result, false); + if (!status.ok()) return status; + return buffer; + } + + bool restore_error() { + if (!error_type_) return false; + PyErr_Restore(error_type_, error_value_, error_traceback_); + error_type_ = error_value_ = error_traceback_ = nullptr; + return true; + } + +private: + arrow::Status failure() const { + if (!error_type_) PyErr_Fetch(&error_type_, &error_value_, &error_traceback_); + else PyErr_Clear(); + return arrow::Status::IOError("Python file operation failed"); + } + + arrow::Result tell() const { + PyRef value(PyObject_CallMethod(file_, "tell", nullptr)); + if (!value.get()) return failure(); + int64_t position = PyLong_AsLongLong(value.get()); + if (PyErr_Occurred()) return failure(); + if (position < 0) return arrow::Status::IOError("negative file position"); + return position; + } + + arrow::Status seek(int64_t position, int whence) const { + PyRef value(PyObject_CallMethod(file_, "seek", "Li", + static_cast(position), whence)); + if (!value.get()) return failure(); + return arrow::Status::OK(); + } + + arrow::Result read(int64_t count, void* out) const { + if (count < 0 || count > PY_SSIZE_T_MAX) + return arrow::Status::Invalid("invalid read size"); + int64_t total = 0; + while (total < count) { + PyRef value(PyObject_CallMethod(file_, "read", "n", + static_cast(count - total))); + if (!value.get()) return failure(); + if (!PyBytes_Check(value.get())) { + PyErr_SetString(PyExc_TypeError, "file.read() must return bytes"); + return failure(); + } + Py_ssize_t size = PyBytes_GET_SIZE(value.get()); + if (size > count - total) { + PyErr_SetString(PyExc_ValueError, "file.read() returned too many bytes"); + return failure(); + } + if (!size) break; + std::memcpy(static_cast(out) + total, + PyBytes_AS_STRING(value.get()), static_cast(size)); + total += size; + } + return total; + } + + PyObject* file_; + mutable std::mutex mutex_; + mutable PyObject* error_type_ = nullptr; + mutable PyObject* error_value_ = nullptr; + mutable PyObject* error_traceback_ = nullptr; +}; + +PyObject* convert_array(const std::shared_ptr& array); + +PyObject* object_value(const std::shared_ptr& array, int64_t index) { + if (array->IsNull(index)) return Py_NewRef(Py_None); + RecursionGuard recursion; + switch (array->type_id()) { +#define AK_SIGNED(ID, CLASS) \ + case arrow::Type::ID: return PyLong_FromLongLong( \ + static_cast(*array).Value(index)); +#define AK_UNSIGNED(ID, CLASS) \ + case arrow::Type::ID: return PyLong_FromUnsignedLongLong( \ + static_cast(*array).Value(index)); + AK_SIGNED(INT8, Int8Array) + AK_SIGNED(INT16, Int16Array) + AK_SIGNED(INT32, Int32Array) + AK_SIGNED(INT64, Int64Array) + AK_UNSIGNED(UINT8, UInt8Array) + AK_UNSIGNED(UINT16, UInt16Array) + AK_UNSIGNED(UINT32, UInt32Array) + AK_UNSIGNED(UINT64, UInt64Array) +#undef AK_SIGNED +#undef AK_UNSIGNED + case arrow::Type::BOOL: + return PyBool_FromLong(static_cast(*array).Value(index)); + case arrow::Type::HALF_FLOAT: + return PyFloat_FromDouble(arrow::util::Float16::FromBits( + static_cast(*array).Value(index)).ToDouble()); + case arrow::Type::FLOAT: + return PyFloat_FromDouble(static_cast(*array).Value(index)); + case arrow::Type::DOUBLE: + return PyFloat_FromDouble(static_cast(*array).Value(index)); + case arrow::Type::STRING: { + auto value = static_cast(*array).GetView(index); + return PyUnicode_DecodeUTF8(value.data(), value.size(), "strict"); + } + case arrow::Type::LARGE_STRING: { + auto value = static_cast(*array).GetView(index); + return PyUnicode_DecodeUTF8(value.data(), value.size(), "strict"); + } + case arrow::Type::BINARY: { + auto value = static_cast(*array).GetView(index); + return PyBytes_FromStringAndSize(value.data(), value.size()); + } + case arrow::Type::LARGE_BINARY: { + auto value = static_cast(*array).GetView(index); + return PyBytes_FromStringAndSize(value.data(), value.size()); + } + case arrow::Type::FIXED_SIZE_BINARY: { + auto value = static_cast(*array).GetView(index); + return PyBytes_FromStringAndSize(value.data(), value.size()); + } + case arrow::Type::STRUCT: { + const auto& values = static_cast(*array); + PyRef result(checked(PyDict_New())); + for (int field = 0; field < values.num_fields(); ++field) { + const auto& name = values.struct_type()->field(field)->name(); + PyRef key(checked(PyUnicode_DecodeUTF8(name.data(), name.size(), "strict"))); + PyRef value(checked(object_value(values.field(field), index))); + if (PyDict_SetItem(result.get(), key.get(), value.get()) < 0) + throw PythonError(); + } + return result.release(); + } + case arrow::Type::LIST: { + const auto& values = static_cast(*array); + return convert_array(values.value_slice(index)); + } + case arrow::Type::LARGE_LIST: { + const auto& values = static_cast(*array); + return convert_array(values.value_slice(index)); + } + case arrow::Type::FIXED_SIZE_LIST: { + const auto& values = static_cast(*array); + return convert_array(values.value_slice(index)); + } + case arrow::Type::DICTIONARY: { + const auto& values = static_cast(*array); + return object_value(values.dictionary(), values.GetValueIndex(index)); + } + case arrow::Type::DECIMAL128: + case arrow::Type::DECIMAL256: + case arrow::Type::DECIMAL32: + case arrow::Type::DECIMAL64: { + auto scalar = unwrap(array->GetScalar(index)); + std::string text = scalar->ToString(); + PyRef module(checked(PyImport_ImportModule("decimal"))); + PyRef cls(checked(PyObject_GetAttrString(module.get(), "Decimal"))); + PyRef value(checked(PyUnicode_FromStringAndSize(text.data(), text.size()))); + return PyObject_CallOneArg(cls.get(), value.get()); + } + case arrow::Type::TIMESTAMP: + case arrow::Type::DATE32: + case arrow::Type::DATE64: { + PyRef values(checked(convert_array(array->Slice(index, 1)))); + return PyArray_GETITEM(reinterpret_cast(values.get()), + static_cast(PyArray_DATA(reinterpret_cast(values.get())))); + } + case arrow::Type::NA: + return Py_NewRef(Py_None); + default: + PyErr_Format(PyExc_NotImplementedError, "unsupported Arrow type: %s", + array->type()->ToString().c_str()); + throw PythonError(); + } +} + +int numpy_type(const arrow::DataType& type, bool nullable) { + switch (type.id()) { + case arrow::Type::INT8: return nullable ? NPY_FLOAT64 : NPY_INT8; + case arrow::Type::INT16: return nullable ? NPY_FLOAT64 : NPY_INT16; + case arrow::Type::INT32: return nullable ? NPY_FLOAT64 : NPY_INT32; + case arrow::Type::INT64: return nullable ? NPY_FLOAT64 : NPY_INT64; + case arrow::Type::UINT8: return nullable ? NPY_FLOAT64 : NPY_UINT8; + case arrow::Type::UINT16: return nullable ? NPY_FLOAT64 : NPY_UINT16; + case arrow::Type::UINT32: return nullable ? NPY_FLOAT64 : NPY_UINT32; + case arrow::Type::UINT64: return nullable ? NPY_FLOAT64 : NPY_UINT64; + case arrow::Type::HALF_FLOAT: return NPY_FLOAT16; + case arrow::Type::FLOAT: return NPY_FLOAT32; + case arrow::Type::DOUBLE: return NPY_FLOAT64; + case arrow::Type::BOOL: return nullable ? NPY_OBJECT : NPY_BOOL; + case arrow::Type::TIMESTAMP: + case arrow::Type::DATE32: + case arrow::Type::DATE64: return NPY_DATETIME; + case arrow::Type::NA: + case arrow::Type::STRING: + case arrow::Type::LARGE_STRING: + case arrow::Type::BINARY: + case arrow::Type::LARGE_BINARY: + case arrow::Type::FIXED_SIZE_BINARY: + case arrow::Type::STRUCT: + case arrow::Type::LIST: + case arrow::Type::LARGE_LIST: + case arrow::Type::FIXED_SIZE_LIST: + case arrow::Type::DICTIONARY: + case arrow::Type::DECIMAL32: + case arrow::Type::DECIMAL64: + case arrow::Type::DECIMAL128: + case arrow::Type::DECIMAL256: return NPY_OBJECT; + default: + PyErr_Format(PyExc_NotImplementedError, "unsupported Arrow type: %s", + type.ToString().c_str()); + throw PythonError(); + } +} + +PyArray_Descr* datetime_dtype(const arrow::DataType& type) { + const char* unit = "D"; + if (type.id() == arrow::Type::DATE64) unit = "ms"; + else if (type.id() == arrow::Type::TIMESTAMP) { + switch (static_cast(type).unit()) { + case arrow::TimeUnit::SECOND: unit = "s"; break; + case arrow::TimeUnit::MILLI: unit = "ms"; break; + case arrow::TimeUnit::MICRO: unit = "us"; break; + case arrow::TimeUnit::NANO: unit = "ns"; break; + } + } + std::string spec = std::string("datetime64[") + unit + "]"; + PyRef value(checked(PyUnicode_FromString(spec.c_str()))); + PyArray_Descr* result = nullptr; + if (!PyArray_DescrConverter(value.get(), &result)) throw PythonError(); + return result; +} + +void copy_chunk(const std::shared_ptr& array, + PyArrayObject* output, int64_t offset, int dtype) { + if (array->length() == 0) return; + char* data = static_cast(PyArray_DATA(output)); + if (dtype == NPY_OBJECT) { + auto values = reinterpret_cast(data); + for (int64_t i = 0; i < array->length(); ++i) { + PyObject* value = checked(object_value(array, i)); + Py_XDECREF(values[offset + i]); + values[offset + i] = value; + } + return; + } + switch (array->type_id()) { +#define AK_COPY(ID, CLASS, CTYPE) \ + case arrow::Type::ID: { \ + const auto& values = static_cast(*array); \ + if (dtype == NPY_FLOAT64) { \ + auto target = reinterpret_cast(data) + offset; \ + for (int64_t i = 0; i < array->length(); ++i) \ + target[i] = array->IsNull(i) ? Py_NAN : static_cast(values.Value(i)); \ + } else { \ + std::memcpy(data + offset * sizeof(CTYPE), values.raw_values(), \ + static_cast(array->length()) * sizeof(CTYPE)); \ + } \ + break; \ + } + AK_COPY(INT8, Int8Array, int8_t) + AK_COPY(INT16, Int16Array, int16_t) + AK_COPY(INT32, Int32Array, int32_t) + AK_COPY(INT64, Int64Array, int64_t) + AK_COPY(UINT8, UInt8Array, uint8_t) + AK_COPY(UINT16, UInt16Array, uint16_t) + AK_COPY(UINT32, UInt32Array, uint32_t) + AK_COPY(UINT64, UInt64Array, uint64_t) +#undef AK_COPY + case arrow::Type::HALF_FLOAT: { + const auto& values = static_cast(*array); + auto target = reinterpret_cast(data) + offset; + for (int64_t i = 0; i < array->length(); ++i) + target[i] = array->IsNull(i) ? 0x7e00 : values.Value(i); + break; + } + case arrow::Type::FLOAT: { + const auto& values = static_cast(*array); + auto target = reinterpret_cast(data) + offset; + for (int64_t i = 0; i < array->length(); ++i) + target[i] = array->IsNull(i) ? Py_NAN : values.Value(i); + break; + } + case arrow::Type::DOUBLE: { + const auto& values = static_cast(*array); + auto target = reinterpret_cast(data) + offset; + for (int64_t i = 0; i < array->length(); ++i) + target[i] = array->IsNull(i) ? Py_NAN : values.Value(i); + break; + } + case arrow::Type::BOOL: { + const auto& values = static_cast(*array); + auto target = reinterpret_cast(data) + offset; + for (int64_t i = 0; i < array->length(); ++i) target[i] = values.Value(i); + break; + } + case arrow::Type::DATE32: { + const auto& values = static_cast(*array); + auto target = reinterpret_cast(data) + offset; + for (int64_t i = 0; i < array->length(); ++i) + target[i] = array->IsNull(i) ? NPY_DATETIME_NAT : values.Value(i); + break; + } + case arrow::Type::DATE64: + case arrow::Type::TIMESTAMP: { + auto target = reinterpret_cast(data) + offset; + auto source = array->data()->GetValues(1); + for (int64_t i = 0; i < array->length(); ++i) + target[i] = array->IsNull(i) ? NPY_DATETIME_NAT : source[i]; + break; + } + default: + PyErr_SetString(PyExc_NotImplementedError, "unsupported numeric Arrow type"); + throw PythonError(); + } +} + +PyObject* convert_column(const arrow::ChunkedArray& column) { + int dtype = numpy_type(*column.type(), column.null_count() != 0); + if (column.length() > NPY_MAX_INTP) { + PyErr_SetString(PyExc_OverflowError, "column is too large for NumPy"); + throw PythonError(); + } + npy_intp size = static_cast(column.length()); + PyArray_Descr* descriptor = dtype == NPY_DATETIME + ? datetime_dtype(*column.type()) : PyArray_DescrFromType(dtype); + if (!descriptor) throw PythonError(); + PyRef result(checked(PyArray_NewFromDescr(&PyArray_Type, descriptor, 1, &size, + nullptr, nullptr, 0, nullptr))); + auto output = reinterpret_cast(result.get()); + int64_t offset = 0; + for (const auto& chunk : column.chunks()) { + copy_chunk(chunk, output, offset, dtype); + offset += chunk->length(); + } + PyArray_CLEARFLAGS(output, NPY_ARRAY_WRITEABLE); + return result.release(); +} + +PyObject* convert_array(const std::shared_ptr& array) { + return convert_column(arrow::ChunkedArray(array)); +} + +PyObject* read_parquet(PyObject*, PyObject* args, PyObject* kwargs) { + try { + PyObject* source; + PyObject* selection = Py_None; + static const char* keywords[] = {"file", "columns_select", nullptr}; + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|$O:parquet_to_arrays", + const_cast(keywords), &source, &selection)) return nullptr; + + std::vector selected; + if (selection != Py_None) { + if (!PyList_Check(selection)) { + PyErr_SetString(PyExc_TypeError, "columns_select must be a list of strings or None"); + throw PythonError(); + } + std::unordered_set seen; + for (Py_ssize_t i = 0; i < PyList_GET_SIZE(selection); ++i) { + PyObject* name = PyList_GET_ITEM(selection, i); + if (!PyUnicode_Check(name)) { + PyErr_SetString(PyExc_TypeError, "columns_select must contain only strings"); + throw PythonError(); + } + Py_ssize_t length; + const char* text = PyUnicode_AsUTF8AndSize(name, &length); + if (!text) throw PythonError(); + std::string value(text, length); + if (!seen.insert(value).second) { + PyErr_Format(PyExc_ValueError, "duplicate column selection: %R", name); + throw PythonError(); + } + selected.push_back(std::move(value)); + } + } + + std::shared_ptr file; + std::shared_ptr python_file; + PyRef path(PyOS_FSPath(source)); + if (path.get()) { + PyRef encoded(PyUnicode_Check(path.get()) + ? PyUnicode_EncodeFSDefault(path.get()) : Py_NewRef(path.get())); + checked(encoded.get()); + char* data; + Py_ssize_t size; + if (PyBytes_AsStringAndSize(encoded.get(), &data, &size) < 0) throw PythonError(); + if (std::memchr(data, '\0', size)) { + PyErr_SetString(PyExc_ValueError, "embedded null byte in path"); + throw PythonError(); + } + arrow::Result> opened; + { + ReleaseGIL released; + opened = arrow::io::ReadableFile::Open(std::string(data, size)); + } + file = unwrap(std::move(opened)); + } else { + if (!PyErr_ExceptionMatches(PyExc_TypeError)) throw PythonError(); + PyErr_Clear(); + for (const char* name : {"read", "seek", "tell"}) { + PyRef method(PyObject_GetAttrString(source, name)); + if (!method.get()) { + PyErr_Clear(); + PyErr_SetString(PyExc_TypeError, "file must be a path or seekable binary file"); + throw PythonError(); + } + if (!PyCallable_Check(method.get())) { + PyErr_SetString(PyExc_TypeError, "file read, seek and tell must be callable"); + throw PythonError(); + } + } + python_file = std::make_shared(source); + file = python_file; + } + + arrow::Result> opened; + { + ReleaseGIL released; + opened = parquet::arrow::OpenFile(file, arrow::default_memory_pool()); + } + if (python_file && python_file->restore_error()) throw PythonError(); + auto reader = unwrap(std::move(opened)); + reader->set_use_threads(!python_file); + std::shared_ptr schema; + check_status(reader->GetSchema(&schema)); + + std::vector fields; + std::vector leaves; + if (selection == Py_None) { + for (int i = 0; i < schema->num_fields(); ++i) fields.push_back(i); + } else { + for (const auto& name : selected) { + auto matches = schema->GetAllFieldIndices(name); + if (matches.size() != 1) { + PyErr_SetString(PyExc_ValueError, "column name is missing or ambiguous"); + throw PythonError(); + } + fields.push_back(matches[0]); + } + } + const auto* parquet_schema = reader->parquet_reader()->metadata()->schema(); + for (int field : fields) { + const auto* root = parquet_schema->group_node()->field(field).get(); + for (int leaf = 0; leaf < parquet_schema->num_columns(); ++leaf) { + if (parquet_schema->GetColumnRoot(leaf) == root) leaves.push_back(leaf); + } + } + if (fields.empty()) return PyList_New(0); + + std::sort(leaves.begin(), leaves.end()); + arrow::Result> table_result; + { + ReleaseGIL released; + table_result = reader->ReadTable(leaves); + } + if (python_file && python_file->restore_error()) throw PythonError(); + auto table = unwrap(std::move(table_result)); + arrow::Status validation; + { + ReleaseGIL released; + validation = table->ValidateFull(); + } + check_status(validation); + // Parquet returns projected fields in file order, not request order. + PyRef result(checked(PyList_New(static_cast(fields.size())))); + std::vector ordered = fields; + std::sort(ordered.begin(), ordered.end()); + for (size_t i = 0; i < fields.size(); ++i) { + auto position = std::lower_bound(ordered.begin(), ordered.end(), fields[i]) + - ordered.begin(); + PyObject* array = checked(convert_column(*table->column(position))); + PyList_SET_ITEM(result.get(), static_cast(i), array); + } + return result.release(); + } catch (const PythonError&) { + if (!PyErr_Occurred()) PyErr_SetString(PyExc_RuntimeError, "Parquet conversion failed"); + } catch (const std::bad_alloc&) { + PyErr_NoMemory(); + } catch (const std::exception& error) { + if (!PyErr_Occurred()) PyErr_SetString(PyExc_RuntimeError, error.what()); + } catch (...) { + if (!PyErr_Occurred()) PyErr_SetString(PyExc_RuntimeError, "unknown native Parquet exception"); + } + return nullptr; +} + +PyMethodDef methods[] = { + {"parquet_to_arrays", reinterpret_cast(read_parquet), + METH_VARARGS | METH_KEYWORDS, + "Read Parquet columns into immutable one-dimensional NumPy arrays."}, + {nullptr, nullptr, 0, nullptr} +}; + +PyModuleDef module = { + PyModuleDef_HEAD_INIT, "_parquet", nullptr, -1, methods, + nullptr, nullptr, nullptr, nullptr +}; + +} // namespace + +PyMODINIT_FUNC PyInit__parquet() { + if (_import_array() < 0) return nullptr; + return PyModule_Create(&module); +} diff --git a/src/parquet_to_arrays.c b/src/parquet_to_arrays.c index 3b2dfe95..0e3a3b83 100644 --- a/src/parquet_to_arrays.c +++ b/src/parquet_to_arrays.c @@ -1,183 +1,35 @@ # include "Python.h" -# define PY_ARRAY_UNIQUE_SYMBOL AK_ARRAY_API -# define NO_IMPORT_ARRAY -# define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION -# include "numpy/arrayobject.h" - # include "parquet_to_arrays.h" PyObject * parquet_to_arrays(PyObject *Py_UNUSED(self), PyObject *args, PyObject *kwargs) { - PyObject *source; - PyObject *columns_select = Py_None; - static char *kwlist[] = {"file", "columns_select", NULL}; - if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|$O:parquet_to_arrays", - kwlist, &source, &columns_select)) { - return NULL; - } - - PyObject *selection = NULL; - PyObject *source_normalized = NULL; - PyObject *module = NULL; - PyObject *reader = NULL; - PyObject *paths = NULL; - PyObject *indices = NULL; - PyObject *schema = NULL; - PyObject *names = NULL; - PyObject *lookup = NULL; - PyObject *seen = NULL; - PyObject *table = NULL; - PyObject *result = NULL; - PyObject *value = NULL; - PyObject *column = NULL; - PyObject *array = NULL; - - if (columns_select != Py_None) { - if (!PyList_Check(columns_select)) { - PyErr_SetString(PyExc_TypeError, - "columns_select must be a list of strings or None"); - goto error; - } - selection = PyList_GetSlice(columns_select, 0, PyList_GET_SIZE(columns_select)); - if (selection == NULL) goto error; - seen = PySet_New(NULL); - if (seen == NULL) goto error; - for (Py_ssize_t i = 0; i < PyList_GET_SIZE(selection); ++i) { - PyObject *name = PyList_GET_ITEM(selection, i); - if (!PyUnicode_Check(name)) { - PyErr_SetString(PyExc_TypeError, - "columns_select must contain only strings"); - goto error; - } - int contains = PySet_Contains(seen, name); - if (contains < 0) goto error; - if (contains) { - PyErr_Format(PyExc_ValueError, "duplicate column selection: %R", name); - goto error; - } - if (PySet_Add(seen, name) < 0) goto error; - } - } - - source_normalized = PyOS_FSPath(source); - if (source_normalized == NULL) { - if (!PyErr_ExceptionMatches(PyExc_TypeError)) goto error; - PyErr_Clear(); - source_normalized = Py_NewRef(source); - } - if (PyBytes_Check(source_normalized)) { - value = PyUnicode_DecodeFSDefaultAndSize( - PyBytes_AS_STRING(source_normalized), PyBytes_GET_SIZE(source_normalized)); - if (value == NULL) goto error; - Py_SETREF(source_normalized, value); - value = NULL; - } - module = PyImport_ImportModule("pyarrow._parquet"); - if (module == NULL) goto error; - value = PyObject_GetAttrString(module, "ParquetReader"); - if (value == NULL) goto error; - reader = PyObject_CallNoArgs(value); - Py_CLEAR(value); - if (reader == NULL) goto error; - value = PyObject_CallMethod(reader, "open", "(O)", source_normalized); - if (value == NULL) goto error; - Py_CLEAR(value); - - if (selection != NULL) { - paths = PyObject_GetAttrString(reader, "column_paths"); - schema = PyObject_GetAttrString(reader, "schema_arrow"); - if (paths == NULL || schema == NULL) goto error; - names = PyObject_GetAttrString(schema, "names"); - if (names == NULL) goto error; - indices = PyList_New(0); - if (indices == NULL) goto error; - lookup = PyDict_New(); - if (lookup == NULL) goto error; - for (Py_ssize_t i = 0; i < PyList_Size(names); ++i) { - PyObject *name = PyList_GetItem(names, i); - int contains = PyDict_Contains(lookup, name); - if (contains < 0) goto error; - value = contains ? Py_NewRef(Py_None) : PyList_New(0); - if (value == NULL || PyDict_SetItem(lookup, name, value) < 0) goto error; - Py_CLEAR(value); - } - // Parquet indices refer to leaves, not top-level Arrow fields. - for (Py_ssize_t i = 0; i < PyList_Size(paths); ++i) { - PyObject *path = PyList_GetItem(paths, i); - PyObject *root = PyList_GetItem(path, 0); - if (root == NULL) goto error; - PyObject *leaves = PyDict_GetItemWithError(lookup, root); - if (leaves == NULL) { - if (!PyErr_Occurred()) { - PyErr_SetString(PyExc_ValueError, "inconsistent Parquet schema"); - } - goto error; - } - if (leaves != Py_None) { - value = PyLong_FromSsize_t(i); - if (value == NULL || PyList_Append(leaves, value) < 0) goto error; - Py_CLEAR(value); + PyObject *module = PyImport_ImportModule("arraykit._parquet"); + if (module == NULL) { + if (PyErr_ExceptionMatches(PyExc_ModuleNotFoundError)) { + PyObject *error = PyErr_GetRaisedException(); + PyObject *name = PyObject_GetAttrString(error, "name"); + int missing_backend = name != NULL && PyUnicode_Check(name) + && PyUnicode_CompareWithASCIIString(name, "arraykit._parquet") == 0; + Py_XDECREF(name); + PyErr_Clear(); + if (missing_backend) { + Py_DECREF(error); + PyErr_SetString(PyExc_ImportError, + "Native Parquet support is not installed; rebuild with " + "ARRAYKIT_WITH_PARQUET=1 and Arrow/Parquet C++ libraries"); } - } - for (Py_ssize_t i = 0; i < PyList_GET_SIZE(selection); ++i) { - PyObject *name = PyList_GET_ITEM(selection, i); - PyObject *leaves = PyDict_GetItemWithError(lookup, name); - if (leaves == NULL && PyErr_Occurred()) goto error; - if (leaves == NULL || leaves == Py_None) { - PyErr_Format(PyExc_ValueError, - "column name is missing or ambiguous: %R", name); - goto error; + else { + PyErr_SetRaisedException(error); } - Py_ssize_t size = PyList_GET_SIZE(indices); - if (PyList_SetSlice(indices, size, size, leaves) < 0) goto error; } + return NULL; } - - // Both reading and NumPy conversion are compiled PyArrow entry points. - table = PyObject_CallMethod(reader, "read_all", "(O)", - indices == NULL ? Py_None : indices); - if (table == NULL) goto error; - value = PyObject_GetAttrString(table, "num_columns"); - if (value == NULL) goto error; - Py_ssize_t count = PyLong_AsSsize_t(value); - Py_CLEAR(value); - if (count == -1 && PyErr_Occurred()) goto error; - result = PyList_New(count); - if (result == NULL) goto error; - for (Py_ssize_t i = 0; i < count; ++i) { - column = PyObject_CallMethod(table, "column", "n", i); - if (column == NULL) goto error; - array = PyObject_CallMethod(column, "to_numpy", NULL); - Py_CLEAR(column); - if (array == NULL) goto error; - if (!PyArray_Check(array) || PyArray_NDIM((PyArrayObject *)array) != 1) { - PyErr_SetString(PyExc_TypeError, "Parquet column did not produce a 1D array"); - goto error; - } - PyArray_CLEARFLAGS((PyArrayObject *)array, NPY_ARRAY_WRITEABLE); - PyList_SET_ITEM(result, i, array); - array = NULL; - } - goto done; - -error: - Py_CLEAR(result); -done: - Py_XDECREF(array); - Py_XDECREF(column); - Py_XDECREF(value); - Py_XDECREF(table); - Py_XDECREF(names); - Py_XDECREF(schema); - Py_XDECREF(lookup); - Py_XDECREF(seen); - Py_XDECREF(indices); - Py_XDECREF(paths); - Py_XDECREF(reader); - Py_XDECREF(module); - Py_XDECREF(selection); - Py_XDECREF(source_normalized); + PyObject *function = PyObject_GetAttrString(module, "parquet_to_arrays"); + Py_DECREF(module); + if (function == NULL) return NULL; + PyObject *result = PyObject_Call(function, args, kwargs); + Py_DECREF(function); return result; } diff --git a/test/test_parquet_build.py b/test/test_parquet_build.py new file mode 100644 index 00000000..b572dd61 --- /dev/null +++ b/test/test_parquet_build.py @@ -0,0 +1,68 @@ +from contextlib import chdir +import os +from pathlib import Path +import runpy +import unittest +from unittest import mock + + +class TestUnit(unittest.TestCase): + @staticmethod + def get_extensions(environment, output=None): + path = Path(__file__).resolve().parents[1] / 'setup.py' + with mock.patch.dict(os.environ, environment, clear=True), chdir(path.parent): + with mock.patch('setuptools.setup') as setup: + with mock.patch('subprocess.check_call'): + with mock.patch('subprocess.check_output', side_effect=output): + runpy.run_path(str(path)) + return setup.call_args.kwargs['ext_modules'] + + def test_parquet_build_default(self): + self.assertEqual(len(self.get_extensions({})), 1) + + def test_parquet_build_prefix(self): + extensions = self.get_extensions( + {'ARRAYKIT_WITH_PARQUET': '1', 'ARROW_HOME': '/opt/arrow'} + ) + native = extensions[-1] + self.assertEqual(native.name, 'arraykit._parquet') + self.assertIn('/opt/arrow/include', native.include_dirs) + self.assertEqual(native.library_dirs, ['/opt/arrow/lib']) + self.assertEqual(native.libraries, ['parquet', 'arrow']) + self.assertEqual(native.language, 'c++') + + def test_parquet_build_explicit_directories(self): + extensions = self.get_extensions( + { + 'ARRAYKIT_WITH_PARQUET': '1', + 'ARROW_INCLUDE_DIR': '/native/include', + 'ARROW_LIBRARY_DIR': '/native/lib64', + } + ) + self.assertIn('/native/include', extensions[-1].include_dirs) + self.assertEqual(extensions[-1].library_dirs, ['/native/lib64']) + + def test_parquet_build_incomplete_directories(self): + with self.assertRaisesRegex(RuntimeError, 'Set both'): + self.get_extensions( + {'ARRAYKIT_WITH_PARQUET': '1', 'ARROW_INCLUDE_DIR': '/native/include'} + ) + + def test_parquet_build_pkg_config(self): + extensions = self.get_extensions( + {'ARRAYKIT_WITH_PARQUET': '1'}, + output=['-I/native/include', '-L/native/lib -lparquet -larrow'], + ) + native = extensions[-1] + self.assertIn('-I/native/include', native.extra_compile_args) + self.assertIn('-L/native/lib', native.extra_link_args) + + def test_parquet_build_missing_libraries(self): + with self.assertRaisesRegex(RuntimeError, 'Native Parquet support requires'): + self.get_extensions( + {'ARRAYKIT_WITH_PARQUET': '1'}, output=FileNotFoundError() + ) + + +if __name__ == '__main__': + unittest.main() diff --git a/test/test_parquet_to_arrays.py b/test/test_parquet_to_arrays.py index 2fc71ec8..267812a8 100644 --- a/test/test_parquet_to_arrays.py +++ b/test/test_parquet_to_arrays.py @@ -1,8 +1,12 @@ import gc +from decimal import Decimal +import importlib.util import io import os from pathlib import Path import tempfile +import subprocess +import sys import unittest from unittest import mock @@ -18,7 +22,11 @@ pq = None -@unittest.skipIf(pa is None, 'PyArrow is not installed') +@unittest.skipIf(pa is None, 'PyArrow is not installed for fixture generation') +@unittest.skipIf( + importlib.util.find_spec('arraykit._parquet') is None, + 'Native Parquet backend is not built', +) class TestUnit(unittest.TestCase): @staticmethod def to_file(table, **kwargs): @@ -130,6 +138,22 @@ def test_parquet_to_arrays_dictionary(self): self.assertEqual(array.tolist(), ['a', 'b', 'a']) self.assertEqual(array.dtype, np.dtype(object)) + def test_parquet_to_arrays_nested_float16(self): + values = [ + {'x': 1.5}, {'x': None}, None, {'x': -0.0}, + {'x': 2**-24}, {'x': float('inf')}, {'x': float('nan')}, + ] + table = pa.table({'a': pa.array(values, type=pa.struct([('x', pa.float16())]))}) + array = parquet_to_arrays(self.to_file(table, row_group_size=1))[0] + self.assertEqual(array.dtype, np.dtype(object)) + self.assertFalse(array.flags.writeable) + self.assertEqual(array[:3].tolist(), values[:3]) + np.testing.assert_array_equal( + [value['x'] for value in array[3:]], + [-0.0, 2**-24, float('inf'), float('nan')], + ) + self.assertTrue(np.signbit(array[3]['x'])) + def test_parquet_to_arrays_wide_projection(self): names = [f'c{i}' for i in range(128)] table = pa.table({name: [i] for i, name in enumerate(names)}) @@ -157,16 +181,105 @@ def test_parquet_to_arrays_errors(self): parquet_to_arrays(Path(directory) / 'missing.parquet') def test_parquet_to_arrays_native_reader(self): - file = self.to_file(pa.table({'a': [1, 2]})) - with mock.patch.object(pq, 'read_table', side_effect=AssertionError): - with mock.patch.object(pq, 'ParquetFile', side_effect=AssertionError): - np.testing.assert_array_equal(parquet_to_arrays(file)[0], [1, 2]) + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / 'data.parquet' + pq.write_table(pa.table({'a': [1, 2], 'b': ['x', 'y']}), path) + code = ''' +import importlib.abc +import io +import sys + +class NoPyArrow(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname == 'pyarrow' or fullname.startswith('pyarrow.'): + raise AssertionError('Reader must not import Python PyArrow') + return None + +sys.meta_path.insert(0, NoPyArrow()) +from arraykit import parquet_to_arrays +arrays = parquet_to_arrays(sys.argv[1], columns_select=['b', 'a']) +assert [a.tolist() for a in arrays] == [['x', 'y'], [1, 2]] +with open(sys.argv[1], 'rb') as file: + arrays = parquet_to_arrays(io.BytesIO(file.read())) +assert [a.tolist() for a in arrays] == [[1, 2], ['x', 'y']] +assert not any(name.startswith('pyarrow') for name in sys.modules) +''' + subprocess.run([sys.executable, '-c', code, str(path)], check=True) + + def test_parquet_to_arrays_file_errors(self): + class ReadError(io.BytesIO): + def read(self, size=-1): + raise RuntimeError('custom read error') + + class SeekError(io.BytesIO): + def seek(self, offset, whence=0): + raise RuntimeError('custom seek error') + + data = self.to_file(pa.table({'a': [1]})).getvalue() + for cls, message in ( + (ReadError, 'custom read error'), + (SeekError, 'custom seek error'), + ): + with self.subTest(cls=cls): + with self.assertRaisesRegex(RuntimeError, message): + parquet_to_arrays(cls(data)) + with self.assertRaises(ValueError): + parquet_to_arrays('data.parquet\x00ignored') + + def test_parquet_to_arrays_dates(self): + table = pa.table( + { + 'date': pa.array([0, None, 100], type=pa.date32()), + 'timestamp_ms': pa.array([0, None, 100], type=pa.timestamp('ms')), + 'timestamp_ns': pa.array([0, None, 100], type=pa.timestamp('ns')), + } + ) + arrays = parquet_to_arrays(self.to_file(table, row_group_size=1)) + for array, expected, unit in zip( + arrays, ([0, 'NaT', 100],) * 3, ('D', 'ms', 'ns') + ): + np.testing.assert_array_equal( + array, np.array(expected, dtype=f'datetime64[{unit}]') + ) + + def test_parquet_to_arrays_nullable_numeric_types(self): + types = ( + pa.int8(), pa.int16(), pa.int32(), pa.int64(), + pa.uint8(), pa.uint16(), pa.uint32(), pa.uint64(), + pa.float16(), pa.float32(), pa.float64(), + ) + for dtype in types: + with self.subTest(dtype=dtype): + table = pa.table({'a': pa.array([1, None, 3], type=dtype)}) + array = parquet_to_arrays(self.to_file(table, row_group_size=1))[0] + expected_dtype = ( + np.float64 if pa.types.is_integer(dtype) else dtype.to_pandas_dtype() + ) + np.testing.assert_array_equal(array, [1, np.nan, 3]) + self.assertEqual(array.dtype, np.dtype(expected_dtype)) + + def test_parquet_to_arrays_decimal(self): + values = [Decimal('1.23'), None, Decimal('-2.50')] + table = pa.table({'a': pa.array(values, type=pa.decimal128(8, 2))}) + array = parquet_to_arrays(self.to_file(table))[0] + self.assertEqual(array.tolist(), values) + self.assertFalse(array.flags.writeable) + + def test_parquet_to_arrays_short_reads(self): + class ShortRead(io.BytesIO): + def read(self, size=-1): + return super().read(min(size, 7)) + + data = self.to_file(pa.table({'a': [1, 2]})).getvalue() + file = ShortRead(data) + self.assertEqual(parquet_to_arrays(file)[0].tolist(), [1, 2]) + self.assertFalse(file.closed) class TestOptionalDependency(unittest.TestCase): - def test_parquet_to_arrays_missing_pyarrow(self): - with mock.patch.dict('sys.modules', {'pyarrow._parquet': None}): - with self.assertRaises(ImportError): + def test_parquet_to_arrays_missing_backend(self): + with mock.patch.dict('sys.modules', {'arraykit._parquet': None}): + with self.assertRaisesRegex(ImportError, 'ARRAYKIT_WITH_PARQUET'): parquet_to_arrays(io.BytesIO())