implement small string optimization

Author: ngoldbaum

stringdtype/stringdtype/src/casts.c

@@ -1057,8 +1057,7 @@ datetime_to_string(PyArrayMethod_Context *context, char *const data[],
         npy_static_string *out_ss = (npy_static_string *)out;
         npy_string_free(out_ss);
         if (*in == NPY_DATETIME_NAT) {
-            /* convert to NA */
-            out_ss = NULL;
+            *out_ss = NPY_NULL_STRING;
         }
         else {
             if (NpyDatetime_ConvertDatetime64ToDatetimeStruct(dt_meta, *in,

stringdtype/stringdtype/src/dtype.c

@@ -20,18 +20,18 @@ new_stringdtype_instance(PyObject *na_object, int coerce)
 
     Py_XINCREF(na_object);
     ((StringDTypeObject *)new)->na_object = na_object;
-    npy_static_string na_name = NPY_NULL_STRING;
+    npy_static_string na_name = NPY_EMPTY_STRING;
+    npy_static_string default_string = NPY_EMPTY_STRING;
     int hasnull = na_object != NULL;
     int has_nan_na = 0;
     int has_string_na = 0;
-    npy_static_string default_string = NPY_EMPTY_STRING;
     if (hasnull) {
         // first check for a string
         if (PyUnicode_Check(na_object)) {
             has_string_na = 1;
             Py_ssize_t size = 0;
             const char *buf = PyUnicode_AsUTF8AndSize(na_object, &size);
-            default_string = NPY_NULL_STRING;
+            default_string = NPY_EMPTY_STRING;
             int res = npy_string_newsize(buf, (size_t)size, &default_string);
             if (res == -1) {
                 PyErr_NoMemory();
@@ -206,8 +206,7 @@ stringdtype_setitem(StringDTypeObject *descr, PyObject *obj, char **dataptr)
     // setting NA *must* check pointer equality since NA types might not
     // allow equality
     if (na_object != NULL && obj == na_object) {
-        // do nothing, npy_string_free already NULLed the struct ssdata points
-        // to so it already contains a NA value
+        *sdata = NPY_NULL_STRING;
     }
     else {
         PyObject *val_obj = get_value(obj, descr->coerce);
@@ -259,7 +258,7 @@ stringdtype_getitem(StringDTypeObject *descr, char **dataptr)
         }
     }
     else {
-        char *data = npy_string_buf(sdata);
+        const char *data = npy_string_buf(sdata);
         size_t size = npy_string_size(sdata);
         val_obj = PyUnicode_FromStringAndSize(data, size);
         if (val_obj == NULL) {
@@ -420,9 +419,7 @@ stringdtype_fill_zero_loop(void *NPY_UNUSED(traverse_context),
                            NpyAuxData *NPY_UNUSED(auxdata))
 {
     while (size--) {
-        if (npy_string_newsize("", 0, (npy_static_string *)(data)) < 0) {
-            return -1;
-        }
+        *(npy_static_string *)(data) = NPY_EMPTY_STRING;
         data += stride;
     }
     return 0;

stringdtype/stringdtype/src/main.c

@@ -58,7 +58,10 @@ _memory_usage(PyObject *NPY_UNUSED(self), PyObject *obj)
         npy_intp count = *innersizeptr;
 
         while (count--) {
-            memory_usage += npy_string_size(((npy_static_string *)in));
+            size_t size = npy_string_size(((npy_static_string *)in));
+            if (size > NPY_SHORT_STRING_MAX_SIZE) {
+                memory_usage += size;
+            }
             in += stride;
         }
 

stringdtype/stringdtype/src/static_string.c

@@ -2,17 +2,43 @@
 
 #include "static_string.h"
 
-// defined this way so NPY_EMPTY_STRING has an in-memory representation that is
-// distinct from a zero-filled struct, allowing us to use a NPY_NULL_STRING
-// to represent a sentinel value
-const npy_static_string NPY_EMPTY_STRING = {0, "\0"};
-const npy_static_string NPY_NULL_STRING = {0, NULL};
+// Since this has no flags set, technically this is a heap-allocated string
+// with size zero practically, that doesn't matter because we always do size
+// checks before accessing heap data, but that may be confusing. The nice part
+// of this choice is a calloc'd array buffer (e.g. from np.empty) is filled
+// with empty elements for free
+const npy_static_string NPY_EMPTY_STRING = {
+        .base = {.direct_buffer = {.flags_and_size = 0, .buf = {0}}}};
+// zero-filled, but with the NULL flag set to distinguish from empty string
+const npy_static_string NPY_NULL_STRING = {
+        .base = {.direct_buffer = {.flags_and_size = NPY_STRING_MISSING,
+                                   .buf = {0}}}};
+
+int
+is_short_string(const npy_static_string *s)
+{
+    unsigned char high_byte = s->base.direct_buffer.flags_and_size;
+    return (high_byte & NPY_STRING_SHORT) == NPY_STRING_SHORT;
+}
+
+int
+npy_string_isnull(const npy_static_string *s)
+{
+    unsigned char high_byte = s->base.direct_buffer.flags_and_size;
+    return (high_byte & NPY_STRING_MISSING) == NPY_STRING_MISSING;
+}
+
+int
+is_not_a_vstring(const npy_static_string *s)
+{
+    return is_short_string(s) || npy_string_isnull(s);
+}
 
 int
 npy_string_newsize(const char *init, size_t size, npy_static_string *to_init)
 {
-    if ((to_init == NULL) || (to_init->buf != NULL) ||
-        (npy_string_size(to_init) != 0)) {
+    if (to_init == NULL || npy_string_size(to_init) != 0 ||
+        size > MAX_STRING_SIZE) {
         return -2;
     }
 
@@ -21,114 +47,128 @@ npy_string_newsize(const char *init, size_t size, npy_static_string *to_init)
         return 0;
     }
 
-    char *ret_buf = (char *)PyMem_RawMalloc(sizeof(char) * size);
-
-    if (ret_buf == NULL) {
-        return -1;
-    }
-
-    to_init->size = size;
+    if (size > NPY_SHORT_STRING_MAX_SIZE) {
+        char *ret_buf = (char *)PyMem_RawMalloc(sizeof(char) * size);
 
-    memcpy(ret_buf, init, size);
+        if (ret_buf == NULL) {
+            return -1;
+        }
 
-    to_init->buf = ret_buf;
+        to_init->base.vstring.size = size;
 
-    return 0;
-}
+        memcpy(ret_buf, init, size);
 
-void
-npy_string_free(npy_static_string *str)
-{
-    if (str->buf != NULL && str->buf != NPY_EMPTY_STRING.buf) {
-        PyMem_RawFree(str->buf);
-        str->buf = NULL;
-    }
-    str->size = 0;
-}
-
-int
-npy_string_dup(const npy_static_string *in, npy_static_string *out)
-{
-    if (npy_string_isnull(in)) {
-        out->size = 0;
-        out->buf = NULL;
-        return 0;
+        to_init->base.vstring.buf = ret_buf;
     }
     else {
-        return npy_string_newsize(in->buf, in->size, out);
+        // size can be no longer than 7 or 15, depending on CPU architecture
+        // in either case, the size data is in at most the least significant 4
+        // bits of the byte so it's safe to | with one of 0x10, 0x20, 0x40, or
+        // 0x80.
+        to_init->base.direct_buffer.flags_and_size = NPY_STRING_SHORT | size;
+        memcpy(&(to_init->base.direct_buffer.buf), init, size);
     }
+
+    return 0;
 }
 
 int
 npy_string_newemptysize(size_t size, npy_static_string *out)
 {
-    if (out->size != 0 || out->buf != NULL) {
+    if (out == NULL || npy_string_size(out) != 0 || size > MAX_STRING_SIZE) {
         return -2;
     }
 
-    out->size = size;
-
     if (size == 0) {
         *out = NPY_EMPTY_STRING;
         return 0;
     }
 
-    char *buf = (char *)PyMem_RawMalloc(sizeof(char) * size);
+    if (size > NPY_SHORT_STRING_MAX_SIZE) {
+        char *buf = (char *)PyMem_RawMalloc(sizeof(char) * size);
 
-    if (buf == NULL) {
-        return -1;
-    }
+        if (buf == NULL) {
+            return -1;
+        }
 
-    out->buf = buf;
+        out->base.vstring.buf = buf;
+        out->base.vstring.size = size;
+    }
+    else {
+        out->base.direct_buffer.flags_and_size = NPY_STRING_SHORT | size;
+    }
 
     return 0;
 }
 
+void
+npy_string_free(npy_static_string *str)
+{
+    if (is_not_a_vstring(str)) {
+        // zero out
+        memcpy(str, &NPY_EMPTY_STRING, sizeof(npy_static_string));
+    }
+    else {
+        if (str->base.vstring.size != 0) {
+            PyMem_RawFree(str->base.vstring.buf);
+        }
+        str->base.vstring.buf = NULL;
+        str->base.vstring.size = 0;
+    }
+}
+
+int
+npy_string_dup(const npy_static_string *in, npy_static_string *out)
+{
+    if (npy_string_isnull(in)) {
+        *out = NPY_NULL_STRING;
+        return 0;
+    }
+
+    return npy_string_newsize(npy_string_buf(in), npy_string_size(in), out);
+}
+
 int
 npy_string_cmp(const npy_static_string *s1, const npy_static_string *s2)
 {
-    size_t minsize = s1->size < s2->size ? s1->size : s2->size;
+    size_t s1_size = npy_string_size(s1);
+    size_t s2_size = npy_string_size(s2);
+
+    char *s1_buf = npy_string_buf(s1);
+    char *s2_buf = npy_string_buf(s2);
 
-    int cmp = strncmp(s1->buf, s2->buf, minsize);
+    size_t minsize = s1_size < s2_size ? s1_size : s2_size;
+
+    int cmp = strncmp(s1_buf, s2_buf, minsize);
 
     if (cmp == 0) {
-        if (s1->size > minsize) {
+        if (s1_size > minsize) {
             return 1;
         }
-        if (s2->size > minsize) {
+        if (s2_size > minsize) {
             return -1;
         }
     }
 
     return cmp;
 }
 
-int
-npy_string_isnull(const npy_static_string *in)
-{
-    if (in->size == 0 && in->buf == NULL) {
-        return 1;
-    }
-    return 0;
-}
-
 size_t
 npy_string_size(const npy_static_string *s)
 {
-    return s->size;
+    if (is_short_string(s)) {
+        unsigned char high_byte = s->base.direct_buffer.flags_and_size;
+        return high_byte & NPY_SHORT_STRING_SIZE_MASK;
+    }
+    return s->base.vstring.size;
 }
 
 char *
 npy_string_buf(const npy_static_string *s)
 {
-    return s->buf;
-}
-
-int
-npy_string_size_and_buf(const npy_static_string *s, size_t *size, char **buf)
-{
-    *size = s->size;
-    *buf = s->buf;
-
-    return 0;
+    if (is_short_string(s)) {
+        // the cast drops const, is there a better way?
+        return (char *)&s->base.direct_buffer.buf[0];
+    }
+    return s->base.vstring.buf;
 }

stringdtype/stringdtype/src/static_string.h

@@ -4,11 +4,59 @@
 #include "stdlib.h"
 #include "string.h"
 
-typedef struct npy_static_string {
+#if NPY_BYTE_ORDER == NPY_LITTLE_ENDIAN
+
+// the high byte in vstring.size is resolved for flags
+// SSSS SSSF
+
+typedef struct _npy_static_string_t {
+    char *buf;
+    size_t size;
+} _npy_static_string_t;
+
+typedef struct _short_string_buffer {
+    char buf[sizeof(_npy_static_string_t) - 1];
+    unsigned char flags_and_size;
+} _short_string_buffer;
+
+#elif NPY_BYTE_ORDER == NPY_BIG_ENDIAN
+
+// the high byte in vstring.size is resolved for flags
+// FSSS SSSS
+
+typedef struct _npy_static_string_t {
     size_t size;
     char *buf;
+} _npy_static_string_t;
+
+typedef struct _short_string_buffer {
+    unsigned char flags_and_size;
+    char buf[sizeof(npy_static_string_t) - 1];
+} _short_string_buffer;
+
+#endif
+
+typedef union _npy_static_string_u {
+    _npy_static_string_t vstring;
+    _short_string_buffer direct_buffer;
+} _npy_static_string_u;
+
+typedef struct npy_static_string {
+    _npy_static_string_u base;
 } npy_static_string;
 
+// room for two more flags with values 0x20 and 0x10
+#define NPY_STRING_MISSING 0x80  // 1000 0000
+#define NPY_STRING_SHORT 0x40    // 0100 0000
+
+// short string sizes fit in a 4-bit integer
+#define NPY_SHORT_STRING_SIZE_MASK 0x0F  // 0000 1111
+#define NPY_SHORT_STRING_MAX_SIZE \
+    (sizeof(npy_static_string) - 1)  // 15 or 7 depending on arch
+
+// one byte in size is reserved for flags and small string optimization
+#define MAX_STRING_SIZE (1 << (sizeof(size_t) - 1)) - 1
+
 // represents the empty string and can be passed safely to npy_static_string
 // API functions
 extern const npy_static_string NPY_EMPTY_STRING;
@@ -65,9 +113,4 @@ npy_string_size(const npy_static_string *s);
 char *
 npy_string_buf(const npy_static_string *s);
 
-// Fills in *size* and *buf* pointers with the values in *s*.
-// Currently always returns 0.
-int
-npy_string_size_and_buf(const npy_static_string *s, size_t *size, char **buf);
-
 #endif /*_NPY_STATIC_STRING_H */