improve allocation for numpy

Author: evertlammerts

.github/workflows/packaging_wheels.yml

@@ -30,7 +30,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python: [ cp314 ]
+        python: [ cp311, cp314 ]
         platform:
           - { os: windows-2022,     arch: amd64,      cibw_system: win }
           - { os: windows-11-arm,   arch: ARM64,      cibw_system: win }

CMakeLists.txt

@@ -39,7 +39,7 @@ endif()
 # internally)
 find_package(
   Python
-  COMPONENTS Interpreter Development.Module
+  COMPONENTS Interpreter Development.Module NumPy
   REQUIRED)
 # Nanobind ships its CMake config inside site-packages/nanobind/cmake, so
 # find_package() can't discover it unless we set it. (scikit-build-core does

pyproject.toml

@@ -65,6 +65,10 @@ requires = [
     "scikit-build-core>=0.11.4",
 	"nanobind>=2.0",
     "setuptools_scm>=8.0",
+    # numpy C API headers (PyArray_Empty in the result path). Building against numpy 2.x yields a
+    # binary compatible with numpy >=1.19 AND 2.x at runtime (numpy 2.0 backward-compat), so the
+    # unpinned runtime numpy range is preserved. Build-time only; the runtime numpy dep is unchanged.
+    "numpy>=2.0",
 ]
 
 [tool.scikit-build]

src/include/duckdb_python/numpy/numpy_array.hpp

@@ -11,8 +11,6 @@
 #include "duckdb_python/nb/casters.hpp"
 #include "duckdb.hpp"
 
-#include <unordered_map>
-
 namespace duckdb {
 
 namespace numpy_internal {
@@ -37,24 +35,12 @@ inline PyTypeObject *NumpyNdarrayType() {
 	return cached;
 }
 
-//! Allocate an uninitialized 1-D numpy array of `count` elements with the given numpy dtype string.
-//! `numpy.empty` and the `np.dtype` objects are cached to avoid a module import, attribute lookup,
-//! and dtype-string parse on every call. This is hot: a LIST/ARRAY column allocates one array per
-//! row. Cached handles are leaked for process lifetime (shutdown-safe: no Python destructor runs
-//! after finalization). Only ever called on the single-threaded, GIL-held result path.
-inline nb::object NumpyEmpty(idx_t count, const string &dtype) {
-	static PyObject *empty_fn = []() -> PyObject * {
-		nb::object fn = nb::module_::import_("numpy").attr("empty");
-		return fn.release().ptr();
-	}();
-	static auto &dtype_cache = *new std::unordered_map<string, PyObject *>();
-	PyObject *&descr = dtype_cache[dtype];
-	if (!descr) {
-		nb::object d = nb::module_::import_("numpy").attr("dtype")(dtype);
-		descr = d.release().ptr();
-	}
-	return nb::borrow<nb::object>(empty_fn)(count, nb::handle(descr));
-}
+//! Allocate an uninitialized 1-D numpy array of `count` elements with the given numpy dtype string
+//! (e.g. "int64", "float32", "object", "datetime64[us]") via the numpy C API (PyArray_Empty). The
+//! parsed np.dtype objects are cached to avoid a dtype-string parse on every call. This is hot: a
+//! LIST/ARRAY column allocates one array per row. Defined in numpy_array.cpp (the single TU that
+//! pulls in the numpy C API). Only ever called on the single-threaded, GIL-held result path.
+nb::object NumpyEmpty(idx_t count, const string &dtype);
 
 } // namespace numpy_internal
 

src/numpy/CMakeLists.txt

@@ -1,7 +1,17 @@
 # this is used for clang-tidy checks
 add_library(
   python_numpy OBJECT
-  type.cpp numpy_scan.cpp array_wrapper.cpp raw_array_wrapper.cpp
-  numpy_bind.cpp numpy_result_conversion.cpp)
+  type.cpp
+  numpy_scan.cpp
+  array_wrapper.cpp
+  raw_array_wrapper.cpp
+  numpy_bind.cpp
+  numpy_result_conversion.cpp
+  numpy_array.cpp)
 
 target_link_libraries(python_numpy PRIVATE _duckdb_dependencies)
+
+# numpy_array.cpp is the single TU that uses the numpy C API (PyArray_Empty), so
+# it needs numpy's headers. Resolved by find_package(Python ... COMPONENTS ...
+# NumPy) in the top-level CMakeLists. Scoped to this object library only.
+target_include_directories(python_numpy PRIVATE ${Python_NumPy_INCLUDE_DIRS})

src/numpy/numpy_array.cpp

@@ -0,0 +1,76 @@
+//===----------------------------------------------------------------------===//
+//                         DuckDB
+//
+// numpy_array.cpp
+//
+// Out-of-line definitions for the NumpyArray facade (numpy_array.hpp). This is the
+// ONLY translation unit that uses the numpy C API, so it does not need
+// PY_ARRAY_UNIQUE_SYMBOL / NO_IMPORT_ARRAY (those coordinate the C-API function
+// pointer table across multiple TUs).
+//===----------------------------------------------------------------------===//
+
+#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
+#include <numpy/arrayobject.h>
+
+#include "duckdb_python/numpy/numpy_array.hpp"
+
+#include <stdexcept>
+#include <unordered_map>
+
+namespace duckdb {
+namespace numpy_internal {
+
+namespace {
+
+//! Lazy, guarded one-time init of the numpy C-API function pointer table. numpy is always
+//! already imported by the time we allocate a result array, so import_array should succeed;
+//! if it does not, the returned value is false and the caller raises. Runs exactly once
+//! (function-local static initializer, GIL held on the result path).
+bool EnsureNumpyCApi() {
+	static bool ok = []() -> bool {
+		// import_array1(ret) expands to `return ret;` on failure, so wrap it in a lambda that
+		// returns int and surface success via the return value.
+		auto do_import = []() -> int {
+			import_array1(-1);
+			return 0;
+		};
+		return do_import() == 0;
+	}();
+	return ok;
+}
+
+} // namespace
+
+nb::object NumpyEmpty(idx_t count, const string &dtype) {
+	// Process-lifetime cache of parsed np.dtype objects, keyed by dtype string. The parse is
+	// otherwise repeated per call; a LIST/ARRAY column allocates one array per row. Leaked on
+	// purpose (numpy is never unloaded; no Python destructor runs after finalization). Only ever
+	// touched on the single-threaded, GIL-held result path.
+	static auto &dtype_cache = *new std::unordered_map<string, PyObject *>();
+	PyObject *&descr = dtype_cache[dtype];
+	if (!descr) {
+		nb::object d = nb::module_::import_("numpy").attr("dtype")(dtype);
+		descr = d.release().ptr();
+	}
+
+	if (!EnsureNumpyCApi()) {
+		throw std::runtime_error("Failed to initialize the numpy C API (import_array failed)");
+	}
+
+	npy_intp dims[1] = {static_cast<npy_intp>(count)};
+	// PyArray_Empty STEALS a reference to descr. descr is a single cached np.dtype reused across
+	// every allocation, so hand PyArray_Empty its own reference to consume.
+	Py_INCREF(descr);
+	PyObject *arr = PyArray_Empty(1, dims, reinterpret_cast<PyArray_Descr *>(descr), 0 /* C order */);
+	if (!arr) {
+		// PyArray_Empty consumed the stolen reference even on failure; balance the INCREF above so
+		// the cached descr is not leaked, then surface the numpy error.
+		Py_DECREF(descr);
+		throw nb::python_error();
+	}
+	// PyArray_Empty returns a NEW reference; hand ownership to nanobind via steal.
+	return nb::steal<nb::object>(arr);
+}
+
+} // namespace numpy_internal
+} // namespace duckdb