Skip to content

Commit 2ab90d7

Browse files
committed
Fix memory leaks in C extension; add test functions
1 parent 3de1f6f commit 2ab90d7

File tree

2 files changed

+93
-9
lines changed

2 files changed

+93
-9
lines changed

accel.c

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -363,18 +363,15 @@ static PyObject *apply_transformer(PyObject *transformer, PyObject *value) {
363363
}
364364

365365
if (transformer == NULL || transformer == Py_None) {
366-
// Increment refcount fo Py_None since we are returning it
367-
if (value == Py_None) {
368-
Py_INCREF(value);
369-
}
366+
// No transformation needed, return value as-is
367+
// We steal the reference from the caller
370368
return value;
371369
}
372370

373371
PyObject *out = PyObject_CallFunction(transformer, "O", value);
374-
// Don't decref Py_None since we assume it was passed in literally and not from a new reference
375-
if (value != NULL && value != Py_None) {
376-
Py_DECREF(value);
377-
}
372+
// Always decref value since we're stealing the reference from the caller
373+
// This includes Py_None when it's passed as a new reference (e.g., from PyIter_Next)
374+
Py_DECREF(value);
378375

379376
return out;
380377
}
@@ -4306,7 +4303,10 @@ static PyObject *load_rowdat_1(PyObject *self, PyObject *args, PyObject *kwargs)
43064303
if (!py_row) goto error;
43074304

43084305
row_id = *(int64_t*)data; data += 8;
4309-
CHECKRC(PyList_Append(py_out_row_ids, PyLong_FromLongLong(row_id)));
4306+
PyObject *py_row_id = PyLong_FromLongLong(row_id);
4307+
if (!py_row_id) goto error;
4308+
CHECKRC(PyList_Append(py_out_row_ids, py_row_id));
4309+
Py_DECREF(py_row_id);
43104310

43114311
for (unsigned long long i = 0; i < colspec_l; i++) {
43124312
is_null = data[0] == '\x01'; data += 1;

singlestoredb/tests/ext_funcs/__init__.py

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,90 @@
3131
from singlestoredb.functions.typing import pyarrow as pat
3232

3333

34+
@udf
35+
def add(x: int, y: int) -> int:
36+
"""
37+
Add two integers.
38+
39+
Parameters
40+
----------
41+
x : int
42+
First integer.
43+
y : int
44+
Second integer.
45+
46+
Returns
47+
-------
48+
int
49+
Sum of x and y.
50+
51+
"""
52+
return x + y
53+
54+
55+
@udf
56+
def add_vec(x: npt.Int64Array, y: npt.Int64Array) -> npt.Int64Array:
57+
"""
58+
Add two numpy arrays of int64.
59+
60+
Parameters
61+
----------
62+
x : numpy.ndarray
63+
First array.
64+
y : numpy.ndarray
65+
Second array.
66+
67+
Returns
68+
-------
69+
numpy.ndarray
70+
Elementwise sum of x and y.
71+
72+
"""
73+
return x + y
74+
75+
76+
@udf
77+
async def async_add(x: int, y: int) -> int:
78+
"""
79+
Asynchronously add two integers.
80+
81+
Parameters
82+
----------
83+
x : int
84+
First integer.
85+
y : int
86+
Second integer.
87+
88+
Returns
89+
-------
90+
int
91+
Sum of x and y.
92+
93+
"""
94+
return x + y
95+
96+
97+
@udf
98+
async def async_add_vec_vec(x: npt.Int64Array, y: npt.Int64Array) -> npt.Int64Array:
99+
"""
100+
Asynchronously add two numpy arrays of int64.
101+
102+
Parameters
103+
----------
104+
x : numpy.ndarray
105+
First array.
106+
y : numpy.ndarray
107+
Second array.
108+
109+
Returns
110+
-------
111+
numpy.ndarray
112+
Elementwise sum of x and y.
113+
114+
"""
115+
return x + y
116+
117+
34118
@udf
35119
def doc_test(x: int, y: float) -> int:
36120
"""

0 commit comments

Comments
 (0)