This is an automated email from the ASF dual-hosted git repository.
paleolimbot pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-nanoarrow.git
The following commit(s) were added to refs/heads/main by this push:
new c66ddc35 feat(python): Add bindings for IPC reader (#388)
c66ddc35 is described below
commit c66ddc35a9ccf0374aadc2d3a8821431ed0c9ca6
Author: Dewey Dunnington <[email protected]>
AuthorDate: Wed Feb 21 13:14:40 2024 -0400
feat(python): Add bindings for IPC reader (#388)
This PR adds bindings to nanoarrow's (high-ish level) IPC reader. There
are some lower-level concepts that might be nice to expose at some point
but the `ArrowArrayStream` implementation lets users realize most of the
benefit.
I am envisioning that there will be a higher level `ArrayStream` class
than the `CArrayStream`, so I am not sure that the current interface
(create a `Stream`, then use `na.c_array_stream()`) will be the one
users actually use to access this. Probably something more like
`na.ArrayStream.from_stream([path or url or file])` would be
appropriate. I held off on writing the Examples section until the first
round of review in case there's a better way to go about this!
nanoarrow's reader is not as fast as pyarrow's using its internal
filesystem (which might do memory mapping and avoids more copies), but
is exactly as fast as pyarrow's reader when using a file-like object.
See also #390, which implements basically the same interface in R.
```python
import numpy as np
import pyarrow as pa
from pyarrow import ipc
import nanoarrow as na
from nanoarrow.ipc import Stream
n = int(1e6)
n_cols = 10
tab = pa.table(
[np.random.random(n) for _ in range(n_cols)],
names=[f"col{i}" for i in range(n_cols)]
)
with ipc.new_stream("file.arrows", tab.schema) as stream:
stream.write_table(tab)
def read_pyarrow():
with ipc.open_stream("file.arrows") as stream:
return list(stream)
%timeit len(read_pyarrow())
#> 1.19 ms ± 48.9 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
def read_pyarrow_pyobj():
with open("file.arrows", "rb") as f, ipc.open_stream(f) as stream:
return list(stream)
%timeit len(read_pyarrow_pyobj())
#> 11.5 ms ± 173 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
def read_nanoarrow():
with Stream.from_path("file.arrows") as input:
stream = na.c_array_stream(input)
return list(stream)
%timeit len(read_nanoarrow())
#> 11.4 ms ± 56.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
```
---------
Co-authored-by: Joris Van den Bossche <[email protected]>
Co-authored-by: Dane Pitkin <[email protected]>
---
.github/workflows/python-wheels.yaml | 5 +-
python/.gitignore | 9 +-
python/MANIFEST.in | 6 +-
python/bootstrap.py | 91 +++----
python/setup.py | 22 +-
python/src/nanoarrow/_ipc_lib.pyx | 152 ++++++++++++
python/src/nanoarrow/_lib.pyx | 56 +++--
.../nanoarrow/{_lib_utils.py => _repr_utils.py} | 43 +++-
python/src/nanoarrow/c_lib.py | 4 +-
python/src/nanoarrow/ipc.py | 265 +++++++++++++++++++++
python/tests/test_c_buffer.py | 7 +-
python/tests/test_ipc.py | 104 ++++++++
python/tests/test_nanoarrow.py | 6 +-
13 files changed, 667 insertions(+), 103 deletions(-)
diff --git a/.github/workflows/python-wheels.yaml
b/.github/workflows/python-wheels.yaml
index 4168b96f..207356cc 100644
--- a/.github/workflows/python-wheels.yaml
+++ b/.github/workflows/python-wheels.yaml
@@ -99,11 +99,8 @@ jobs:
python -m cibuildwheel --output-dir wheelhouse python
env:
CIBW_ARCHS_MACOS: x86_64 arm64
- # Optional (test suite will pass if these are not available)
- # Commenting this for now because not all the tests pass yet (fixes
in another PR)
- # CIBW_BEFORE_TEST: pip install --only-binary ":all:" pyarrow numpy
|| pip install --only-binary ":all:" numpy || true
CIBW_TEST_REQUIRES: pytest
- CIBW_TEST_COMMAND: pytest {package}/tests
+ CIBW_TEST_COMMAND: pytest {package}/tests -vv
- uses: actions/upload-artifact@v4
with:
diff --git a/python/.gitignore b/python/.gitignore
index 09279802..69037356 100644
--- a/python/.gitignore
+++ b/python/.gitignore
@@ -16,12 +16,9 @@
# specific language governing permissions and limitations
# under the License.
-src/nanoarrow/nanoarrow.c
-src/nanoarrow/nanoarrow.h
-src/nanoarrow/nanoarrow_device.h
-src/nanoarrow/nanoarrow_testing.hpp
-src/nanoarrow/nanoarrow_c.pxd
-src/nanoarrow/*.c
+src/nanoarrow/vendor
+src/nanoarrow/_lib.c
+src/nanoarrow/_ipc_lib.c
# Byte-compiled / optimized / DLL files
__pycache__/
diff --git a/python/MANIFEST.in b/python/MANIFEST.in
index 21fdc935..c8309a54 100644
--- a/python/MANIFEST.in
+++ b/python/MANIFEST.in
@@ -16,9 +16,5 @@
# under the License.
exclude bootstrap.py
-include src/nanoarrow/nanoarrow.c
-include src/nanoarrow/nanoarrow.h
-include src/nanoarrow/nanoarrow_c.pxd
-include src/nanoarrow/nanoarrow_device.c
-include src/nanoarrow/nanoarrow_device.h
+recursive-include src/nanoarrow/vendor *.c *.h *.pxd
include src/nanoarrow/nanoarrow_device_c.pxd
diff --git a/python/bootstrap.py b/python/bootstrap.py
index f1c9fbee..6395d128 100644
--- a/python/bootstrap.py
+++ b/python/bootstrap.py
@@ -181,49 +181,59 @@ class NanoarrowPxdGenerator:
def copy_or_generate_nanoarrow_c():
this_dir = os.path.abspath(os.path.dirname(__file__))
source_dir = os.path.dirname(this_dir)
-
- maybe_nanoarrow_h = os.path.join(this_dir, "src/nanoarrow/nanoarrow.h")
- maybe_nanoarrow_c = os.path.join(this_dir, "src/nanoarrow/nanoarrow.c")
- maybe_nanoarrow_device_h = os.path.join(
- this_dir, "src/nanoarrow/nanoarrow_device.h"
- )
- maybe_nanoarrow_device_c = os.path.join(
- this_dir, "src/nanoarrow/nanoarrow_device.c"
- )
-
- for f in (
- maybe_nanoarrow_c,
- maybe_nanoarrow_h,
- maybe_nanoarrow_device_h,
- maybe_nanoarrow_device_c,
- ):
+ vendor_dir = os.path.join(this_dir, "src", "nanoarrow", "vendor")
+
+ vendored_files = [
+ "nanoarrow.h",
+ "nanoarrow.c",
+ "nanoarrow_ipc.h",
+ "nanoarrow_ipc.c",
+ "nanoarrow_device.h",
+ "nanoarrow_device.c",
+ ]
+ dst = {name: os.path.join(vendor_dir, name) for name in vendored_files}
+
+ for f in dst.values():
if os.path.exists(f):
os.unlink(f)
is_cmake_dir = "CMakeLists.txt" in os.listdir(source_dir)
- is_in_nanoarrow_repo = "nanoarrow.h" in os.listdir(
+ is_in_nanoarrow_repo = is_cmake_dir and "nanoarrow.h" in os.listdir(
os.path.join(source_dir, "src", "nanoarrow")
)
+
+ if not is_in_nanoarrow_repo:
+ raise ValueError(
+ "Attempt to build source distribution outside the nanoarrow repo"
+ )
+
cmake_bin = os.getenv("CMAKE_BIN")
if not cmake_bin:
cmake_bin = "cmake"
has_cmake = os.system(f"{cmake_bin} --version") == 0
+ if not has_cmake:
+ raise ValueError("Attempt to build source distribution without CMake")
- with tempfile.TemporaryDirectory() as build_dir:
- if is_in_nanoarrow_repo:
- device_ext_src = os.path.join(
- source_dir, "extensions/nanoarrow_device/src/nanoarrow"
- )
- shutil.copyfile(
- os.path.join(device_ext_src, "nanoarrow_device.h"),
- maybe_nanoarrow_device_h,
- )
- shutil.copyfile(
- os.path.join(device_ext_src, "nanoarrow_device.c"),
- maybe_nanoarrow_device_c,
- )
+ # The C library, IPC extension, and Device extension all currently have
slightly
+ # different methods of bundling (hopefully this can be unified)
+
+ if not os.path.exists(vendor_dir):
+ os.mkdir(vendor_dir)
+
+ # Copy device files
+ device_ext_src = os.path.join(
+ source_dir, "extensions/nanoarrow_device/src/nanoarrow"
+ )
+ for device_file in ["nanoarrow_device.h", "nanoarrow_device.c"]:
+ shutil.copyfile(
+ os.path.join(device_ext_src, device_file),
+ dst[device_file],
+ )
- if has_cmake and is_cmake_dir and is_in_nanoarrow_repo:
+ ipc_source_dir = os.path.join(source_dir, "extensions/nanoarrow_ipc")
+
+ for cmake_project in [source_dir, ipc_source_dir]:
+ with tempfile.TemporaryDirectory() as build_dir:
try:
subprocess.run(
[
@@ -231,7 +241,8 @@ def copy_or_generate_nanoarrow_c():
"-B",
build_dir,
"-S",
- source_dir,
+ cmake_project,
+ "-DNANOARROW_IPC_BUNDLE=ON",
"-DNANOARROW_BUNDLE=ON",
"-DNANOARROW_NAMESPACE=PythonPkg",
]
@@ -242,29 +253,21 @@ def copy_or_generate_nanoarrow_c():
"--install",
build_dir,
"--prefix",
- os.path.join(this_dir, "src", "nanoarrow"),
+ vendor_dir,
]
)
except Exception as e:
warnings.warn(f"cmake call failed: {e}")
- else:
- raise ValueError(
- "Attempt to build source distribution outside the nanoarrow
repo"
- )
- if not os.path.exists(os.path.join(this_dir, "src/nanoarrow/nanoarrow.h")):
+ if not os.path.exists(dst["nanoarrow.h"]):
raise ValueError("Attempt to vendor nanoarrow.c/h failed")
- maybe_nanoarrow_hpp = os.path.join(this_dir, "src/nanoarrow/nanoarrow.hpp")
- if os.path.exists(maybe_nanoarrow_hpp):
- os.unlink(maybe_nanoarrow_hpp)
-
# Runs the pxd generator with some information about the file name
def generate_nanoarrow_pxd():
this_dir = os.path.abspath(os.path.dirname(__file__))
- maybe_nanoarrow_h = os.path.join(this_dir, "src/nanoarrow/nanoarrow.h")
- maybe_nanoarrow_pxd = os.path.join(this_dir,
"src/nanoarrow/nanoarrow_c.pxd")
+ maybe_nanoarrow_h = os.path.join(this_dir,
"src/nanoarrow/vendor/nanoarrow.h")
+ maybe_nanoarrow_pxd = os.path.join(this_dir,
"src/nanoarrow/vendor/nanoarrow_c.pxd")
NanoarrowPxdGenerator().generate_nanoarrow_pxd(
maybe_nanoarrow_h, maybe_nanoarrow_pxd
diff --git a/python/setup.py b/python/setup.py
index cdffda2b..5153fec7 100644
--- a/python/setup.py
+++ b/python/setup.py
@@ -69,17 +69,31 @@ setup(
ext_modules=[
Extension(
name="nanoarrow._lib",
- include_dirs=["src/nanoarrow"],
+ include_dirs=["src/nanoarrow", "src/nanoarrow/vendor"],
language="c",
sources=[
"src/nanoarrow/_lib.pyx",
- "src/nanoarrow/nanoarrow.c",
- "src/nanoarrow/nanoarrow_device.c",
+ "src/nanoarrow/vendor/nanoarrow.c",
+ "src/nanoarrow/vendor/nanoarrow_device.c",
],
extra_compile_args=extra_compile_args,
extra_link_args=extra_link_args,
define_macros=extra_define_macros,
- )
+ ),
+ Extension(
+ name="nanoarrow._ipc_lib",
+ include_dirs=["src/nanoarrow", "src/nanoarrow/vendor"],
+ language="c",
+ sources=[
+ "src/nanoarrow/_ipc_lib.pyx",
+ "src/nanoarrow/vendor/nanoarrow.c",
+ "src/nanoarrow/vendor/nanoarrow_ipc.c",
+ "src/nanoarrow/vendor/flatcc.c",
+ ],
+ extra_compile_args=extra_compile_args,
+ extra_link_args=extra_link_args,
+ define_macros=extra_define_macros,
+ ),
],
version=version,
)
diff --git a/python/src/nanoarrow/_ipc_lib.pyx
b/python/src/nanoarrow/_ipc_lib.pyx
new file mode 100644
index 00000000..ea15921b
--- /dev/null
+++ b/python/src/nanoarrow/_ipc_lib.pyx
@@ -0,0 +1,152 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# cython: language_level = 3
+# cython: linetrace=True
+
+from libc.stdint cimport uint8_t, int64_t, uintptr_t
+from libc.errno cimport EIO
+from libc.stdio cimport snprintf
+from cpython.ref cimport PyObject, Py_INCREF, Py_DECREF
+from cpython cimport Py_buffer, PyBuffer_FillInfo
+
+from nanoarrow_c cimport (
+ ArrowErrorCode,
+ ArrowError,
+ NANOARROW_OK,
+ ArrowArrayStream,
+)
+
+
+cdef extern from "nanoarrow_ipc.h" nogil:
+ struct ArrowIpcInputStream:
+ ArrowErrorCode (*read)(ArrowIpcInputStream* stream, uint8_t* buf,
+ int64_t buf_size_bytes, int64_t* size_read_out,
+ ArrowError* error)
+ void (*release)(ArrowIpcInputStream* stream)
+ void* private_data
+
+ struct ArrowIpcArrayStreamReaderOptions:
+ int64_t field_index
+ int use_shared_buffers
+
+ ArrowErrorCode ArrowIpcArrayStreamReaderInit(
+ ArrowArrayStream* out, ArrowIpcInputStream* input_stream,
+ ArrowIpcArrayStreamReaderOptions* options)
+
+
+cdef class PyInputStreamPrivate:
+ cdef object obj
+ cdef object obj_method
+ cdef void* addr
+ cdef Py_ssize_t size_bytes
+ cdef int close_stream
+
+ def __cinit__(self, obj, close_stream=False):
+ self.obj = obj
+ self.obj_method = obj.readinto
+ self.addr = NULL
+ self.size_bytes = 0
+ self.close_stream = close_stream
+
+ # Implement the buffer protocol so that this object can be used as
+ # the argument to xxx.readinto(). This ensures that no extra copies
+ # (beyond any buffering done by the upstream file-like object) are held
+ # since the upstream object has access to the preallocated output buffer.
+ # In this case, the preallocation is done by the ArrowArrayStream
+ # implementation before issuing each read call (two per message, with
+ # an extra call for a RecordBatch message to get the actual buffer data).
+ def __getbuffer__(self, Py_buffer* buffer, int flags):
+ PyBuffer_FillInfo(buffer, self, self.addr, self.size_bytes, 0, flags)
+
+ def __releasebuffer__(self, Py_buffer* buffer):
+ pass
+
+
+cdef ArrowErrorCode py_input_stream_read(ArrowIpcInputStream* stream, uint8_t*
buf,
+ int64_t buf_size_bytes, int64_t*
size_read_out,
+ ArrowError* error) noexcept:
+ cdef PyInputStreamPrivate stream_private = <object>stream.private_data
+ stream_private.addr = buf
+ stream_private.size_bytes = buf_size_bytes
+
+ try:
+ size_read_out[0] = stream_private.obj_method(stream_private)
+ return NANOARROW_OK
+ except Exception as e:
+ cls = type(e).__name__.encode()
+ msg = str(e).encode()
+ snprintf(
+ error.message,
+ sizeof(error.message),
+ "%s: %s",
+ <const char*>cls,
+ <const char*>msg
+ )
+ return EIO
+
+
+cdef void py_input_stream_release(ArrowIpcInputStream* stream) noexcept:
+ cdef PyInputStreamPrivate stream_private = <object>stream.private_data
+ if stream_private.close_stream:
+ stream_private.obj.close()
+
+ Py_DECREF(stream_private)
+ stream.private_data = NULL
+ stream.release = NULL
+
+
+cdef class CIpcInputStream:
+ cdef ArrowIpcInputStream _stream
+
+ def __cinit__(self):
+ self._stream.release = NULL
+
+ def is_valid(self):
+ return self._stream.release != NULL
+
+ def __dealloc__(self):
+ # Duplicating release() to avoid Python API calls in the deallocator
+ if self._stream.release != NULL:
+ self._stream.release(&self._stream)
+
+ def release(self):
+ if self._stream.release != NULL:
+ self._stream.release(&self._stream)
+ return True
+ else:
+ return False
+
+ @staticmethod
+ def from_readable(obj, close_stream=False):
+ cdef CIpcInputStream stream = CIpcInputStream()
+ cdef PyInputStreamPrivate private_data = PyInputStreamPrivate(obj,
close_stream)
+
+ stream._stream.private_data = <PyObject*>private_data
+ Py_INCREF(private_data)
+ stream._stream.read = &py_input_stream_read
+ stream._stream.release = &py_input_stream_release
+ return stream
+
+
+def init_array_stream(CIpcInputStream input_stream, uintptr_t out):
+ cdef ArrowArrayStream* out_ptr = <ArrowArrayStream*>out
+
+ # There are some options here that could be exposed at some point
+ cdef int code = ArrowIpcArrayStreamReaderInit(out_ptr,
&input_stream._stream, NULL)
+ if code != NANOARROW_OK:
+ raise RuntimeError(f"ArrowIpcArrayStreamReaderInit() failed with code
[{code}]")
diff --git a/python/src/nanoarrow/_lib.pyx b/python/src/nanoarrow/_lib.pyx
index 00b4cf45..e1307d0b 100644
--- a/python/src/nanoarrow/_lib.pyx
+++ b/python/src/nanoarrow/_lib.pyx
@@ -52,7 +52,7 @@ from nanoarrow_device_c cimport *
from sys import byteorder as sys_byteorder
from struct import unpack_from, iter_unpack, calcsize, Struct
-from nanoarrow import _lib_utils
+from nanoarrow import _repr_utils
def c_version():
"""Return the nanoarrow C library version string
@@ -525,7 +525,7 @@ cdef class CDevice:
return CDeviceArray(holder, <uintptr_t>device_array_ptr, schema)
def __repr__(self):
- return _lib_utils.device_repr(self)
+ return _repr_utils.device_repr(self)
@property
def device_type(self):
@@ -658,7 +658,7 @@ cdef class CSchema:
return out_str
def __repr__(self):
- return _lib_utils.schema_repr(self)
+ return _repr_utils.schema_repr(self)
@property
def format(self):
@@ -858,7 +858,7 @@ cdef class CSchemaView:
def __repr__(self):
- return _lib_utils.schema_view_repr(self)
+ return _repr_utils.schema_view_repr(self)
cdef class CSchemaBuilder:
@@ -1121,7 +1121,7 @@ cdef class CArray:
return None
def __repr__(self):
- return _lib_utils.array_repr(self)
+ return _repr_utils.array_repr(self)
cdef class CArrayView:
@@ -1244,7 +1244,7 @@ cdef class CArrayView:
)
def __repr__(self):
- return _lib_utils.array_view_repr(self)
+ return _repr_utils.array_view_repr(self)
@staticmethod
def from_cpu_array(CArray array):
@@ -1466,7 +1466,8 @@ cdef class CBufferView:
pass
def __repr__(self):
- return f"CBufferView({_lib_utils.buffer_view_repr(self)})"
+ class_label = _repr_utils.make_class_label(self,
module="nanoarrow.c_lib")
+ return f"{class_label}({_repr_utils.buffer_view_repr(self)})"
cdef class CBuffer:
@@ -1628,10 +1629,11 @@ cdef class CBuffer:
self._get_buffer_count -= 1
def __repr__(self):
+ class_label = _repr_utils.make_class_label(self,
module="nanoarrow.c_lib")
if self._ptr == NULL:
- return "CBuffer(<invalid>)"
+ return f"{class_label}(<invalid>)"
- return f"CBuffer({_lib_utils.buffer_view_repr(self._view)})"
+ return f"{class_label}({_repr_utils.buffer_view_repr(self._view)})"
cdef class CBufferBuilder:
@@ -1736,6 +1738,10 @@ cdef class CBufferBuilder:
self._buffer = CBuffer.empty()
return out
+ def __repr__(self):
+ class_label = _repr_utils.make_class_label(self,
module="nanoarrow.c_lib")
+ return f"{class_label}({self.size_bytes}/{self.capacity_bytes})"
+
cdef class CArrayBuilder:
cdef CArray c_array
@@ -1926,6 +1932,10 @@ cdef class CArrayStream:
self._ptr = <ArrowArrayStream*>addr
self._cached_schema = None
+ def release(self):
+ if self.is_valid():
+ self._ptr.release(self._ptr)
+
@staticmethod
def _import_from_c_capsule(stream_capsule):
"""
@@ -1983,10 +1993,15 @@ cdef class CArrayStream:
def _get_schema(self, CSchema schema):
self._assert_valid()
cdef Error error = Error()
- cdef int code = self._ptr.get_schema(self._ptr, schema._ptr)
- Error.raise_error_not_ok("ArrowArrayStream::get_schema()", code)
+ cdef int code = ArrowArrayStreamGetSchema(self._ptr, schema._ptr,
&error.c_error)
+ error.raise_message_not_ok("ArrowArrayStream::get_schema()", code)
+
+ def _get_cached_schema(self):
+ if self._cached_schema is None:
+ self._cached_schema = CSchema.allocate()
+ self._get_schema(self._cached_schema)
- self._cached_schema = schema
+ return self._cached_schema
def get_schema(self):
"""Get the schema associated with this stream
@@ -2006,14 +2021,11 @@ cdef class CArrayStream:
# Array that is returned. This is independent of get_schema(),
# which is guaranteed to call the C object's callback and
# faithfully pass on the returned value.
- if self._cached_schema is None:
- self._cached_schema = CSchema.allocate()
- self._get_schema(self._cached_schema)
cdef Error error = Error()
- cdef CArray array = CArray.allocate(self._cached_schema)
+ cdef CArray array = CArray.allocate(self._get_cached_schema())
cdef int code = ArrowArrayStreamGetNext(self._ptr, array._ptr,
&error.c_error)
- Error.raise_error_not_ok("ArrowArrayStream::get_next()", code)
+ error.raise_message_not_ok("ArrowArrayStream::get_next()", code)
if not array.is_valid():
raise StopIteration()
@@ -2026,8 +2038,14 @@ cdef class CArrayStream:
def __next__(self):
return self.get_next()
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *args, **kwargs):
+ self.release()
+
def __repr__(self):
- return _lib_utils.array_stream_repr(self)
+ return _repr_utils.array_stream_repr(self)
cdef class CDeviceArray:
@@ -2053,4 +2071,4 @@ cdef class CDeviceArray:
return CArray(self, <uintptr_t>&self._ptr.array, self._schema)
def __repr__(self):
- return _lib_utils.device_array_repr(self)
+ return _repr_utils.device_array_repr(self)
diff --git a/python/src/nanoarrow/_lib_utils.py
b/python/src/nanoarrow/_repr_utils.py
similarity index 85%
rename from python/src/nanoarrow/_lib_utils.py
rename to python/src/nanoarrow/_repr_utils.py
index cdb1933e..26a274aa 100644
--- a/python/src/nanoarrow/_lib_utils.py
+++ b/python/src/nanoarrow/_repr_utils.py
@@ -20,14 +20,21 @@
# after editing when working with an editable installation)
+def make_class_label(obj, module=None):
+ if module is None:
+ module = obj.__class__.__module__
+ return f"{module}.{obj.__class__.__name__}"
+
+
def schema_repr(schema, indent=0):
indent_str = " " * indent
+ class_label = make_class_label(schema, module="nanoarrow.c_lib")
if schema._addr() == 0:
- return "<NULL nanoarrow.c_lib.CSchema>"
+ return f"<{class_label} <NULL>>"
elif not schema.is_valid():
- return "<released nanoarrow.c_lib.CSchema>"
+ return f"<{class_label} <released>>"
- lines = [f"<nanoarrow.c_lib.CSchema {schema._to_string()}>"]
+ lines = [f"<{class_label} {schema._to_string()}>"]
for attr in ("format", "name", "flags"):
attr_repr = repr(getattr(schema, attr))
@@ -60,15 +67,16 @@ def array_repr(array, indent=0, max_char_width=80):
max_char_width = 20
indent_str = " " * indent
+ class_label = make_class_label(array, module="nanoarrow.c_lib")
if array._addr() == 0:
- return "<NULL nanoarrow.c_lib.CArray>"
+ return f"<{class_label} <NULL>>"
elif not array.is_valid():
- return "<released nanoarrow.c_lib.CArray>"
+ return f"<{class_label} <released>>"
schema_string = array.schema._to_string(
max_chars=max_char_width - indent - 23, recursive=True
)
- lines = [f"<nanoarrow.c_lib.CArray {schema_string}>"]
+ lines = [f"<{class_label} {schema_string}>"]
for attr in ("length", "offset", "null_count", "buffers"):
attr_repr = repr(getattr(array, attr))
lines.append(f"{indent_str}- {attr}: {attr_repr}")
@@ -88,8 +96,10 @@ def array_repr(array, indent=0, max_char_width=80):
def schema_view_repr(schema_view):
+ class_label = make_class_label(schema_view, module="nanoarrow.c_lib")
+
lines = [
- "<nanoarrow.c_lib.CSchemaView>",
+ f"<{class_label}>",
f"- type: {repr(schema_view.type)}",
f"- storage_type: {repr(schema_view.storage_type)}",
]
@@ -109,8 +119,9 @@ def schema_view_repr(schema_view):
def array_view_repr(array_view, max_char_width=80, indent=0):
indent_str = " " * indent
+ class_label = make_class_label(array_view, module="nanoarrow.c_lib")
- lines = ["<nanoarrow.c_lib.CArrayView>"]
+ lines = [f"<{class_label}>"]
for attr in ("storage_type", "length", "offset", "null_count"):
attr_repr = repr(getattr(array_view, attr))
@@ -190,12 +201,14 @@ def buffer_view_preview_cpu(buffer_view, max_char_width):
def array_stream_repr(array_stream, max_char_width=80):
+ class_label = make_class_label(array_stream, module="nanoarrow.c_lib")
+
if array_stream._addr() == 0:
- return "<NULL nanoarrow.c_lib.CArrayStream>"
+ return f"<{class_label} <NULL>>"
elif not array_stream.is_valid():
- return "<released nanoarrow.c_lib.CArrayStream>"
+ return f"<{class_label} <released>>"
- lines = ["<nanoarrow.c_lib.CArrayStream>"]
+ lines = [f"<{class_label}>"]
try:
schema = array_stream.get_schema()
schema_string = schema._to_string(max_chars=max_char_width - 16,
recursive=True)
@@ -207,7 +220,9 @@ def array_stream_repr(array_stream, max_char_width=80):
def device_array_repr(device_array):
- title_line = "<nanoarrow.device.c_lib.CDeviceArray>"
+ class_label = make_class_label(device_array, module="nanoarrow.device")
+
+ title_line = f"<{class_label}>"
device_type = f"- device_type: {device_array.device_type}"
device_id = f"- device_id: {device_array.device_id}"
array = f"- array: {array_repr(device_array.array, indent=2)}"
@@ -215,7 +230,9 @@ def device_array_repr(device_array):
def device_repr(device):
- title_line = "<nanoarrow.device.Device>"
+ class_label = make_class_label(device, module="nanoarrow.device")
+
+ title_line = f"<{class_label}>"
device_type = f"- device_type: {device.device_type}"
device_id = f"- device_id: {device.device_id}"
return "\n".join([title_line, device_type, device_id])
diff --git a/python/src/nanoarrow/c_lib.py b/python/src/nanoarrow/c_lib.py
index 84631034..20af4678 100644
--- a/python/src/nanoarrow/c_lib.py
+++ b/python/src/nanoarrow/c_lib.py
@@ -452,9 +452,9 @@ def c_buffer(obj, schema=None) -> CBuffer:
>>> import nanoarrow as na
>>> na.c_buffer(b"1234")
- CBuffer(uint8[4 b] 49 50 51 52)
+ nanoarrow.c_lib.CBuffer(uint8[4 b] 49 50 51 52)
>>> na.c_buffer([1, 2, 3], na.int32())
- CBuffer(int32[12 b] 1 2 3)
+ nanoarrow.c_lib.CBuffer(int32[12 b] 1 2 3)
"""
if isinstance(obj, CBuffer) and schema is None:
return obj
diff --git a/python/src/nanoarrow/ipc.py b/python/src/nanoarrow/ipc.py
new file mode 100644
index 00000000..82b719ba
--- /dev/null
+++ b/python/src/nanoarrow/ipc.py
@@ -0,0 +1,265 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import io
+
+from nanoarrow._ipc_lib import CIpcInputStream, init_array_stream
+from nanoarrow._lib import CArrayStream
+
+from nanoarrow import _repr_utils
+
+
+class Stream:
+ """Stream of serialized Arrow data
+
+ Reads file paths or otherwise readable file objects that contain
+ serialized Arrow data. Arrow documentation typically refers to this format
+ as "Arrow IPC" because its origin was as a means to transmit tables between
+ processes; however, this format can also be written to and read from files
+ or URLs and is essentially a high-performance equivalent of a CSV file that
+ does a better job maintaining type fidelity.
+
+ Use :staticmethod:`from_readable`, :staticmethod:`from_path`, or
+ :staticmethod:`from_url` to construct these streams.
+
+ Examples
+ --------
+
+ >>> import nanoarrow as na
+ >>> from nanoarrow.ipc import Stream
+ >>> with Stream.example() as inp, na.c_array_stream(inp) as stream:
+ ... stream
+ <nanoarrow.c_lib.CArrayStream>
+ - get_schema(): struct<some_col: int32>
+ """
+
+ def __init__(self):
+ self._stream = None
+ self._desc = None
+
+ def _is_valid(self) -> bool:
+ return self._stream is not None and self._stream.is_valid()
+
+ def __arrow_c_stream__(self, requested_schema=None):
+ """Export this stream as an ArrowArrayStream
+
+ Implements the Arrow PyCapsule interface by transferring ownership of
this
+ input stream to an ArrowArrayStream wrapped by a PyCapsule.
+ """
+ if not self._is_valid():
+ raise RuntimeError("nanoarrow.ipc.Stream is no longer valid")
+
+ with CArrayStream.allocate() as array_stream:
+ init_array_stream(self._stream, array_stream._addr())
+ array_stream._get_cached_schema()
+ return
array_stream.__arrow_c_stream__(requested_schema=requested_schema)
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *args, **kwargs):
+ if self._stream is not None:
+ self._stream.release()
+
+ @staticmethod
+ def from_readable(obj):
+ """Wrap an open readable object as an Arrow stream
+
+ Wraps a readable object (specificially, an object that implements a
+ ``readinto()`` method) as a non-owning Stream. Closing ``obj`` remains
+ the caller's responsibility: neither this stream nor the resulting
array
+ stream will call ``obj.close()``.
+
+ Parameters
+ ----------
+ obj : readable file-like
+ An object implementing ``readinto()``.
+
+ Examples
+ --------
+
+ >>> import io
+ >>> import nanoarrow as na
+ >>> from nanoarrow.ipc import Stream
+ >>> with io.BytesIO(Stream.example_bytes()) as f:
+ ... inp = Stream.from_readable(f)
+ ... na.c_array_stream(inp)
+ <nanoarrow.c_lib.CArrayStream>
+ - get_schema(): struct<some_col: int32>
+ """
+ out = Stream()
+ out._stream = CIpcInputStream.from_readable(obj)
+ out._desc = repr(obj)
+ return out
+
+ @staticmethod
+ def from_path(obj, *args, **kwargs):
+ """Wrap a local file as an Arrow stream
+
+ Wraps a pathlike object (specificially, one that can be passed to
``open()``)
+ as an owning Stream. The file will be opened in binary mode and will
be closed
+ when this stream or the resulting array stream is released.
+
+ Parameters
+ ----------
+ obj : path-like
+ A string or path-like object that can be passed to ``open()``
+
+ Examples
+ --------
+
+ >>> import tempfile
+ >>> import os
+ >>> import nanoarrow as na
+ >>> from nanoarrow.ipc import Stream
+ >>> with tempfile.TemporaryDirectory() as td:
+ ... path = os.path.join(td, "test.arrows")
+ ... with open(path, "wb") as f:
+ ... nbytes = f.write(Stream.example_bytes())
+ ...
+ ... with Stream.from_path(path) as inp, na.c_array_stream(inp) as
stream:
+ ... stream
+ <nanoarrow.c_lib.CArrayStream>
+ - get_schema(): struct<some_col: int32>
+ """
+ out = Stream()
+ out._stream = CIpcInputStream.from_readable(
+ open(obj, "rb", *args, **kwargs), close_stream=True
+ )
+ out._desc = repr(obj)
+ return out
+
+ @staticmethod
+ def from_url(obj, *args, **kwargs):
+ """Wrap a URL as an Arrow stream
+
+ Wraps a URL (specificially, one that can be passed to
+ ``urllib.request.urlopen()``) as an owning Stream. The URL will be
+ closed when this stream or the resulting array stream is released.
+
+ Parameters
+ ----------
+ obj : str
+ A URL that can be passed to ``urllib.request.urlopen()``
+
+ Examples
+ --------
+
+ >>> import pathlib
+ >>> import tempfile
+ >>> import os
+ >>> import nanoarrow as na
+ >>> from nanoarrow.ipc import Stream
+ >>> with tempfile.TemporaryDirectory() as td:
+ ... path = os.path.join(td, "test.arrows")
+ ... with open(path, "wb") as f:
+ ... nbytes = f.write(Stream.example_bytes())
+ ...
+ ... uri = pathlib.Path(path).as_uri()
+ ... with Stream.from_url(uri) as inp, na.c_array_stream(inp) as
stream:
+ ... stream
+ <nanoarrow.c_lib.CArrayStream>
+ - get_schema(): struct<some_col: int32>
+ """
+ import urllib.request
+
+ out = Stream()
+ out._stream = CIpcInputStream.from_readable(
+ urllib.request.urlopen(obj, *args, **kwargs), close_stream=True
+ )
+ out._desc = repr(obj)
+ return out
+
+ @staticmethod
+ def example():
+ """Example Stream
+
+ A self-contained example whose value is the serialized version of
+ ``DataFrame({"some_col": [1, 2, 3]})``. This may be used for testing
+ and documentation and is useful because nanoarrow does not implement
+ a writer to generate test data.
+
+ Examples
+ --------
+
+ >>> from nanoarrow.ipc import Stream
+ >>> Stream.example()
+ <nanoarrow.ipc.Stream <_io.BytesIO object at ...>>
+ """
+ return Stream.from_readable(io.BytesIO(Stream.example_bytes()))
+
+ @staticmethod
+ def example_bytes():
+ """Example stream bytes
+
+ The underlying bytes of the :staticmethod:`example` Stream. This is
useful
+ for writing files or creating other types of test input.
+
+ Examples
+ --------
+
+ >>> import os
+ >>> import tempfile
+ >>> from nanoarrow.ipc import Stream
+ >>> with tempfile.TemporaryDirectory() as td:
+ ... path = os.path.join(td, "test.arrows")
+ ... with open(path, "wb") as f:
+ ... f.write(Stream.example_bytes())
+ 440
+ """
+ return _EXAMPLE_IPC_SCHEMA + _EXAMPLE_IPC_BATCH
+
+ def __repr__(self) -> str:
+ class_label = _repr_utils.make_class_label(self)
+ if self._is_valid():
+ return f"<{class_label} {self._desc}>"
+ else:
+ return f"<{class_label} <invalid>>"
+
+
+# A self-contained example whose value is the serialized verison of
+# DataFrame({"some_col": [1, 2, 3]}). Used to make the tests and documentation
+# self-contained since we don't have an IPC writer.
+_EXAMPLE_IPC_SCHEMA = (
+
b"\xff\xff\xff\xff\x10\x01\x00\x00\x10\x00\x00\x00\x00\x00\x0a\x00\x0e\x00\x06"
+
b"\x00\x05\x00\x08\x00\x0a\x00\x00\x00\x00\x01\x04\x00\x10\x00\x00\x00\x00\x00"
+
b"\x0a\x00\x0c\x00\x00\x00\x04\x00\x08\x00\x0a\x00\x00\x00\x3c\x00\x00\x00\x04"
+
b"\x00\x00\x00\x01\x00\x00\x00\x04\x00\x00\x00\x84\xff\xff\xff\x18\x00\x00\x00"
+
b"\x04\x00\x00\x00\x0a\x00\x00\x00\x73\x6f\x6d\x65\x5f\x76\x61\x6c\x75\x65\x00"
+
b"\x00\x08\x00\x00\x00\x73\x6f\x6d\x65\x5f\x6b\x65\x79\x00\x00\x00\x00\x01\x00"
+
b"\x00\x00\x18\x00\x00\x00\x00\x00\x12\x00\x18\x00\x08\x00\x06\x00\x07\x00\x0c"
+
b"\x00\x00\x00\x10\x00\x14\x00\x12\x00\x00\x00\x00\x00\x01\x02\x14\x00\x00\x00"
+
b"\x70\x00\x00\x00\x08\x00\x00\x00\x18\x00\x00\x00\x00\x00\x00\x00\x08\x00\x00"
+
b"\x00\x73\x6f\x6d\x65\x5f\x63\x6f\x6c\x00\x00\x00\x00\x01\x00\x00\x00\x0c\x00"
+
b"\x00\x00\x08\x00\x0c\x00\x04\x00\x08\x00\x08\x00\x00\x00\x20\x00\x00\x00\x04"
+
b"\x00\x00\x00\x10\x00\x00\x00\x73\x6f\x6d\x65\x5f\x76\x61\x6c\x75\x65\x5f\x66"
+
b"\x69\x65\x6c\x64\x00\x00\x00\x00\x0e\x00\x00\x00\x73\x6f\x6d\x65\x5f\x6b\x65"
+
b"\x79\x5f\x66\x69\x65\x6c\x64\x00\x00\x08\x00\x0c\x00\x08\x00\x07\x00\x08\x00"
+ b"\x00\x00\x00\x00\x00\x01\x20\x00\x00\x00\x00\x00\x00\x00"
+)
+
+_EXAMPLE_IPC_BATCH = (
+
b"\xff\xff\xff\xff\x88\x00\x00\x00\x14\x00\x00\x00\x00\x00\x00\x00\x0c\x00\x16"
+
b"\x00\x06\x00\x05\x00\x08\x00\x0c\x00\x0c\x00\x00\x00\x00\x03\x04\x00\x18\x00"
+
b"\x00\x00\x10\x00\x00\x00\x00\x00\x00\x00\x00\x00\x0a\x00\x18\x00\x0c\x00\x04"
+
b"\x00\x08\x00\x0a\x00\x00\x00\x3c\x00\x00\x00\x10\x00\x00\x00\x03\x00\x00\x00"
+
b"\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
+
b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x0c\x00"
+
b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x03\x00\x00\x00\x00"
+
b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x02\x00\x00\x00"
+ b"\x03\x00\x00\x00\x00\x00\x00\x00"
+)
diff --git a/python/tests/test_c_buffer.py b/python/tests/test_c_buffer.py
index 9fea15ff..0487a382 100644
--- a/python/tests/test_c_buffer.py
+++ b/python/tests/test_c_buffer.py
@@ -34,7 +34,7 @@ def test_buffer_invalid():
with pytest.raises(RuntimeError, match="CBuffer is not valid"):
memoryview(invalid)
- assert repr(invalid) == "CBuffer(<invalid>)"
+ assert repr(invalid) == "nanoarrow.c_lib.CBuffer(<invalid>)"
def test_c_buffer_constructor():
@@ -66,7 +66,7 @@ def test_c_buffer_empty():
assert empty.size_bytes == 0
assert bytes(empty) == b""
- assert repr(empty) == "CBuffer(binary[0 b] b'')"
+ assert repr(empty) == "nanoarrow.c_lib.CBuffer(binary[0 b] b'')"
# Export it via the Python buffer protocol wrapped in a new CBuffer
empty_roundtrip = na.c_buffer(empty)
@@ -83,7 +83,7 @@ def test_c_buffer_pybuffer():
assert buffer.size_bytes == len(data)
assert bytes(buffer) == b"abcdefghijklmnopqrstuvwxyz"
- assert repr(buffer).startswith("CBuffer(uint8[26 b] 97 98")
+ assert repr(buffer).startswith("nanoarrow.c_lib.CBuffer(uint8[26 b] 97 98")
def test_c_buffer_unsupported_type():
@@ -199,6 +199,7 @@ def test_c_buffer_builder():
builder = CBufferBuilder()
assert builder.size_bytes == 0
assert builder.capacity_bytes == 0
+ assert repr(builder) == "nanoarrow.c_lib.CBufferBuilder(0/0)"
builder.reserve_bytes(123)
assert builder.size_bytes == 0
diff --git a/python/tests/test_ipc.py b/python/tests/test_ipc.py
new file mode 100644
index 00000000..d23abdfb
--- /dev/null
+++ b/python/tests/test_ipc.py
@@ -0,0 +1,104 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import io
+import os
+import pathlib
+import tempfile
+
+import pytest
+from nanoarrow._lib import NanoarrowException
+from nanoarrow.ipc import Stream
+
+import nanoarrow as na
+
+
+def test_ipc_stream_example():
+
+ with Stream.example() as input:
+ assert input._is_valid() is True
+ assert "BytesIO object" in repr(input)
+
+ stream = na.c_array_stream(input)
+ assert input._is_valid() is False
+ assert stream.is_valid() is True
+ assert repr(input) == "<nanoarrow.ipc.Stream <invalid>>"
+ with pytest.raises(RuntimeError, match="no longer valid"):
+ stream = na.c_array_stream(input)
+
+ with stream:
+ schema = stream.get_schema()
+ assert schema.format == "+s"
+ assert schema.child(0).format == "i"
+ batches = list(stream)
+ assert stream.is_valid() is True
+
+ assert stream.is_valid() is False
+ assert len(batches) == 1
+ batch = na.c_array_view(batches[0])
+ assert list(batch.child(0).buffer(1)) == [1, 2, 3]
+
+
+def test_ipc_stream_from_path():
+ with tempfile.TemporaryDirectory() as td:
+ path = os.path.join(td, "test.arrows")
+ with open(path, "wb") as f:
+ f.write(Stream.example_bytes())
+
+ with Stream.from_path(path) as input:
+ assert repr(path) in repr(input)
+ with na.c_array_stream(input) as stream:
+ batches = list(stream)
+ assert len(batches) == 1
+ assert batches[0].length == 3
+
+
+def test_ipc_stream_from_url():
+ with tempfile.TemporaryDirectory() as td:
+ path = os.path.join(td, "test.arrows")
+ with open(path, "wb") as f:
+ f.write(Stream.example_bytes())
+
+ uri = pathlib.Path(path).as_uri()
+ with Stream.from_url(uri) as input:
+ with na.c_array_stream(input) as stream:
+ batches = list(stream)
+ assert len(batches) == 1
+ assert batches[0].length == 3
+
+
+def test_ipc_stream_python_exception_on_read():
+ class ExtraordinarilyInconvenientFile:
+ def readinto(self, obj):
+ raise RuntimeError("I error for all read requests")
+
+ input = Stream.from_readable(ExtraordinarilyInconvenientFile())
+ with pytest.raises(
+ NanoarrowException, match="RuntimeError: I error for all read requests"
+ ):
+ na.c_array_stream(input)
+
+
+def test_ipc_stream_error_on_read():
+ with io.BytesIO(Stream.example_bytes()[:100]) as f:
+ with Stream.from_readable(f) as input:
+
+ with pytest.raises(
+ NanoarrowException,
+ match="Expected >= 280 bytes of remaining data",
+ ):
+ na.c_array_stream(input)
diff --git a/python/tests/test_nanoarrow.py b/python/tests/test_nanoarrow.py
index 6b6f1254..ea3d95ec 100644
--- a/python/tests/test_nanoarrow.py
+++ b/python/tests/test_nanoarrow.py
@@ -72,7 +72,7 @@ def test_c_schema_basic():
schema = na.allocate_c_schema()
assert schema.is_valid() is False
assert schema._to_string() == "[invalid: schema is released]"
- assert repr(schema) == "<released nanoarrow.c_lib.CSchema>"
+ assert repr(schema) == "<nanoarrow.c_lib.CSchema <released>>"
schema = na.c_schema(pa.schema([pa.field("some_name", pa.int32())]))
@@ -175,7 +175,7 @@ def test_c_schema_view_extra_params():
def test_c_array_empty():
array = na.allocate_c_array()
assert array.is_valid() is False
- assert repr(array) == "<released nanoarrow.c_lib.CArray>"
+ assert repr(array) == "<nanoarrow.c_lib.CArray <released>>"
def test_c_array():
@@ -467,7 +467,7 @@ def test_buffers_interval_month_day_nano():
def test_c_array_stream():
array_stream = na.allocate_c_array_stream()
assert na.c_array_stream(array_stream) is array_stream
- assert repr(array_stream) == "<released nanoarrow.c_lib.CArrayStream>"
+ assert repr(array_stream) == "<nanoarrow.c_lib.CArrayStream <released>>"
assert array_stream.is_valid() is False
with pytest.raises(RuntimeError):
