diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b2e61525..cf3f41fd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -25,6 +25,29 @@ jobs: name: dist-sdist path: dist/* + parquet: + name: Native Parquet / ubuntu-24.04 / Python 3.14 + runs-on: ubuntu-24.04 + steps: + - uses: actions/checkout@master + - uses: actions/setup-python@master + with: + python-version: 3.14 + - run: echo '::add-matcher::.github/problem-matchers/gcc.json' + - name: Install Arrow/Parquet C++ libraries + run: | + wget -q https://packages.apache.org/artifactory/arrow/ubuntu/apache-arrow-apt-source-latest-noble.deb + sudo apt-get install -y ./apache-arrow-apt-source-latest-noble.deb + sudo apt-get update + sudo apt-get install -y libarrow-dev libparquet-dev + pkg-config --modversion arrow parquet + - run: python -m pip install -r requirements-dev-3_14.txt pyarrow + - run: python -m pip install -v . + env: + ARRAYKIT_WITH_PARQUET: '1' + - run: python -c "import arraykit._parquet" + - run: python -m pytest test/test_parquet_to_arrays.py test/test_parquet_build.py + matrix_config: name: Matrix Runner Config runs-on: ubuntu-22.04 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 1ee56c47..7d60e9aa 100644 --- a/README.rst +++ b/README.rst @@ -32,6 +32,70 @@ 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. 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 +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 native extension reads paths through Arrow's random-access file API. +Open Python files use binary ``read`` (returning any bytes-like object), +``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``, timestamps and dates use ``datetime64`` with ``NaT``, and +durations use ``timedelta64`` with ``NaT``. Timestamp time zones are dropped; +values are UTC. Large nullable integers can lose precision when converted to +floating point. Dictionary columns are decoded to values, decimal columns use +``decimal.Decimal`` objects, time-of-day columns use ``datetime.time`` objects +(nanosecond values that are not whole microseconds raise ``ValueError``), and +nested struct/list columns use object arrays. Temporal values nested in structs +use NumPy ``datetime64`` or ``timedelta64`` scalars. 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/performance/parquet/fixtures.py b/performance/parquet/fixtures.py new file mode 100644 index 00000000..ad499910 --- /dev/null +++ b/performance/parquet/fixtures.py @@ -0,0 +1,125 @@ +import atexit +import io +import os +import shutil +import tempfile +import typing as tp + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq + + +COLUMNS = 10 +NULL_INTERVAL = 10 # every n-th value is null in nullable fixtures + +DIRECTORY = tempfile.mkdtemp(prefix='arraykit-parquet-') +atexit.register(shutil.rmtree, DIRECTORY, ignore_errors=True) + + +class PayLoad: + def __init__(self, label: str, size: int, table: pa.Table): + self.table = table + self.size = size + self.path = os.path.join(DIRECTORY, f'{label}-{size}.parquet') + pq.write_table(table, self.path) + with open(self.path, 'rb') as f: + self.data = f.read() + self.columns = table.column_names + self.columns_select = self.columns[::2] + + def bytes_io(self) -> io.BytesIO: + return io.BytesIO(self.data) + + +# ------------------------------------------------------------------------------- +def get_table(get_array: tp.Callable[[int, int], pa.Array], size: int) -> pa.Table: + return pa.table({f'c{i}': get_array(size, i) for i in range(COLUMNS)}) + + +def get_null_mask(size: int) -> np.ndarray: + return (np.arange(size) % NULL_INTERVAL) == 0 + + +class FixtureFactory: + NAME = '' + SORT = 0 + CACHE = {} # can be shared for all classes + + @staticmethod + def get_array(size: int, column: int) -> pa.Array: + raise NotImplementedError() + + @classmethod + def get_label_payload(cls, size: int) -> tp.Tuple[str, PayLoad]: + key = (cls, size) + if key not in cls.CACHE: + table = get_table(cls.get_array, size) + cls.CACHE[key] = PayLoad(cls.NAME, size, table) + return cls.NAME, cls.CACHE[key] + + +class FFInt64(FixtureFactory): + NAME = 'int64' + SORT = 0 + + @staticmethod + def get_array(size: int, column: int) -> pa.Array: + return pa.array(np.arange(column, column + size, dtype=np.int64)) + + +class FFInt64Null(FixtureFactory): + NAME = 'int64 (nulls)' + SORT = 1 + + @staticmethod + def get_array(size: int, column: int) -> pa.Array: + values = np.arange(column, column + size, dtype=np.int64) + return pa.array(values, mask=get_null_mask(size)) + + +class FFFloat64(FixtureFactory): + NAME = 'float64' + SORT = 2 + + @staticmethod + def get_array(size: int, column: int) -> pa.Array: + return pa.array(np.arange(column, column + size, dtype=np.float64) * 0.5) + + +class FFBool(FixtureFactory): + NAME = 'bool' + SORT = 3 + + @staticmethod + def get_array(size: int, column: int) -> pa.Array: + return pa.array((np.arange(size) + column) % 3 == 0) + + +class FFDTns(FixtureFactory): + NAME = 'dt64[ns]' + SORT = 4 + + @staticmethod + def get_array(size: int, column: int) -> pa.Array: + values = np.arange(size, dtype=np.int64) * 1_000_000_007 + column + return pa.array(values.astype('datetime64[ns]')) + + +class FFString(FixtureFactory): + NAME = 'string' + SORT = 5 + + @staticmethod + def get_array(size: int, column: int) -> pa.Array: + return pa.array([hex(e) * 2 for e in range(column, column + size)]) + + +class FFStringDict(FixtureFactory): + NAME = 'string (dict)' + SORT = 6 + + @staticmethod + def get_array(size: int, column: int) -> pa.Array: + labels = np.array([f'label-{i}' for i in range(100)], dtype=object) + return pa.array(labels[(np.arange(size) + column) % 100]).dictionary_encode() diff --git a/performance/parquet/parquet-to-arrays.py b/performance/parquet/parquet-to-arrays.py new file mode 100644 index 00000000..01cb8ddf --- /dev/null +++ b/performance/parquet/parquet-to-arrays.py @@ -0,0 +1,301 @@ +import os +import sys +import timeit + +import arraykit +from arraykit import parquet_to_arrays + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq + +sys.path.append(os.getcwd()) + +from fixtures import ( + PayLoad, + COLUMNS, + FFInt64, + FFInt64Null, + FFFloat64, + FFBool, + FFDTns, + FFString, + FFStringDict, +) + + +class ParquetProcessor: + NAME = '' + SORT = -1 + + def __init__(self, pl: PayLoad): + self.pl = pl + self.path = pl.path + self.size = pl.size + self.columns_select = pl.columns_select + + +def table_to_arrays(table: pa.Table) -> list: + # ChunkedArray.to_numpy() applies the same null conversions as ArrayKit: + # nullable integers to float64 with NaN, nullable bool and strings to object. + return [column.to_numpy() for column in table.columns] + + +# ------------------------------------------------------------------------------- +class AKPath(ParquetProcessor): + NAME = 'AK: path' + SORT = 0 + + def __call__(self): + arrays = parquet_to_arrays(self.path) + assert len(arrays[0]) == self.size + + +class AKBytesIO(ParquetProcessor): + NAME = 'AK: BytesIO' + SORT = 1 + + def __call__(self): + arrays = parquet_to_arrays(self.pl.bytes_io()) + assert len(arrays[0]) == self.size + + +class PAPath(ParquetProcessor): + NAME = 'PA: path' + SORT = 2 + + def __call__(self): + arrays = table_to_arrays(pq.read_table(self.path)) + assert len(arrays[0]) == self.size + + +class PAPathSingleThread(ParquetProcessor): + NAME = 'PA: path, 1 thread' + SORT = 3 + + def __call__(self): + arrays = table_to_arrays(pq.read_table(self.path, use_threads=False)) + assert len(arrays[0]) == self.size + + +class PABytesIO(ParquetProcessor): + NAME = 'PA: BytesIO' + SORT = 4 + + def __call__(self): + arrays = table_to_arrays(pq.read_table(self.pl.bytes_io())) + assert len(arrays[0]) == self.size + + +# ------------------------------------------------------------------------------- +class AKPathSelect(ParquetProcessor): + NAME = 'AK: path select' + SORT = 0 + + def __call__(self): + arrays = parquet_to_arrays(self.path, columns_select=self.columns_select) + assert len(arrays) == len(self.columns_select) + + +class AKBytesIOSelect(ParquetProcessor): + NAME = 'AK: BytesIO select' + SORT = 1 + + def __call__(self): + arrays = parquet_to_arrays( + self.pl.bytes_io(), columns_select=self.columns_select + ) + assert len(arrays) == len(self.columns_select) + + +class PAPathSelect(ParquetProcessor): + NAME = 'PA: path select' + SORT = 2 + + def __call__(self): + table = pq.read_table(self.path, columns=self.columns_select) + arrays = table_to_arrays(table) + assert len(arrays) == len(self.columns_select) + + +class PAPathSingleThreadSelect(ParquetProcessor): + NAME = 'PA: path select, 1 thread' + SORT = 3 + + def __call__(self): + table = pq.read_table( + self.path, columns=self.columns_select, use_threads=False + ) + arrays = table_to_arrays(table) + assert len(arrays) == len(self.columns_select) + + +class PABytesIOSelect(ParquetProcessor): + NAME = 'PA: BytesIO select' + SORT = 4 + + def __call__(self): + table = pq.read_table(self.pl.bytes_io(), columns=self.columns_select) + arrays = table_to_arrays(table) + assert len(arrays) == len(self.columns_select) + + +def get_versions() -> str: + import platform + + return ( + f'OS: {platform.system()} / ArrayKit: {arraykit.__version__} / ' + f'NumPy: {np.__version__} / PyArrow: {pa.__version__} / ' + f'Columns: {COLUMNS}\n' + ) + + +CLS_FF = ( + FFInt64, + FFInt64Null, + FFFloat64, + FFBool, + FFDTns, + FFString, + FFStringDict, +) +FF_ORDER = [f.NAME for f in sorted(CLS_FF, key=lambda ff: ff.SORT)] + +# ------------------------------------------------------------------------------- +NUMBER = 5 + + +def seconds_to_display(seconds: float) -> str: + seconds /= NUMBER + if seconds < 1e-4: + return f'{seconds * 1e6: .1f} (µs)' + if seconds < 1e-1: + return f'{seconds * 1e3: .1f} (ms)' + return f'{seconds: .1f} (s)' + + +def plot_performance(frame, suffix: str = ''): + fixture_total = len(frame['fixture'].unique()) + cat_total = len(frame['size'].unique()) + processor_total = len(frame['cls_processor'].unique()) + fig, axes = plt.subplots(cat_total, fixture_total) + + cmap = plt.get_cmap('plasma') + color = cmap(np.arange(processor_total) / processor_total) + + # category is the number of rows + for cat_count, (cat_label, cat) in enumerate(frame.groupby('size')): + # fixture is the column type fixture + fixture_data = {fix_label: fix for fix_label, fix in cat.groupby('fixture')} + for fixture_count, fixture_label in enumerate(FF_ORDER): + fixture = fixture_data[fixture_label] + ax = axes[cat_count][fixture_count] + + # set order by cls_processor, i.e., the type of test being done + fixture['sort'] = [f.SORT for f in fixture['cls_processor']] + fixture = fixture.sort_values('sort') + + results = fixture['time'].values.tolist() + names = [cls.NAME for cls in fixture['cls_processor']] + names_display = names + post = ax.bar(names_display, results, color=color) + + title = f'{cat_label:.0e}\n{fixture_label}' + + ax.set_title(title, fontsize=6) + ax.set_box_aspect(0.8) + time_max = fixture['time'].max() + time_min = fixture['time'].min() + y_ticks = [0, time_min, time_max * 0.5, time_max] + y_labels = [ + '', + seconds_to_display(time_min), + seconds_to_display(time_max * 0.5), + seconds_to_display(time_max), + ] + if time_min > time_max * 0.25: + # remove the min if it is greater than quarter + y_ticks.pop(1) + y_labels.pop(1) + + ax.set_yticks(y_ticks) + ax.set_yticklabels(y_labels, fontsize=4) + ax.tick_params( + axis='x', + bottom=False, + labelbottom=False, + ) + ax.tick_params( + axis='y', + length=2, + width=0.5, + pad=1, + ) + fig.set_size_inches(10, 3.5) # width, height + fig.legend(post, names_display, loc='center right', fontsize=6) + # horizontal, vertical + fig.text( + 0.05, 0.96, f'Parquet to Arrays {suffix.title()}: {NUMBER} Iterations', fontsize=10 + ) + fig.text(0.05, 0.90, get_versions(), fontsize=6) + + fp = f'/tmp/arraykit-parquet-{suffix}.png' + plt.subplots_adjust( + left=0.075, + bottom=0.05, + right=0.82, + top=0.80, + wspace=1.0, # width + hspace=0.5, + ) + plt.savefig(fp, dpi=300) + + if sys.platform.startswith('linux'): + os.system(f'eog {fp}&') + else: + os.system(f'open {fp}') + + +def run_test(processors, suffix): + records = [] + for size in (10_000, 100_000, 1_000_000): + for ff in CLS_FF: + fixture_label, fixture = ff.get_label_payload(size) + for cls in processors: + runner = cls(fixture) + runner() # warm up file system cache and imports + + record = [cls, NUMBER, fixture_label, size] + print(record) + result = timeit.timeit(f'runner()', globals=locals(), number=NUMBER) + record.append(result) + records.append(record) + + f = pd.DataFrame.from_records( + records, columns=('cls_processor', 'number', 'fixture', 'size', 'time') + ) + print(f) + plot_performance(f, suffix) + + +if __name__ == '__main__': + cls_read = ( + AKPath, + AKBytesIO, + PAPath, + PAPathSingleThread, + PABytesIO, + ) + + cls_select = ( + AKPathSelect, + AKBytesIOSelect, + PAPathSelect, + PAPathSingleThreadSelect, + PABytesIOSelect, + ) + + run_test(cls_read, 'read') + run_test(cls_select, 'select') diff --git a/setup.py b/setup.py index 46e91660..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() @@ -33,6 +35,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', ], @@ -43,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/__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..9d228264 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 @@ -191,6 +192,11 @@ def delimited_to_arrays( thousandschar: str = ',', decimalchar: str = '.', ) -> tp.List[np.array]: ... +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]: ... def split_after_count( string: str, *, diff --git a/src/_arraykit.c b/src/_arraykit.c index 0d14b23f..4e4f2550 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 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..da53a90b --- /dev/null +++ b/src/parquet_reader.cpp @@ -0,0 +1,735 @@ +#define PY_SSIZE_T_CLEAN +#include +#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. +// Lock order is always mutex then GIL: never call into Arrow APIs that may do +// file I/O while holding the GIL, or a reader thread holding the mutex and +// waiting for the GIL can deadlock. +class PythonFile final : public arrow::io::RandomAccessFile { +public: + explicit PythonFile(PyObject* file) : file_(Py_NewRef(file)) {} + // Errors not reported through restore_error() are discarded: restoring one + // here could set an exception on a successful return. + ~PythonFile() override { + AcquireGIL gil; + Py_DECREF(file_); + Py_XDECREF(error_); + } + + 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_) return false; + PyErr_SetRaisedException(error_); + error_ = nullptr; + return true; + } + +private: + arrow::Status failure() const { + if (!error_) error_ = PyErr_GetRaisedException(); + 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(); + Py_buffer view; + if (PyObject_GetBuffer(value.get(), &view, PyBUF_SIMPLE) < 0) { + PyErr_SetString(PyExc_TypeError, "file.read() must return a bytes-like object"); + return failure(); + } + Py_ssize_t size = view.len; + if (size > count - total) { + PyBuffer_Release(&view); + PyErr_SetString(PyExc_ValueError, "file.read() returned too many bytes"); + return failure(); + } + if (size) std::memcpy(static_cast(out) + total, view.buf, + static_cast(size)); + PyBuffer_Release(&view); + if (!size) break; + total += size; + } + return total; + } + + PyObject* file_; + mutable std::mutex mutex_; + mutable PyObject* error_ = nullptr; +}; + +PyObject* decimal_type = nullptr; + +PyObject* convert_array(const std::shared_ptr& array); + +PyObject* time_value(int64_t value, arrow::TimeUnit::type unit) { + int64_t micros = value; + switch (unit) { + case arrow::TimeUnit::SECOND: micros = value * 1000000; break; + case arrow::TimeUnit::MILLI: micros = value * 1000; break; + case arrow::TimeUnit::MICRO: break; + case arrow::TimeUnit::NANO: + if (value % 1000) { + PyErr_Format(PyExc_ValueError, + "time value %lld has non-zero nanoseconds", static_cast(value)); + throw PythonError(); + } + micros = value / 1000; + break; + } + if (micros < 0 || micros >= 86400LL * 1000000) { + PyErr_Format(PyExc_ValueError, "time value %lld is out of range", + static_cast(value)); + throw PythonError(); + } + int64_t seconds = micros / 1000000; + return PyTime_FromTime(static_cast(seconds / 3600), + static_cast(seconds / 60 % 60), static_cast(seconds % 60), + static_cast(micros % 1000000)); +} + +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 value(checked(PyUnicode_FromStringAndSize(text.data(), text.size()))); + return PyObject_CallOneArg(decimal_type, value.get()); + } + // NumPy scalars keep every unit exactly; PyArray_GETITEM would return + // datetime objects for some units and ints for nanoseconds. + case arrow::Type::TIMESTAMP: + case arrow::Type::DATE32: + case arrow::Type::DATE64: + case arrow::Type::DURATION: { + PyRef values(checked(convert_array(array->Slice(index, 1)))); + auto values_array = reinterpret_cast(values.get()); + return PyArray_Scalar(PyArray_DATA(values_array), + PyArray_DESCR(values_array), values.get()); + } + case arrow::Type::TIME32: { + const auto& values = static_cast(*array); + return time_value(values.Value(index), + static_cast(*array->type()).unit()); + } + case arrow::Type::TIME64: { + const auto& values = static_cast(*array); + return time_value(values.Value(index), + static_cast(*array->type()).unit()); + } + 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::DURATION: return NPY_TIMEDELTA; + case arrow::Type::TIME32: + case arrow::Type::TIME64: + 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(); + } +} + +// Timestamp time zones are dropped: values are UTC and become naive datetime64. +PyArray_Descr* temporal_dtype(const arrow::DataType& type) { + const char* unit = "D"; + if (type.id() == arrow::Type::DATE64) unit = "ms"; + else if (type.id() == arrow::Type::TIMESTAMP || type.id() == arrow::Type::DURATION) { + auto time_unit = type.id() == arrow::Type::TIMESTAMP + ? static_cast(type).unit() + : static_cast(type).unit(); + switch (time_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; + } + } + const char* kind = type.id() == arrow::Type::DURATION ? "timedelta64[" : "datetime64["; + std::string spec = std::string(kind) + 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: + case arrow::Type::DURATION: { + 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 || dtype == NPY_TIMEDELTA + ? temporal_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(); + size_t count = leaves.size(); + for (int leaf = 0; leaf < parquet_schema->num_columns(); ++leaf) { + if (parquet_schema->GetColumnRoot(leaf) == root) leaves.push_back(leaf); + } + // Columns are matched to fields by position below, so every field + // must produce exactly one table column. + if (leaves.size() == count) { + PyErr_Format(PyExc_NotImplementedError, + "column %s has no leaf columns", root->name().c_str()); + throw PythonError(); + } + } + 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->Validate(); + } + 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; + PyDateTime_IMPORT; + if (!PyDateTimeAPI) return nullptr; + if (!decimal_type) { + PyObject* decimal = PyImport_ImportModule("decimal"); + if (!decimal) return nullptr; + decimal_type = PyObject_GetAttrString(decimal, "Decimal"); + Py_DECREF(decimal); + if (!decimal_type) return nullptr; + } + PyObject* m = PyModule_Create(&module); +#ifdef Py_GIL_DISABLED + if (m) PyUnstable_Module_SetGIL(m, Py_MOD_GIL_NOT_USED); +#endif + return m; +} diff --git a/src/parquet_to_arrays.c b/src/parquet_to_arrays.c new file mode 100644 index 00000000..0e3a3b83 --- /dev/null +++ b/src/parquet_to_arrays.c @@ -0,0 +1,35 @@ +# include "Python.h" + +# include "parquet_to_arrays.h" + +PyObject * +parquet_to_arrays(PyObject *Py_UNUSED(self), PyObject *args, PyObject *kwargs) +{ + 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"); + } + else { + PyErr_SetRaisedException(error); + } + } + return NULL; + } + 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/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_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 new file mode 100644 index 00000000..8556e5df --- /dev/null +++ b/test/test_parquet_to_arrays.py @@ -0,0 +1,349 @@ +import datetime +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 + +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 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): + 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_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)}) + 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): + 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_durations(self): + for unit in ('s', 'ms', 'us', 'ns'): + with self.subTest(unit=unit): + table = pa.table({'a': pa.array([5, None, -3], type=pa.duration(unit))}) + array = parquet_to_arrays(self.to_file(table, row_group_size=1))[0] + self.assertEqual(array.dtype, np.dtype(f'timedelta64[{unit}]')) + np.testing.assert_array_equal( + array, np.array([5, 'NaT', -3], dtype=f'timedelta64[{unit}]') + ) + self.assertFalse(array.flags.writeable) + + def test_parquet_to_arrays_times(self): + types = (pa.time32('s'), pa.time32('ms'), pa.time64('us'), pa.time64('ns')) + for dtype in types: + with self.subTest(dtype=dtype): + values = [datetime.time(1, 2, 3), None, datetime.time(23, 59, 59)] + table = pa.table({'a': pa.array(values, type=dtype)}) + array = parquet_to_arrays(self.to_file(table))[0] + self.assertEqual(array.dtype, np.dtype(object)) + self.assertEqual(array.tolist(), values) + table = pa.table({'a': pa.array([1001], type=pa.time64('ns'))}) + with self.assertRaisesRegex(ValueError, 'non-zero nanoseconds'): + parquet_to_arrays(self.to_file(table)) + + def test_parquet_to_arrays_nested_temporal(self): + struct = pa.struct([ + ('ns', pa.timestamp('ns')), + ('us', pa.timestamp('us')), + ('d', pa.duration('ms')), + ]) + table = pa.table({'s': pa.array( + [{'ns': 1, 'us': 2, 'd': 3}, {'ns': None, 'us': None, 'd': None}], + type=struct, + )}) + array = parquet_to_arrays(self.to_file(table))[0] + self.assertEqual(array[0], { + 'ns': np.datetime64(1, 'ns'), + 'us': np.datetime64(2, 'us'), + 'd': np.timedelta64(3, 'ms'), + }) + self.assertIsInstance(array[0]['ns'], np.datetime64) + self.assertEqual(array[0]['ns'].dtype, np.dtype('datetime64[ns]')) + self.assertEqual(array[1], {'ns': None, 'us': None, 'd': None}) + + def test_parquet_to_arrays_read_bytes_like(self): + for wrap in (bytearray, memoryview): + class BytesLike(io.BytesIO): + def read(self, size=-1, wrap=wrap): + return wrap(super().read(size)) + + with self.subTest(wrap=wrap): + data = self.to_file(pa.table({'a': [1, 2]})).getvalue() + self.assertEqual(parquet_to_arrays(BytesLike(data))[0].tolist(), [1, 2]) + + class NotBytes(io.BytesIO): + def read(self, size=-1): + return 'text' + + with self.assertRaisesRegex(TypeError, 'bytes-like'): + parquet_to_arrays(NotBytes(data)) + + 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_backend(self): + with mock.patch.dict('sys.modules', {'arraykit._parquet': None}): + with self.assertRaisesRegex(ImportError, 'ARRAYKIT_WITH_PARQUET'): + parquet_to_arrays(io.BytesIO()) + + +if __name__ == '__main__': + unittest.main()