Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 18 additions & 18 deletions src/python_udf.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -530,26 +530,26 @@ struct PythonUDFData {
ScalarFunction GetFunction(const nb::callable &udf, PythonExceptionHandling exception_handling, bool side_effects,
const ClientProperties &client_properties) {

// Import this module, because importing this from a non-main thread causes a segfault

auto &import_cache = *DuckDBPyConnection::ImportCache();
nb::handle core;
auto numpy = import_cache.numpy();
if (!numpy) {
throw InvalidInputException("'numpy' is required for this operation, but it wasn't installed");
}
// numpy.__version__ is a string; nb::cast<nb::tuple> rejects a non-tuple, so convert it explicitly.
nb::object numpy_version_str = numpy.attr("__version__");
auto numpy_version = nb::tuple(numpy_version_str);
if (NumpyDeprecatesAccessToCore(numpy_version)) {
core = numpy.attr("_core");
} else {
core = numpy.attr("core");
}
(void)core.attr("multiarray");

scalar_function_t func;
if (vectorized) {
// Only the vectorized (pyarrow) path needs numpy; import it here rather than before
// the branch. Importing off the main thread causes a segfault.
auto &import_cache = *DuckDBPyConnection::ImportCache();
nb::handle core;
auto numpy = import_cache.numpy();
if (!numpy) {
throw InvalidInputException("'numpy' is required for this operation, but it wasn't installed");
}
// numpy.__version__ is a string; nb::cast<nb::tuple> rejects a non-tuple, so convert it explicitly.
nb::object numpy_version_str = numpy.attr("__version__");
auto numpy_version = nb::tuple(numpy_version_str);
if (NumpyDeprecatesAccessToCore(numpy_version)) {
core = numpy.attr("_core");
} else {
core = numpy.attr("core");
}
(void)core.attr("multiarray");

func = CreateVectorizedFunction(udf.ptr(), exception_handling, null_handling);
} else {
func = CreateNativeFunction(udf.ptr(), exception_handling, client_properties, null_handling);
Expand Down
21 changes: 21 additions & 0 deletions tests/fast/udf/test_scalar_arrow.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
import subprocess
import sys

import pytest

import duckdb
Expand Down Expand Up @@ -214,3 +217,21 @@ def func(data):

res = duckdb_cursor.sql("select func(1).y").fetchone()
assert res == ("this is not an inlined string",)

def test_arrow_udf_requires_numpy(self):
# Subprocess isolation, same reasoning as test_native_udf_without_numpy in
# test_scalar_native.py (numpy's import cache is a process-lifetime singleton).
script = """
import sys
sys.modules["numpy"] = None
import duckdb
from duckdb.sqltypes import VARCHAR

try:
duckdb.create_function("arrow_no_numpy", lambda: "x", [], VARCHAR, type="arrow")
raise SystemExit("expected InvalidInputException, but create_function succeeded")
except duckdb.InvalidInputException as e:
assert "numpy" in str(e), f"unexpected error message: {e}"
"""
result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True)
assert result.returncode == 0, result.stderr
23 changes: 23 additions & 0 deletions tests/fast/udf/test_scalar_native.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
import subprocess
import sys

import pytest

import duckdb
Expand Down Expand Up @@ -254,3 +257,23 @@ def example():
(res,) = duckdb_cursor.sql("SELECT example()").fetchone()
for key, val in res.items():
assert key == val

def test_native_udf_without_numpy(self):
# Runs in a subprocess because the numpy import cache is a process-lifetime singleton -
# once something else in this interpreter has imported numpy, sys.modules trickery here
# can't make it forget that.
script = """
import sys
sys.modules["numpy"] = None
import duckdb
from duckdb.sqltypes import VARCHAR

def generate_name():
return "hello-world"

duckdb.create_function("native_no_numpy", generate_name, [], VARCHAR)
res = duckdb.sql("select native_no_numpy()").fetchall()
assert res == [("hello-world",)]
"""
result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True)
assert result.returncode == 0, result.stderr