Skip to content

Commit ca25471

Browse files
ngoldbaumngoldbaum
committed
Refactor so malloc/free are accessed via an allocator
1 parent ec7058f commit ca25471

File tree

6 files changed

+165
-91
lines changed

6 files changed

+165
-91
lines changed

‎stringdtype/stringdtype/src/casts.c

Lines changed: 45 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -72,13 +72,12 @@ string_to_string(PyArrayMethod_Context *context, char *const data[],
7272
npy_intp const dimensions[], npy_intp const strides[],
7373
NpyAuxData *NPY_UNUSED(auxdata))
7474
{
75-
StringDTypeObject *in_descr =
76-
((StringDTypeObject *)context->descriptors[0]);
77-
StringDTypeObject *out_descr =
78-
((StringDTypeObject *)context->descriptors[1]);
79-
int in_hasnull = in_descr->na_object != NULL;
80-
int out_hasnull = out_descr->na_object != NULL;
81-
const npy_static_string *in_na_name = &in_descr->na_name;
75+
StringDTypeObject *idescr = (StringDTypeObject *)context->descriptors[0];
76+
StringDTypeObject *odescr = (StringDTypeObject *)context->descriptors[1];
77+
npy_string_allocator *allocator = odescr->allocator;
78+
int in_hasnull = idescr->na_object != NULL;
79+
int out_hasnull = odescr->na_object != NULL;
80+
const npy_static_string *in_na_name = &idescr->na_name;
8281
npy_intp N = dimensions[0];
8382
char *in = data[0];
8483
char *out = data[1];
@@ -89,17 +88,17 @@ string_to_string(PyArrayMethod_Context *context, char *const data[],
8988
const npy_packed_static_string *s = (npy_packed_static_string *)in;
9089
npy_packed_static_string *os = (npy_packed_static_string *)out;
9190
if (in != out) {
92-
npy_string_free(os);
91+
npy_string_free(os, allocator);
9392
if (in_hasnull && !out_hasnull && npy_string_isnull(s)) {
9493
// lossy but this is an unsafe cast so this is OK
95-
if (npy_string_newsize(in_na_name->buf, in_na_name->size, os) <
96-
0) {
94+
if (npy_string_newsize(
95+
in_na_name->buf, in_na_name->size, os, allocator) < 0) {
9796
gil_error(PyExc_MemoryError,
9897
"Failed to allocate string in string to string "
9998
"cast.");
10099
}
101100
}
102-
else if (npy_string_dup(s, os) < 0) {
101+
else if (npy_string_dup(s, os, allocator) < 0) {
103102
gil_error(PyExc_MemoryError, "npy_string_dup failed");
104103
return -1;
105104
}
@@ -211,6 +210,8 @@ unicode_to_string(PyArrayMethod_Context *context, char *const data[],
211210
NpyAuxData *NPY_UNUSED(auxdata))
212211
{
213212
PyArray_Descr **descrs = context->descriptors;
213+
npy_string_allocator *allocator =
214+
((StringDTypeObject *)descrs[1])->allocator;
214215
long max_in_size = (descrs[0]->elsize) / 4;
215216

216217
npy_intp N = dimensions[0];
@@ -230,8 +231,8 @@ unicode_to_string(PyArrayMethod_Context *context, char *const data[],
230231
return -1;
231232
}
232233
npy_packed_static_string *out_pss = (npy_packed_static_string *)out;
233-
npy_string_free(out_pss);
234-
if (npy_string_newemptysize(out_num_bytes, out_pss) < 0) {
234+
npy_string_free(out_pss, allocator);
235+
if (npy_string_newemptysize(out_num_bytes, out_pss, allocator) < 0) {
235236
gil_error(PyExc_MemoryError,
236237
"Failed to allocate string in unicode to string cast");
237238
return -1;
@@ -493,7 +494,7 @@ static char *s2b_name = "cast_StringDType_to_Bool";
493494
// bool to string
494495

495496
static int
496-
bool_to_string(PyArrayMethod_Context *NPY_UNUSED(context), char *const data[],
497+
bool_to_string(PyArrayMethod_Context *context, char *const data[],
497498
npy_intp const dimensions[], npy_intp const strides[],
498499
NpyAuxData *NPY_UNUSED(auxdata))
499500
{
@@ -504,9 +505,12 @@ bool_to_string(PyArrayMethod_Context *NPY_UNUSED(context), char *const data[],
504505
npy_intp in_stride = strides[0];
505506
npy_intp out_stride = strides[1];
506507

508+
StringDTypeObject *descr = (StringDTypeObject *)context->descriptors[1];
509+
npy_string_allocator *allocator = descr->allocator;
510+
507511
while (N--) {
508512
npy_packed_static_string *out_pss = (npy_packed_static_string *)out;
509-
npy_string_free(out_pss);
513+
npy_string_free(out_pss, allocator);
510514
char *ret_val = NULL;
511515
size_t size = 0;
512516
if ((npy_bool)(*in) == 1) {
@@ -522,7 +526,7 @@ bool_to_string(PyArrayMethod_Context *NPY_UNUSED(context), char *const data[],
522526
"invalid value encountered in bool to string cast");
523527
return -1;
524528
}
525-
if (npy_string_newsize(ret_val, size, out_pss) < 0) {
529+
if (npy_string_newsize(ret_val, size, out_pss, allocator) < 0) {
526530
// execution should never get here because this will be a small
527531
// string on all platforms
528532
gil_error(PyExc_MemoryError,
@@ -605,7 +609,7 @@ string_to_int(char *in, npy_longlong *value, int hasnull,
605609
}
606610

607611
static int
608-
pyobj_to_string(PyObject *obj, char *out)
612+
pyobj_to_string(PyObject *obj, char *out, npy_string_allocator *allocator)
609613
{
610614
if (obj == NULL) {
611615
return -1;
@@ -622,8 +626,8 @@ pyobj_to_string(PyObject *obj, char *out)
622626
return -1;
623627
}
624628
npy_packed_static_string *out_ss = (npy_packed_static_string *)out;
625-
npy_string_free(out_ss);
626-
if (npy_string_newsize(cstr_val, length, out_ss) < 0) {
629+
npy_string_free(out_ss, allocator);
630+
if (npy_string_newsize(cstr_val, length, out_ss, allocator) < 0) {
627631
PyErr_SetString(PyExc_MemoryError,
628632
"Failed to allocate numpy string when converting from "
629633
"python string.");
@@ -636,17 +640,18 @@ pyobj_to_string(PyObject *obj, char *out)
636640
}
637641

638642
static int
639-
int_to_string(long long in, char *out)
643+
int_to_string(long long in, char *out, npy_string_allocator *allocator)
640644
{
641645
PyObject *pylong_val = PyLong_FromLongLong(in);
642-
return pyobj_to_string(pylong_val, out);
646+
return pyobj_to_string(pylong_val, out, allocator);
643647
}
644648

645649
static int
646-
uint_to_string(unsigned long long in, char *out)
650+
uint_to_string(unsigned long long in, char *out,
651+
npy_string_allocator *allocator)
647652
{
648653
PyObject *pylong_val = PyLong_FromUnsignedLongLong(in);
649-
return pyobj_to_string(pylong_val, out);
654+
return pyobj_to_string(pylong_val, out, allocator);
650655
}
651656

652657
#define STRING_INT_CASTS(typename, typekind, shortname, numpy_tag, \
@@ -720,7 +725,7 @@ uint_to_string(unsigned long long in, char *out)
720725
static char *s2##shortname##_name = "cast_StringDType_to_" #typename; \
721726
\
722727
static int typename##_to_string( \
723-
PyArrayMethod_Context *NPY_UNUSED(context), char *const data[], \
728+
PyArrayMethod_Context *context, char *const data[], \
724729
npy_intp const dimensions[], npy_intp const strides[], \
725730
NpyAuxData *NPY_UNUSED(auxdata)) \
726731
{ \
@@ -731,8 +736,12 @@ uint_to_string(unsigned long long in, char *out)
731736
npy_intp in_stride = strides[0] / sizeof(npy_##typename); \
732737
npy_intp out_stride = strides[1]; \
733738
\
739+
StringDTypeObject *descr = \
740+
(StringDTypeObject *)context->descriptors[1]; \
741+
npy_string_allocator *allocator = descr->allocator; \
742+
\
734743
while (N--) { \
735-
if (typekind##_to_string((longtype)*in, out) != 0) { \
744+
if (typekind##_to_string((longtype)*in, out, allocator) != 0) { \
736745
return -1; \
737746
} \
738747
\
@@ -917,9 +926,13 @@ string_to_pyfloat(char *in, int hasnull,
917926
npy_intp in_stride = strides[0] / sizeof(npy_##typename); \
918927
npy_intp out_stride = strides[1]; \
919928
\
929+
StringDTypeObject *descr = \
930+
(StringDTypeObject *)context->descriptors[1]; \
931+
npy_string_allocator *allocator = descr->allocator; \
932+
\
920933
while (N--) { \
921934
PyObject *scalar_val = PyArray_Scalar(in, float_descr, NULL); \
922-
if (pyobj_to_string(scalar_val, out) == -1) { \
935+
if (pyobj_to_string(scalar_val, out, allocator) == -1) { \
923936
return -1; \
924937
} \
925938
\
@@ -1096,9 +1109,12 @@ datetime_to_string(PyArrayMethod_Context *context, char *const data[],
10961109
// buffer passed to numpy to build datetime string
10971110
char datetime_buf[NPY_DATETIME_MAX_ISO8601_STRLEN];
10981111

1112+
StringDTypeObject *sdescr = (StringDTypeObject *)context->descriptors[1];
1113+
npy_string_allocator *allocator = sdescr->allocator;
1114+
10991115
while (N--) {
11001116
npy_packed_static_string *out_pss = (npy_packed_static_string *)out;
1101-
npy_string_free(out_pss);
1117+
npy_string_free(out_pss, allocator);
11021118
if (*in == NPY_DATETIME_NAT) {
11031119
*out_pss = *NPY_NULL_STRING;
11041120
}
@@ -1117,8 +1133,8 @@ datetime_to_string(PyArrayMethod_Context *context, char *const data[],
11171133
return -1;
11181134
}
11191135

1120-
if (npy_string_newsize(datetime_buf, strlen(datetime_buf),
1121-
out_pss) < 0) {
1136+
if (npy_string_newsize(datetime_buf, strlen(datetime_buf), out_pss,
1137+
allocator) < 0) {
11221138
PyErr_SetString(PyExc_MemoryError,
11231139
"Failed to allocate string when converting "
11241140
"from a datetime.");

‎stringdtype/stringdtype/src/dtype.c

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,9 @@ new_stringdtype_instance(PyObject *na_object, int coerce)
1818
return NULL;
1919
}
2020

21+
npy_string_allocator *allocator =
22+
npy_string_new_allocator(PyMem_RawMalloc, PyMem_RawFree);
23+
2124
Py_XINCREF(na_object);
2225
((StringDTypeObject *)new)->na_object = na_object;
2326
npy_packed_static_string packed_na_name = *NPY_EMPTY_STRING;
@@ -31,8 +34,8 @@ new_stringdtype_instance(PyObject *na_object, int coerce)
3134
has_string_na = 1;
3235
Py_ssize_t size = 0;
3336
const char *buf = PyUnicode_AsUTF8AndSize(na_object, &size);
34-
if (npy_string_newsize(buf, (size_t)size, &packed_default_string) <
35-
0) {
37+
if (npy_string_newsize(buf, (size_t)size, &packed_default_string,
38+
allocator) < 0) {
3639
PyErr_NoMemory();
3740
Py_DECREF(new);
3841
return NULL;
@@ -65,7 +68,8 @@ new_stringdtype_instance(PyObject *na_object, int coerce)
6568

6669
Py_ssize_t size = 0;
6770
const char *utf8_ptr = PyUnicode_AsUTF8AndSize(na_pystr, &size);
68-
if (npy_string_newsize(utf8_ptr, (size_t)size, &packed_na_name)) {
71+
if (npy_string_newsize(utf8_ptr, (size_t)size, &packed_na_name,
72+
allocator)) {
6973
PyErr_NoMemory();
7074
Py_DECREF(new);
7175
Py_DECREF(na_pystr);
@@ -81,6 +85,7 @@ new_stringdtype_instance(PyObject *na_object, int coerce)
8185
snew->packed_default_string = packed_default_string;
8286
snew->packed_na_name = packed_na_name;
8387
snew->coerce = coerce;
88+
snew->allocator = allocator;
8489

8590
npy_static_string default_string = {0, NULL};
8691
npy_load_string(&snew->packed_default_string, &default_string);
@@ -197,7 +202,7 @@ stringdtype_setitem(StringDTypeObject *descr, PyObject *obj, char **dataptr)
197202

198203
// free if dataptr holds preexisting string data,
199204
// npy_string_free does a NULL check and checks for small strings
200-
npy_string_free(sdata);
205+
npy_string_free(sdata, descr->allocator);
201206

202207
// borrow reference
203208
PyObject *na_object = descr->na_object;
@@ -221,7 +226,7 @@ stringdtype_setitem(StringDTypeObject *descr, PyObject *obj, char **dataptr)
221226
return -1;
222227
}
223228

224-
if (npy_string_newsize(val, length, sdata) < 0) {
229+
if (npy_string_newsize(val, length, sdata, descr->allocator) < 0) {
225230
PyErr_NoMemory();
226231
Py_DECREF(val_obj);
227232
return -1;
@@ -375,13 +380,14 @@ stringdtype_ensure_canonical(StringDTypeObject *self)
375380

376381
static int
377382
stringdtype_clear_loop(void *NPY_UNUSED(traverse_context),
378-
PyArray_Descr *NPY_UNUSED(descr), char *data,
379-
npy_intp size, npy_intp stride,
380-
NpyAuxData *NPY_UNUSED(auxdata))
383+
PyArray_Descr *descr, char *data, npy_intp size,
384+
npy_intp stride, NpyAuxData *NPY_UNUSED(auxdata))
381385
{
386+
StringDTypeObject *sdescr = (StringDTypeObject *)descr;
382387
while (size--) {
388+
npy_packed_static_string *sdata = (npy_packed_static_string *)data;
383389
if (data != NULL) {
384-
npy_string_free((npy_packed_static_string *)data);
390+
npy_string_free(sdata, sdescr->allocator);
385391
memset(data, 0, sizeof(npy_packed_static_string));
386392
}
387393
data += stride;
@@ -560,8 +566,9 @@ static void
560566
stringdtype_dealloc(StringDTypeObject *self)
561567
{
562568
Py_XDECREF(self->na_object);
563-
npy_string_free(&self->packed_default_string);
564-
npy_string_free(&self->packed_na_name);
569+
npy_string_free(&self->packed_default_string, self->allocator);
570+
npy_string_free(&self->packed_na_name, self->allocator);
571+
npy_string_free_allocator(self->allocator);
565572
PyArrayDescr_Type.tp_dealloc((PyObject *)self);
566573
}
567574

‎stringdtype/stringdtype/src/dtype.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ typedef struct {
2929
npy_packed_static_string packed_default_string;
3030
npy_static_string na_name;
3131
npy_packed_static_string packed_na_name;
32+
npy_string_allocator *allocator;
3233
} StringDTypeObject;
3334

3435
typedef struct {

0 commit comments

Comments
 (0)