Skip to content

Commit c74bb37

Browse files
authored
refactor(python): Split buffer Cython internals into a separate module (#549)
Continuing the saga, this splits the buffer internals into their own module. The definition of `Extension`s in setup.py was also getting unwieldy, so I factored that out into a function. This should better support future linking to a common nanoarrow static or shared library (instead of just compiling in the sources for whatever we need in that module).
1 parent a24aa23 commit c74bb37

File tree

9 files changed

+1139
-1045
lines changed

9 files changed

+1139
-1045
lines changed

python/setup.py

Lines changed: 50 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -53,12 +53,13 @@ def get_version(pkg_path):
5353

5454

5555
# Set some extra flags for compiling with coverage support
56-
extra_include_dirs = []
5756
extra_compile_args = []
5857
extra_link_args = []
5958
extra_define_macros = []
60-
library_dirs = []
61-
libraries = []
59+
device_include_dirs = []
60+
device_library_dirs = []
61+
device_libraries = []
62+
device_define_macros = []
6263

6364
if os.getenv("NANOARROW_PYTHON_COVERAGE") == "1":
6465
extra_compile_args.append("--coverage")
@@ -87,81 +88,57 @@ def get_version(pkg_path):
8788
lib_dirs_err = ", ".join(f"'{d}" for d in possible_libs)
8889
raise ValueError(f"Can't find CUDA library directory. Checked {lib_dirs_err}")
8990

90-
extra_include_dirs.append(str(include_dir))
91-
library_dirs.append(str(lib_dirs[0].parent))
92-
libraries.append("cuda")
93-
extra_define_macros.append(("NANOARROW_DEVICE_WITH_CUDA", 1))
91+
device_include_dirs.append(str(include_dir))
92+
device_library_dirs.append(str(lib_dirs[0].parent))
93+
device_libraries.append("cuda")
94+
device_define_macros.append(("NANOARROW_DEVICE_WITH_CUDA", 1))
95+
96+
97+
def nanoarrow_extension(
98+
name, *, nanoarrow_c=False, nanoarrow_device=False, nanoarrow_ipc=False
99+
):
100+
sources = ["src/" + name.replace(".", "/") + ".pyx"]
101+
libraries = []
102+
library_dirs = []
103+
include_dirs = ["src/nanoarrow", "vendor"]
104+
define_macros = list(extra_define_macros)
105+
106+
if nanoarrow_c:
107+
sources.append("vendor/nanoarrow.c")
108+
109+
if nanoarrow_device:
110+
sources.append("vendor/nanoarrow_device.c")
111+
include_dirs.extend(device_include_dirs)
112+
libraries.extend(device_libraries)
113+
library_dirs.extend(device_library_dirs)
114+
define_macros.extend(device_define_macros)
115+
116+
if nanoarrow_ipc:
117+
sources.extend(["vendor/nanoarrow_ipc.c", "vendor/flatcc.c"])
118+
119+
return Extension(
120+
name=name,
121+
include_dirs=include_dirs,
122+
language="c",
123+
sources=sources,
124+
extra_compile_args=extra_compile_args,
125+
extra_link_args=extra_link_args,
126+
define_macros=define_macros,
127+
library_dirs=library_dirs,
128+
libraries=libraries,
129+
)
94130

95131

96132
setup(
97133
ext_modules=[
98-
Extension(
99-
name="nanoarrow._device",
100-
include_dirs=["src/nanoarrow", "vendor"],
101-
language="c",
102-
sources=[
103-
"src/nanoarrow/_device.pyx",
104-
"vendor/nanoarrow.c",
105-
"vendor/nanoarrow_device.c",
106-
],
107-
extra_compile_args=extra_compile_args,
108-
extra_link_args=extra_link_args,
109-
define_macros=extra_define_macros,
110-
library_dirs=library_dirs,
111-
libraries=libraries,
112-
),
113-
Extension(
114-
name="nanoarrow._types",
115-
include_dirs=["src/nanoarrow", "vendor"],
116-
language="c",
117-
sources=[
118-
"src/nanoarrow/_types.pyx",
119-
],
120-
extra_compile_args=extra_compile_args,
121-
extra_link_args=extra_link_args,
122-
define_macros=extra_define_macros,
123-
),
124-
Extension(
125-
name="nanoarrow._utils",
126-
include_dirs=extra_include_dirs + ["src/nanoarrow", "vendor"],
127-
language="c",
128-
sources=[
129-
"src/nanoarrow/_utils.pyx",
130-
"vendor/nanoarrow.c",
131-
],
132-
extra_compile_args=extra_compile_args,
133-
extra_link_args=extra_link_args,
134-
define_macros=extra_define_macros,
135-
),
136-
Extension(
137-
name="nanoarrow._lib",
138-
include_dirs=extra_include_dirs + ["src/nanoarrow", "vendor"],
139-
language="c",
140-
sources=[
141-
"src/nanoarrow/_lib.pyx",
142-
"vendor/nanoarrow.c",
143-
"vendor/nanoarrow_device.c",
144-
],
145-
extra_compile_args=extra_compile_args,
146-
extra_link_args=extra_link_args,
147-
define_macros=extra_define_macros,
148-
library_dirs=library_dirs,
149-
libraries=libraries,
150-
),
151-
Extension(
152-
name="nanoarrow._ipc_lib",
153-
include_dirs=extra_include_dirs + ["src/nanoarrow", "vendor"],
154-
language="c",
155-
sources=[
156-
"src/nanoarrow/_ipc_lib.pyx",
157-
"vendor/nanoarrow.c",
158-
"vendor/nanoarrow_ipc.c",
159-
"vendor/flatcc.c",
160-
],
161-
extra_compile_args=extra_compile_args,
162-
extra_link_args=extra_link_args,
163-
define_macros=extra_define_macros,
134+
nanoarrow_extension("nanoarrow._types"),
135+
nanoarrow_extension("nanoarrow._utils", nanoarrow_c=True),
136+
nanoarrow_extension(
137+
"nanoarrow._device", nanoarrow_c=True, nanoarrow_device=True
164138
),
139+
nanoarrow_extension("nanoarrow._buffer", nanoarrow_c=True),
140+
nanoarrow_extension("nanoarrow._lib", nanoarrow_c=True, nanoarrow_device=True),
141+
nanoarrow_extension("nanoarrow._ipc_lib", nanoarrow_c=True, nanoarrow_ipc=True),
165142
],
166143
version=version,
167144
)

python/src/nanoarrow/_buffer.pxd

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
# cython: language_level = 3
19+
20+
from libc.stdint cimport int64_t
21+
22+
from nanoarrow_c cimport (
23+
ArrowBuffer,
24+
ArrowBufferView,
25+
ArrowType,
26+
)
27+
28+
from nanoarrow._device cimport Device
29+
30+
31+
cdef class CBufferView:
32+
cdef object _base
33+
cdef ArrowBufferView _ptr
34+
cdef ArrowType _data_type
35+
cdef Device _device
36+
cdef Py_ssize_t _element_size_bits
37+
cdef Py_ssize_t _shape
38+
cdef Py_ssize_t _strides
39+
cdef int64_t _n_elements
40+
cdef char _format[128]
41+
42+
cdef _check_copy_into_bounds(self, Py_buffer* dest, int64_t offset, int64_t length,
43+
int64_t dest_offset, int64_t dest_itemsize)
44+
45+
cdef Py_ssize_t _item_size(self)
46+
47+
cdef _do_getbuffer(self, Py_buffer *buffer, int flags)
48+
49+
cdef _do_releasebuffer(self, Py_buffer* buffer)
50+
51+
cdef class CBuffer:
52+
cdef object _base
53+
cdef ArrowBuffer* _ptr
54+
cdef ArrowType _data_type
55+
cdef int _element_size_bits
56+
cdef char _format[32]
57+
cdef Device _device
58+
cdef CBufferView _view
59+
cdef int _get_buffer_count
60+
61+
cdef _assert_valid(self)
62+
63+
cdef _assert_buffer_count_zero(self)
64+
65+
cdef _populate_view(self)

0 commit comments

Comments
 (0)