Skip to content

Commit 85581ee

Browse files
mxwlimxwli
authored
Python UDF and C++ UDF improvements (#3483)
1 parent ceabb4a commit 85581ee

6 files changed

Lines changed: 121 additions & 13 deletions

File tree

‎src_cpp/include/py_connection.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,8 @@ class PyConnection {
4242
static bool isPandasDataframe(const py::object& object);
4343

4444
void createScalarFunction(const std::string& name, const py::function& udf,
45-
const py::list& params, const std::string& retval);
45+
const py::list& params, const std::string& retval, bool defaultNull, bool catchExceptions);
46+
void removeScalarFunction(const std::string& name);
4647

4748
static Value transformPythonValue(const py::handle& val);
4849
static Value transformPythonValueAs(const py::handle& val, const LogicalType& type);

‎src_cpp/include/py_udf.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,5 +13,6 @@ class PyUDF {
1313

1414
public:
1515
static function_set toFunctionSet(const std::string& name, const py::function& udf,
16-
const py::list& paramTypes, const std::string& resultType);
16+
const py::list& paramTypes, const std::string& resultType, bool defaultNull,
17+
bool catchExceptions);
1718
};

‎src_cpp/py_connection.cpp‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,9 @@ void PyConnection::initialize(py::handle& m) {
3636
py::arg("np_array"), py::arg("src_table_name"), py::arg("rel_name"),
3737
py::arg("dst_table_name"), py::arg("query_batch_size"))
3838
.def("create_function", &PyConnection::createScalarFunction, py::arg("name"),
39-
py::arg("udf"), py::arg("params_type"), py::arg("return_value"));
39+
py::arg("udf"), py::arg("params_type"), py::arg("return_value"),
40+
py::arg("default_null"), py::arg("catch_exceptions"))
41+
.def("remove_function", &PyConnection::removeScalarFunction, py::arg("name"));
4042
PyDateTime_IMPORT;
4143
}
4244

@@ -478,6 +480,11 @@ std::unique_ptr<PyQueryResult> PyConnection::checkAndWrapQueryResult(
478480
}
479481

480482
void PyConnection::createScalarFunction(const std::string& name, const py::function& udf,
481-
const py::list& params, const std::string& retval) {
482-
conn->addUDFFunctionSet(name, PyUDF::toFunctionSet(name, udf, params, retval));
483+
const py::list& params, const std::string& retval, bool defaultNull, bool catchExceptions) {
484+
conn->addUDFFunctionSet(name,
485+
PyUDF::toFunctionSet(name, udf, params, retval, defaultNull, catchExceptions));
486+
}
487+
488+
void PyConnection::removeScalarFunction(const std::string& name) {
489+
conn->removeUDFFunction(name);
483490
}

‎src_cpp/py_udf.cpp‎

Lines changed: 26 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -124,7 +124,8 @@ static PyUDFSignature analyzeSignature(const py::function& udf) {
124124
return UDFSignature;
125125
}
126126

127-
static scalar_func_exec_t getUDFExecFunc(const py::function& udf) {
127+
static scalar_func_exec_t getUDFExecFunc(const py::function& udf, bool defaultNull,
128+
bool catchExceptions) {
128129
return [=](const std::vector<std::shared_ptr<ValueVector>>& params, ValueVector& result,
129130
void* /* dataPtr */) -> void {
130131
py::gil_scoped_acquire acquire;
@@ -133,16 +134,33 @@ static scalar_func_exec_t getUDFExecFunc(const py::function& udf) {
133134
for (auto i = 0u; i < resultSelVector.getSelSize(); ++i) {
134135
auto resultPos = resultSelVector[i];
135136
py::list pyParams;
137+
bool hasNull = false;
136138
for (const auto& param : params) {
137139
auto paramPos =
138140
param->state->isFlat() ? param->state->getSelVector()[0] : resultPos;
139141
auto value = param->getAsValue(paramPos);
142+
if (value->isNull()) {
143+
hasNull = true;
144+
}
140145
auto pyValue = PyQueryResult::convertValueToPyObject(*value);
141146
pyParams.append(pyValue);
142147
}
143-
auto pyResult = udf(*pyParams);
144-
auto resultValue = PyConnection::transformPythonValueAs(pyResult, result.dataType);
145-
result.copyFromValue(resultPos, resultValue);
148+
if (defaultNull && hasNull) {
149+
result.setNull(resultPos, true);
150+
} else {
151+
try {
152+
auto pyResult = udf(*pyParams);
153+
auto resultValue =
154+
PyConnection::transformPythonValueAs(pyResult, result.dataType);
155+
result.copyFromValue(resultPos, resultValue);
156+
} catch (py::error_already_set& e) {
157+
if (catchExceptions) {
158+
result.setNull(resultPos, true);
159+
} else {
160+
throw common::RuntimeException(e.what());
161+
}
162+
}
163+
}
146164
}
147165
};
148166
}
@@ -155,7 +173,8 @@ static scalar_bind_func getUDFBindFunc(const PyUDFSignature& signature) {
155173
}
156174

157175
function_set PyUDF::toFunctionSet(const std::string& name, const py::function& udf,
158-
const py::list& paramTypes, const std::string& resultType) {
176+
const py::list& paramTypes, const std::string& resultType, bool defaultNull,
177+
bool catchExceptions) {
159178
auto pySignature = analyzeSignature(udf);
160179
auto explicitParamTypes = pyListToParams(paramTypes);
161180
if (explicitParamTypes.size() > 0) {
@@ -183,7 +202,7 @@ function_set PyUDF::toFunctionSet(const std::string& name, const py::function& u
183202

184203
function_set definitions;
185204
definitions.push_back(std::make_unique<ScalarFunction>(name, paramIDTypes,
186-
pySignature.resultType.getLogicalTypeID(), getUDFExecFunc(udf),
187-
getUDFBindFunc(pySignature)));
205+
pySignature.resultType.getLogicalTypeID(),
206+
getUDFExecFunc(udf, defaultNull, catchExceptions), getUDFBindFunc(pySignature)));
188207
return definitions;
189208
}

‎src_py/connection.py‎

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -235,6 +235,9 @@ def create_function(
235235
udf: Callable[[...], Any],
236236
params_type: list[Type | str] | None = None,
237237
return_type: Type | str = "",
238+
*,
239+
default_null_handling: bool = True,
240+
catch_exceptions: bool = False
238241
) -> None:
239242
"""
240243
Sets a User Defined Function (UDF) to use in cypher queries.
@@ -252,10 +255,35 @@ def create_function(
252255
253256
return_type: Optional[Type]
254257
a Type enum to describe the returned value
258+
259+
default_null_handling: Optional[bool]
260+
if true, when any parameter is null, the resulting value will be null
261+
262+
catch_exceptions: Optional[bool]
263+
if true, when an exception is thrown from python, the function output will be null
264+
Otherwise, the exception will be rethrown
255265
"""
256266
if params_type is None:
257267
params_type = []
258268
parsed_params_type = [x if type(x) is str else x.value for x in params_type]
259269
if type(return_type) is not str:
260270
return_type = return_type.value
261-
self._connection.create_function(name, udf, parsed_params_type, return_type)
271+
272+
self._connection.create_function(
273+
name=name,
274+
udf=udf,
275+
params_type=parsed_params_type,
276+
return_value=return_type,
277+
default_null=default_null_handling,
278+
catch_exceptions=catch_exceptions)
279+
280+
def remove_function(self, name: str) -> None:
281+
"""
282+
Removes a User Defined Function (UDF).
283+
284+
Parameters
285+
----------
286+
name: str
287+
name of function to be removed.
288+
"""
289+
self._connection.remove_function(name)

‎test/test_udf.py‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,10 @@
44

55
import pandas as pd
66
import pyarrow as pa
7+
import pytest
8+
from datetime import date, datetime, timedelta # noqa: TCH003
9+
10+
import kuzu
711
from kuzu import Type
812
from type_aliases import ConnDB
913

@@ -141,3 +145,51 @@ def selectIfSeven(a: int) -> bool:
141145
[("a", 1), ("b", 2), ("c", 3), ("x", -1), ("y", -2), ("z", -3)],
142146
[("l", 1), ("m", 2), ("n", 3), ("one", -1), ("two", -2), ("three", -3)],
143147
]
148+
149+
def test_udf_null(conn_db_readwrite: ConnDB) -> None:
150+
conn, db = conn_db_readwrite
151+
152+
def get5(x: int) -> int:
153+
return 5
154+
155+
conn.create_function("get5", get5)
156+
assert conn.execute("RETURN get5(NULL)").get_next() == [None]
157+
158+
conn.remove_function("get5")
159+
conn.create_function("get5", get5, default_null_handling = False)
160+
assert conn.execute("RETURN get5(NULL)").get_next() == [5]
161+
162+
def test_udf_except(conn_db_readwrite: ConnDB) -> None:
163+
class TestException(Exception):
164+
pass
165+
166+
conn, db = conn_db_readwrite
167+
168+
def throw() -> int:
169+
errmsg = "test"
170+
raise TestException(errmsg)
171+
172+
conn.create_function("testexcept", throw)
173+
174+
pytest.raises(RuntimeError, conn.execute, "RETURN testexcept()")
175+
176+
conn.remove_function("testexcept")
177+
conn.create_function("testexcept", throw, catch_exceptions = True)
178+
179+
assert conn.execute("RETURN testexcept()").get_next() == [None]
180+
181+
def test_udf_remove(conn_db_readwrite: ConnDB) -> None:
182+
conn, db = conn_db_readwrite
183+
184+
def myfunction() -> int:
185+
return 1
186+
187+
conn.create_function("myfunction", myfunction)
188+
189+
with pytest.raises(RuntimeError, match="Catalog exception: function notmyfunction doesn't exist."):
190+
conn.remove_function("notmyfunction")
191+
192+
conn.remove_function("myfunction")
193+
194+
with pytest.raises(RuntimeError, match="Catalog exception: function list_create doesn't exist."):
195+
conn.remove_function("list_create")

0 commit comments

Comments
 (0)