Skip to content

Commit 186c1c2

Browse files
kesmit13kesmit13
committed
Fix error and memory handling
1 parent e2363ce commit 186c1c2

File tree

1 file changed

+70
-21
lines changed

1 file changed

+70
-21
lines changed

‎accel.c‎

Lines changed: 70 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -2678,8 +2678,19 @@ static PyObject *load_rowdat_1_numpy(PyObject *self, PyObject *args, PyObject *k
26782678

26792679
exit:
26802680
if (ctypes) free(ctypes);
2681-
if (out_cols) free(out_cols);
2682-
if (mask_cols) free(mask_cols);
2681+
if (out_cols) {
2682+
for (i = 0; i < n_cols; i++) {
2683+
if (out_cols[i]) free(out_cols[i]);
2684+
}
2685+
free(out_cols);
2686+
}
2687+
if (mask_cols) {
2688+
for (i = 0; i < n_cols; i++) {
2689+
if (mask_cols[i]) free(mask_cols[i]);
2690+
}
2691+
free(mask_cols);
2692+
}
2693+
if (out_row_ids) free(out_row_ids);
26832694
if (data_formats) free(data_formats);
26842695
if (item_sizes) free(item_sizes);
26852696

@@ -2943,11 +2954,17 @@ static PyObject *dump_rowdat_1_numpy(PyObject *self, PyObject *args, PyObject *k
29432954
out_l = 256 * n_cols;
29442955
out_idx = 0;
29452956
out = malloc(out_l);
2946-
if (!out) goto error;
2957+
if (!out) {
2958+
PyErr_SetString(PyExc_MemoryError, "failed to allocate output buffer");
2959+
goto error;
2960+
}
29472961

29482962
// Get return types
29492963
returns = malloc(sizeof(int) * n_cols);
2950-
if (!returns) goto error;
2964+
if (!returns) {
2965+
PyErr_SetString(PyExc_MemoryError, "failed to allocate returns array");
2966+
goto error;
2967+
}
29512968

29522969
for (i = 0; i < n_cols; i++) {
29532970
PyObject *py_item = PySequence_GetItem(py_returns, i);
@@ -2959,11 +2976,20 @@ static PyObject *dump_rowdat_1_numpy(PyObject *self, PyObject *args, PyObject *k
29592976

29602977
// Get column array memory
29612978
cols = calloc(sizeof(char*), n_cols);
2962-
if (!cols) goto error;
2979+
if (!cols) {
2980+
PyErr_SetString(PyExc_MemoryError, "failed to allocate cols array");
2981+
goto error;
2982+
}
29632983
col_types = calloc(sizeof(NumpyColType), n_cols);
2964-
if (!col_types) goto error;
2984+
if (!col_types) {
2985+
PyErr_SetString(PyExc_MemoryError, "failed to allocate col_types array");
2986+
goto error;
2987+
}
29652988
masks = calloc(sizeof(char*), n_cols);
2966-
if (!masks) goto error;
2989+
if (!masks) {
2990+
PyErr_SetString(PyExc_MemoryError, "failed to allocate masks array");
2991+
goto error;
2992+
}
29672993
for (i = 0; i < n_cols; i++) {
29682994
PyObject *py_item = PyList_GetItem(py_cols, i);
29692995
if (!py_item) goto error;
@@ -2996,8 +3022,12 @@ static PyObject *dump_rowdat_1_numpy(PyObject *self, PyObject *args, PyObject *k
29963022
#define CHECKMEM(x) \
29973023
if ((out_idx + x) > out_l) { \
29983024
out_l = out_l * 2 + x; \
2999-
out = realloc(out, out_l); \
3000-
if (!out) goto error; \
3025+
char *new_out = realloc(out, out_l); \
3026+
if (!new_out) { \
3027+
PyErr_SetString(PyExc_MemoryError, "failed to reallocate output buffer"); \
3028+
goto error; \
3029+
} \
3030+
out = new_out; \
30013031
}
30023032

30033033
for (j = 0; j < n_rows; j++) {
@@ -4079,10 +4109,10 @@ static PyObject *dump_rowdat_1_numpy(PyObject *self, PyObject *args, PyObject *k
40794109
}
40804110
}
40814111

4082-
py_out = PyMemoryView_FromMemory(out, out_idx, PyBUF_WRITE);
4083-
if (!py_out) goto error;
4112+
py_out = PyBytes_FromStringAndSize(out, out_idx);
40844113

40854114
exit:
4115+
if (out) free(out);
40864116
if (returns) free(returns);
40874117
if (masks) free(masks);
40884118
if (cols) free(cols);
@@ -4091,7 +4121,6 @@ static PyObject *dump_rowdat_1_numpy(PyObject *self, PyObject *args, PyObject *k
40914121
return py_out;
40924122

40934123
error:
4094-
if (!py_out && out) free(out);
40954124
Py_XDECREF(py_out);
40964125
py_out = NULL;
40974126

@@ -4471,8 +4500,12 @@ static PyObject *dump_rowdat_1(PyObject *self, PyObject *args, PyObject *kwargs)
44714500
#define CHECKMEM(x) \
44724501
if ((out_idx + x) > out_l) { \
44734502
out_l = out_l * 2 + x; \
4474-
out = realloc(out, out_l); \
4475-
if (!out) goto error; \
4503+
char *new_out = realloc(out, out_l); \
4504+
if (!new_out) { \
4505+
PyErr_SetString(PyExc_MemoryError, "failed to reallocate output buffer"); \
4506+
goto error; \
4507+
} \
4508+
out = new_out; \
44764509
}
44774510

44784511
py_rows_iter = PyObject_GetIter(py_rows);
@@ -4483,12 +4516,20 @@ static PyObject *dump_rowdat_1(PyObject *self, PyObject *args, PyObject *kwargs)
44834516

44844517
while ((py_row = PyIter_Next(py_rows_iter))) {
44854518
py_row_iter = PyObject_GetIter(py_row);
4486-
if (!py_row_iter) goto error;
4519+
if (!py_row_iter) {
4520+
Py_DECREF(py_row);
4521+
goto error;
4522+
}
44874523

44884524
// First item is always a row ID
44894525
py_item = PyIter_Next(py_row_ids_iter);
4490-
if (!py_item) goto error;
4526+
if (!py_item) {
4527+
Py_DECREF(py_row_iter);
4528+
Py_DECREF(py_row);
4529+
goto error;
4530+
}
44914531
row_id = (int64_t)PyLong_AsLongLong(py_item);
4532+
Py_DECREF(py_item);
44924533

44934534
CHECKMEM(8);
44944535
memcpy(out+out_idx, &row_id, 8);
@@ -4631,12 +4672,16 @@ static PyObject *dump_rowdat_1(PyObject *self, PyObject *args, PyObject *kwargs)
46314672
out_idx += 8;
46324673
} else {
46334674
PyObject *py_bytes = PyUnicode_AsEncodedString(py_item, "utf-8", "strict");
4634-
if (!py_bytes) goto error;
4675+
if (!py_bytes) {
4676+
Py_DECREF(py_item);
4677+
goto error;
4678+
}
46354679

46364680
char *str = NULL;
46374681
Py_ssize_t str_l = 0;
46384682
if (PyBytes_AsStringAndSize(py_bytes, &str, &str_l) < 0) {
46394683
Py_DECREF(py_bytes);
4684+
Py_DECREF(py_item);
46404685
goto error;
46414686
}
46424687

@@ -4671,6 +4716,7 @@ static PyObject *dump_rowdat_1(PyObject *self, PyObject *args, PyObject *kwargs)
46714716
char *str = NULL;
46724717
Py_ssize_t str_l = 0;
46734718
if (PyBytes_AsStringAndSize(py_item, &str, &str_l) < 0) {
4719+
Py_DECREF(py_item);
46744720
goto error;
46754721
}
46764722

@@ -4684,6 +4730,7 @@ static PyObject *dump_rowdat_1(PyObject *self, PyObject *args, PyObject *kwargs)
46844730
break;
46854731

46864732
default:
4733+
Py_DECREF(py_item);
46874734
goto error;
46884735
}
46894736

@@ -4693,14 +4740,17 @@ static PyObject *dump_rowdat_1(PyObject *self, PyObject *args, PyObject *kwargs)
46934740
i++;
46944741
}
46954742

4743+
Py_DECREF(py_row_iter);
46964744
Py_DECREF(py_row);
4745+
py_row_iter = NULL;
46974746
py_row = NULL;
46984747
}
46994748

4700-
py_out = PyMemoryView_FromMemory(out, out_idx, PyBUF_WRITE);
4701-
if (!py_out) goto error;
4749+
// Convert the output buffer to a Python bytes object and free the buffer
4750+
py_out = PyBytes_FromStringAndSize(out, out_idx);
47024751

47034752
exit:
4753+
if (out) free(out);
47044754
if (returns) free(returns);
47054755

47064756
Py_XDECREF(py_item);
@@ -4712,7 +4762,6 @@ static PyObject *dump_rowdat_1(PyObject *self, PyObject *args, PyObject *kwargs)
47124762
return py_out;
47134763

47144764
error:
4715-
if (!py_out && out) free(out);
47164765
Py_XDECREF(py_out);
47174766
py_out = NULL;
47184767

@@ -4839,7 +4888,7 @@ PyMODINIT_FUNC PyInit__singlestoredb_accel(void) {
48394888

48404889
PyObj.create_numpy_array_kwargs = PyDict_New();
48414890
if (!PyObj.create_numpy_array_kwargs) goto error;
4842-
if (PyDict_SetItemString(PyObj.create_numpy_array_kwargs, "copy", Py_False)) {
4891+
if (PyDict_SetItemString(PyObj.create_numpy_array_kwargs, "copy", Py_True)) {
48434892
goto error;
48444893
}
48454894

0 commit comments

Comments
 (0)