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).

python/setup.py

@@ -53,12 +53,13 @@ def get_version(pkg_path):
 
 
 # Set some extra flags for compiling with coverage support
-extra_include_dirs = []
 extra_compile_args = []
 extra_link_args = []
 extra_define_macros = []
-library_dirs = []
-libraries = []
+device_include_dirs = []
+device_library_dirs = []
+device_libraries = []
+device_define_macros = []
 
 if os.getenv("NANOARROW_PYTHON_COVERAGE") == "1":
     extra_compile_args.append("--coverage")
@@ -87,81 +88,57 @@ def get_version(pkg_path):
         lib_dirs_err = ", ".join(f"'{d}" for d in possible_libs)
         raise ValueError(f"Can't find CUDA library directory. Checked {lib_dirs_err}")
 
-    extra_include_dirs.append(str(include_dir))
-    library_dirs.append(str(lib_dirs[0].parent))
-    libraries.append("cuda")
-    extra_define_macros.append(("NANOARROW_DEVICE_WITH_CUDA", 1))
+    device_include_dirs.append(str(include_dir))
+    device_library_dirs.append(str(lib_dirs[0].parent))
+    device_libraries.append("cuda")
+    device_define_macros.append(("NANOARROW_DEVICE_WITH_CUDA", 1))
+
+
+def nanoarrow_extension(
+    name, *, nanoarrow_c=False, nanoarrow_device=False, nanoarrow_ipc=False
+):
+    sources = ["src/" + name.replace(".", "/") + ".pyx"]
+    libraries = []
+    library_dirs = []
+    include_dirs = ["src/nanoarrow", "vendor"]
+    define_macros = list(extra_define_macros)
+
+    if nanoarrow_c:
+        sources.append("vendor/nanoarrow.c")
+
+    if nanoarrow_device:
+        sources.append("vendor/nanoarrow_device.c")
+        include_dirs.extend(device_include_dirs)
+        libraries.extend(device_libraries)
+        library_dirs.extend(device_library_dirs)
+        define_macros.extend(device_define_macros)
+
+    if nanoarrow_ipc:
+        sources.extend(["vendor/nanoarrow_ipc.c", "vendor/flatcc.c"])
+
+    return Extension(
+        name=name,
+        include_dirs=include_dirs,
+        language="c",
+        sources=sources,
+        extra_compile_args=extra_compile_args,
+        extra_link_args=extra_link_args,
+        define_macros=define_macros,
+        library_dirs=library_dirs,
+        libraries=libraries,
+    )
 
 
 setup(
     ext_modules=[
-        Extension(
-            name="nanoarrow._device",
-            include_dirs=["src/nanoarrow", "vendor"],
-            language="c",
-            sources=[
-                "src/nanoarrow/_device.pyx",
-                "vendor/nanoarrow.c",
-                "vendor/nanoarrow_device.c",
-            ],
-            extra_compile_args=extra_compile_args,
-            extra_link_args=extra_link_args,
-            define_macros=extra_define_macros,
-            library_dirs=library_dirs,
-            libraries=libraries,
-        ),
-        Extension(
-            name="nanoarrow._types",
-            include_dirs=["src/nanoarrow", "vendor"],
-            language="c",
-            sources=[
-                "src/nanoarrow/_types.pyx",
-            ],
-            extra_compile_args=extra_compile_args,
-            extra_link_args=extra_link_args,
-            define_macros=extra_define_macros,
-        ),
-        Extension(
-            name="nanoarrow._utils",
-            include_dirs=extra_include_dirs + ["src/nanoarrow", "vendor"],
-            language="c",
-            sources=[
-                "src/nanoarrow/_utils.pyx",
-                "vendor/nanoarrow.c",
-            ],
-            extra_compile_args=extra_compile_args,
-            extra_link_args=extra_link_args,
-            define_macros=extra_define_macros,
-        ),
-        Extension(
-            name="nanoarrow._lib",
-            include_dirs=extra_include_dirs + ["src/nanoarrow", "vendor"],
-            language="c",
-            sources=[
-                "src/nanoarrow/_lib.pyx",
-                "vendor/nanoarrow.c",
-                "vendor/nanoarrow_device.c",
-            ],
-            extra_compile_args=extra_compile_args,
-            extra_link_args=extra_link_args,
-            define_macros=extra_define_macros,
-            library_dirs=library_dirs,
-            libraries=libraries,
-        ),
-        Extension(
-            name="nanoarrow._ipc_lib",
-            include_dirs=extra_include_dirs + ["src/nanoarrow", "vendor"],
-            language="c",
-            sources=[
-                "src/nanoarrow/_ipc_lib.pyx",
-                "vendor/nanoarrow.c",
-                "vendor/nanoarrow_ipc.c",
-                "vendor/flatcc.c",
-            ],
-            extra_compile_args=extra_compile_args,
-            extra_link_args=extra_link_args,
-            define_macros=extra_define_macros,
+        nanoarrow_extension("nanoarrow._types"),
+        nanoarrow_extension("nanoarrow._utils", nanoarrow_c=True),
+        nanoarrow_extension(
+            "nanoarrow._device", nanoarrow_c=True, nanoarrow_device=True
         ),
+        nanoarrow_extension("nanoarrow._buffer", nanoarrow_c=True),
+        nanoarrow_extension("nanoarrow._lib", nanoarrow_c=True, nanoarrow_device=True),
+        nanoarrow_extension("nanoarrow._ipc_lib", nanoarrow_c=True, nanoarrow_ipc=True),
     ],
     version=version,
 )

python/src/nanoarrow/_buffer.pxd

@@ -0,0 +1,65 @@
+# 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
+
+from libc.stdint cimport int64_t
+
+from nanoarrow_c cimport (
+    ArrowBuffer,
+    ArrowBufferView,
+    ArrowType,
+)
+
+from nanoarrow._device cimport Device
+
+
+cdef class CBufferView:
+    cdef object _base
+    cdef ArrowBufferView _ptr
+    cdef ArrowType _data_type
+    cdef Device _device
+    cdef Py_ssize_t _element_size_bits
+    cdef Py_ssize_t _shape
+    cdef Py_ssize_t _strides
+    cdef int64_t _n_elements
+    cdef char _format[128]
+
+    cdef _check_copy_into_bounds(self, Py_buffer* dest, int64_t offset, int64_t length,
+                                 int64_t dest_offset, int64_t dest_itemsize)
+
+    cdef Py_ssize_t _item_size(self)
+
+    cdef _do_getbuffer(self, Py_buffer *buffer, int flags)
+
+    cdef _do_releasebuffer(self, Py_buffer* buffer)
+
+cdef class CBuffer:
+    cdef object _base
+    cdef ArrowBuffer* _ptr
+    cdef ArrowType _data_type
+    cdef int _element_size_bits
+    cdef char _format[32]
+    cdef Device _device
+    cdef CBufferView _view
+    cdef int _get_buffer_count
+
+    cdef _assert_valid(self)
+
+    cdef _assert_buffer_count_zero(self)
+
+    cdef _populate_view(self)