#include <Python.h>

#include <bytesobject.h>

#include <string.h>

#define NO_IMPORT_ARRAY

#include "numpyapi.h"

#include "common.h"

#include "markers.h"

#include "encoder.h"

#include "python_funcs.h"

static char bytes_array_prefix[] = {ARRAY_START, CONTAINER_TYPE, TYPE_BYTE, CONTAINER_COUNT};

#define POWER_TWO(x) ((long long) 1 << (x))

#if defined(_MSC_VER) && !defined(fpclassify)

#define USE__FPCLASS

#define BUFFER_INITIAL_SIZE 64

#define BUFFER_FP_SIZE 256

static PyObject* EncoderException = NULL;

static PyTypeObject* PyDec_Type = NULL;

#define PyDec_Check(v) PyObject_TypeCheck(v, PyDec_Type)

static int _encoder_buffer_write(_bjdata_encoder_buffer_t* buffer, const char* const chunk, size_t chunk_len);

#define RECURSE_AND_BAIL_ON_NONZERO(action, recurse_msg) {\

int ret;\

BAIL_ON_NONZERO(Py_EnterRecursiveCall(recurse_msg));\

ret = (action);\

Py_LeaveRecursiveCall();\

BAIL_ON_NONZERO(ret);\

}

#define WRITE_OR_BAIL(str, len) BAIL_ON_NONZERO(_encoder_buffer_write(buffer, (str), len))

#define WRITE_CHAR_OR_BAIL(c) {\

char ctmp = (c);\

WRITE_OR_BAIL(&ctmp, 1);\

}

static int _encode_PyBytes(PyObject* obj, _bjdata_encoder_buffer_t* buffer);

static int _encode_PyObject_as_PyDecimal(PyObject* obj, _bjdata_encoder_buffer_t* buffer);

static int _encode_PyDecimal(PyObject* obj, _bjdata_encoder_buffer_t* buffer);

static int _encode_PyUnicode(PyObject* obj, _bjdata_encoder_buffer_t* buffer);

static int _encode_PyFloat(PyObject* obj, _bjdata_encoder_buffer_t* buffer);

static int _encode_PyLong(PyObject* obj, _bjdata_encoder_buffer_t* buffer);

static int _encode_longlong(long long num, _bjdata_encoder_buffer_t* buffer);

#if PY_MAJOR_VERSION < 3

static int _encode_PyInt(PyObject* obj, _bjdata_encoder_buffer_t* buffer);

static int _encode_PySequence(PyObject* obj, _bjdata_encoder_buffer_t* buffer);

static int _encode_mapping_key(PyObject* obj, _bjdata_encoder_buffer_t* buffer);

static int _encode_PyMapping(PyObject* obj, _bjdata_encoder_buffer_t* buffer);

static int _encode_NDarray(PyObject* obj, _bjdata_encoder_buffer_t* buffer);

static int _encode_soa(PyArrayObject* arr, _bjdata_encoder_buffer_t* buffer, int is_row_major);

static int _write_field_schema_recursive(PyArray_Descr* fd, _bjdata_encoder_buffer_t* buffer);

static int _write_index_value(Py_ssize_t idx, int size, _bjdata_encoder_buffer_t* buffer);

static int _analyze_string_field(PyArrayObject* flat, PyObject* field_name,

Py_ssize_t field_index, npy_intp count,

double threshold, _string_field_info_t* info);

static const int numpytypes[][2] = {

{NPY_BOOL, TYPE_UINT8},

{NPY_BYTE, TYPE_INT8},

{NPY_INT8, TYPE_INT8},

{NPY_SHORT, TYPE_INT16},

{NPY_INT16, TYPE_INT16},

{NPY_INT, TYPE_INT32},

{NPY_INT32, TYPE_INT32},

{NPY_LONGLONG, TYPE_INT64},

{NPY_INT64, TYPE_INT64},

{NPY_UINT8, TYPE_UINT8},

{NPY_UBYTE, TYPE_UINT8},

{NPY_USHORT, TYPE_UINT16},

{NPY_UINT16, TYPE_UINT16},

{NPY_UINT, TYPE_UINT32},

{NPY_UINT32, TYPE_UINT32},

{NPY_ULONGLONG, TYPE_UINT64},

{NPY_UINT64, TYPE_UINT64},

{NPY_HALF, TYPE_FLOAT16},

{NPY_FLOAT16, TYPE_FLOAT16},

{NPY_FLOAT, TYPE_FLOAT32},

{NPY_FLOAT32, TYPE_FLOAT32},

{NPY_DOUBLE, TYPE_FLOAT64},

{NPY_FLOAT64, TYPE_FLOAT64},

{NPY_CFLOAT, TYPE_FLOAT32},

{NPY_COMPLEX64, TYPE_FLOAT32},

{NPY_CDOUBLE, TYPE_FLOAT64},

{NPY_COMPLEX128, TYPE_FLOAT64},

{NPY_STRING, TYPE_STRING},

{NPY_UNICODE, TYPE_STRING}

};

_bjdata_encoder_buffer_t* _bjdata_encoder_buffer_create(_bjdata_encoder_prefs_t* prefs, PyObject* fp_write) {

_bjdata_encoder_buffer_t* buffer;

if (NULL == (buffer = calloc(1, sizeof(_bjdata_encoder_buffer_t)))) {

PyErr_NoMemory();

return NULL;

}

buffer->len = (NULL != fp_write) ? BUFFER_FP_SIZE : BUFFER_INITIAL_SIZE;

BAIL_ON_NULL(buffer->obj = PyBytes_FromStringAndSize(NULL, buffer->len));

buffer->raw = PyBytes_AS_STRING(buffer->obj);

buffer->pos = 0;

BAIL_ON_NULL(buffer->markers = PySet_New(NULL));

buffer->prefs = *prefs;

buffer->fp_write = fp_write;

Py_XINCREF(fp_write);

// treat Py_None as no default_func being supplied

if (Py_None == buffer->prefs.default_func) {

buffer->prefs.default_func = NULL;

}

return buffer;

bail:

_bjdata_encoder_buffer_free(&buffer);

return NULL;

}

void _bjdata_encoder_buffer_free(_bjdata_encoder_buffer_t** buffer) {

if (NULL != buffer && NULL != *buffer) {

Py_XDECREF((*buffer)->obj);

Py_XDECREF((*buffer)->fp_write);

Py_XDECREF((*buffer)->markers);

free(*buffer);

*buffer = NULL;

}

}

static int _encoder_buffer_write(_bjdata_encoder_buffer_t* buffer, const char* const chunk, size_t chunk_len) {

size_t new_len;

PyObject* fp_write_ret;

if (0 == chunk_len) {

return 0;

}

// no write method, use buffer only

if (NULL == buffer->fp_write) {

// increase buffer size if too small

if (chunk_len > (buffer->len - buffer->pos)) {

for (new_len = buffer->len; new_len < (buffer->pos + chunk_len); new_len *= 2);

BAIL_ON_NONZERO(_PyBytes_Resize(&buffer->obj, new_len));

buffer->raw = PyBytes_AS_STRING(buffer->obj);

buffer->len = new_len;

}

memcpy(&(buffer->raw[buffer->pos]), chunk, sizeof(char) * chunk_len);

buffer->pos += chunk_len;

} else {

// increase buffer to fit all first

if (chunk_len > (buffer->len - buffer->pos)) {

BAIL_ON_NONZERO(_PyBytes_Resize(&buffer->obj, (buffer->pos + chunk_len)));

buffer->raw = PyBytes_AS_STRING(buffer->obj);

buffer->len = buffer->pos + chunk_len;

}

memcpy(&(buffer->raw[buffer->pos]), chunk, sizeof(char) * chunk_len);

buffer->pos += chunk_len;

// flush buffer to write method

if (buffer->pos >= buffer->len) {

BAIL_ON_NULL(fp_write_ret = PyObject_CallFunctionObjArgs(buffer->fp_write, buffer->obj, NULL));

Py_DECREF(fp_write_ret);

Py_DECREF(buffer->obj);

buffer->len = BUFFER_FP_SIZE;

BAIL_ON_NULL(buffer->obj = PyBytes_FromStringAndSize(NULL, buffer->len));

buffer->raw = PyBytes_AS_STRING(buffer->obj);

buffer->pos = 0;

}

}

return 0;

bail:

return 1;

}

PyObject* _bjdata_encoder_buffer_finalise(_bjdata_encoder_buffer_t* buffer) {

PyObject* fp_write_ret;

// shrink buffer to fit

if (buffer->pos < buffer->len) {

BAIL_ON_NONZERO(_PyBytes_Resize(&buffer->obj, buffer->pos));

buffer->len = buffer->pos;

}

if (NULL == buffer->fp_write) {

Py_INCREF(buffer->obj);

return buffer->obj;

} else {

if (buffer->pos > 0) {

BAIL_ON_NULL(fp_write_ret = PyObject_CallFunctionObjArgs(buffer->fp_write, buffer->obj, NULL));

Py_DECREF(fp_write_ret);

}

Py_RETURN_NONE;

}

bail:

return NULL;

}

static int _encode_PyBytes(PyObject* obj, _bjdata_encoder_buffer_t* buffer) {

const char* raw;

Py_ssize_t len;

raw = PyBytes_AS_STRING(obj);

len = PyBytes_GET_SIZE(obj);

WRITE_OR_BAIL(bytes_array_prefix, sizeof(bytes_array_prefix));

BAIL_ON_NONZERO(_encode_longlong(len, buffer));

WRITE_OR_BAIL(raw, len);

// no ARRAY_END since length was specified

return 0;

bail:

return 1;

}

static int _encode_PyByteArray(PyObject* obj, _bjdata_encoder_buffer_t* buffer) {

const char* raw;

Py_ssize_t len;

raw = PyByteArray_AS_STRING(obj);

len = PyByteArray_GET_SIZE(obj);

WRITE_OR_BAIL(bytes_array_prefix, sizeof(bytes_array_prefix));

BAIL_ON_NONZERO(_encode_longlong(len, buffer));

WRITE_OR_BAIL(raw, len);

// no ARRAY_END since length was specified

return 0;

bail:

return 1;

}

static int _lookup_marker(npy_intp numpytypeid) {

int i, len = (sizeof(numpytypes) >> 3);

for (i = 0; i < len; i++) {

if (numpytypeid == (npy_intp)numpytypes[i][0]) {

return numpytypes[i][1];

}

}

return -1;

}

static int _get_soa_type_marker(int dtype_num) {

/* Boolean needs special handling in SOA */

if (dtype_num == NPY_BOOL) {

return TYPE_BOOL_TRUE;

}

/* For other types, use the unified lookup */

int marker = _lookup_marker(dtype_num);

/* Filter out string types - they need special handling in SOA */

if (marker == TYPE_STRING) {

return -1;

}

return marker;

}

static inline int _is_string_type(int type_num) {

return (type_num == NPY_UNICODE || type_num == NPY_STRING);

}

static int _is_soa_compatible_dtype(PyArray_Descr* fd) {

int type_num = DESCR_TYPE_NUM(fd);

/* String types are supported */

if (_is_string_type(type_num)) {

return 1;

}

/* Numeric types (including bool) */

if (_get_soa_type_marker(type_num) >= 0) {

return 1;

}

if (type_num == NPY_VOID) {

/* Check for sub-array first */

PyObject* subdtype = PyObject_GetAttrString((PyObject*)fd, "subdtype");

if (subdtype && subdtype != Py_None && PyTuple_Check(subdtype) && PyTuple_GET_SIZE(subdtype) >= 2) {

PyArray_Descr* base = (PyArray_Descr*)PyTuple_GET_ITEM(subdtype, 0);

int result = _is_soa_compatible_dtype(base);

Py_DECREF(subdtype);

return result;

}

Py_XDECREF(subdtype);

PyErr_Clear();

/* Check for nested struct */

PyObject* names = PyObject_GetAttrString((PyObject*)fd, "names");

if (names && names != Py_None && PyTuple_Check(names) && PyTuple_GET_SIZE(names) > 0) {

PyObject* fields = PyObject_GetAttrString((PyObject*)fd, "fields");

if (!fields || !PyMapping_Check(fields)) {

Py_DECREF(names);

Py_XDECREF(fields);

PyErr_Clear();

return 0;

}

int result = 1;

Py_ssize_t num_names = PyTuple_GET_SIZE(names);

for (Py_ssize_t i = 0; i < num_names && result; i++) {

PyObject* fname = PyTuple_GET_ITEM(names, i);

PyObject* info = PyObject_GetItem(fields, fname);

if (!info || !PyTuple_Check(info) || PyTuple_GET_SIZE(info) < 1) {

Py_XDECREF(info);

result = 0;

break;

}

PyArray_Descr* nested_fd = (PyArray_Descr*)PyTuple_GET_ITEM(info, 0);

result = _is_soa_compatible_dtype(nested_fd);

Py_DECREF(info);

}

Py_DECREF(fields);

Py_DECREF(names);

return result;

}

Py_XDECREF(names);

PyErr_Clear();

}

return 0;

}

static int _can_encode_as_soa(PyArrayObject* arr) {

PyArray_Descr* dtype = PyArray_DESCR(arr);

PyObject* names = PyObject_GetAttrString((PyObject*)dtype, "names");

if (!names || names == Py_None || !PyTuple_Check(names)) {

Py_XDECREF(names);

PyErr_Clear();

return 0;

}

Py_ssize_t num_names = PyTuple_GET_SIZE(names);

if (num_names == 0) {

Py_DECREF(names);

return 0;

}

PyObject* fields = PyObject_GetAttrString((PyObject*)dtype, "fields");

if (!fields || !PyMapping_Check(fields)) {

Py_DECREF(names);

Py_XDECREF(fields);

PyErr_Clear();

return 0;

}

int result = 1;

for (Py_ssize_t i = 0; i < num_names && result; i++) {

PyObject* fname = PyTuple_GET_ITEM(names, i);

PyObject* info = PyObject_GetItem(fields, fname);

if (!info || !PyTuple_Check(info) || PyTuple_GET_SIZE(info) < 1) {

Py_XDECREF(info);

result = 0;

break;

}

PyArray_Descr* fd = (PyArray_Descr*)PyTuple_GET_ITEM(info, 0);

result = _is_soa_compatible_dtype(fd);

Py_DECREF(info);

}

Py_DECREF(fields);

Py_DECREF(names);

return result;

}

static int _encode_field_name(PyObject* name, _bjdata_encoder_buffer_t* buffer) {

PyObject* encoded = PyUnicode_AsEncodedString(name, "utf-8", NULL);

if (!encoded) {

return 1;

}

int ret = _encode_longlong(PyBytes_GET_SIZE(encoded), buffer);

if (ret == 0) {

WRITE_OR_BAIL(PyBytes_AS_STRING(encoded), PyBytes_GET_SIZE(encoded));

}

Py_DECREF(encoded);

return ret;

bail:

Py_DECREF(encoded);

return 1;

}

static int _write_field_schema_recursive(PyArray_Descr* fd, _bjdata_encoder_buffer_t* buffer) {

int type_num = DESCR_TYPE_NUM(fd);

Py_ssize_t i;

/* Handle NPY_VOID: could be sub-array or nested struct */

if (type_num == NPY_VOID) {

/* Check for sub-array first */

PyObject* subdtype = PyObject_GetAttrString((PyObject*)fd, "subdtype");

if (subdtype && subdtype != Py_None && PyTuple_Check(subdtype) && PyTuple_GET_SIZE(subdtype) >= 2) {

PyArray_Descr* base_dtype = (PyArray_Descr*)PyTuple_GET_ITEM(subdtype, 0);

PyObject* shape = PyTuple_GET_ITEM(subdtype, 1);

Py_ssize_t num_elem = 1;

if (PyTuple_Check(shape)) {

for (i = 0; i < PyTuple_GET_SIZE(shape); i++) {

num_elem *= PyLong_AsLong(PyTuple_GET_ITEM(shape, i));

}

}

int base_type = DESCR_TYPE_NUM(base_dtype);

Py_DECREF(subdtype);

/* Write sub-array schema: [TTT...] */

WRITE_CHAR_OR_BAIL(ARRAY_START);

if (base_type == NPY_BOOL) {

for (i = 0; i < num_elem; i++) {

WRITE_CHAR_OR_BAIL(TYPE_BOOL_TRUE);

}

} else if (_is_string_type(base_type)) {

for (i = 0; i < num_elem; i++) {

WRITE_CHAR_OR_BAIL(TYPE_CHAR);

}

} else {

int marker = _get_soa_type_marker(base_type);

if (marker < 0) {

PyErr_Format(PyExc_ValueError, "Unsupported sub-array element type: %d", base_type);

goto bail;

}

for (i = 0; i < num_elem; i++) {

WRITE_CHAR_OR_BAIL((char)marker);

}

}

WRITE_CHAR_OR_BAIL(ARRAY_END);

return 0;

}

Py_XDECREF(subdtype);

PyErr_Clear();

/* Check for nested struct */

PyObject* names = PyObject_GetAttrString((PyObject*)fd, "names");

if (names && names != Py_None && PyTuple_Check(names) && PyTuple_GET_SIZE(names) > 0) {

PyObject* fields = PyObject_GetAttrString((PyObject*)fd, "fields");

if (fields && PyMapping_Check(fields)) {

WRITE_CHAR_OR_BAIL(OBJECT_START);

Py_ssize_t nf = PyTuple_GET_SIZE(names);

for (i = 0; i < nf; i++) {

PyObject* nm = PyTuple_GET_ITEM(names, i);

BAIL_ON_NONZERO(_encode_field_name(nm, buffer));

PyObject* info = PyObject_GetItem(fields, nm);

if (!info) {

Py_DECREF(fields);

Py_DECREF(names);

goto bail;

}

PyArray_Descr* nested_fd = (PyArray_Descr*)PyTuple_GET_ITEM(info, 0);

int ret = _write_field_schema_recursive(nested_fd, buffer);

Py_DECREF(info);

if (ret != 0) {

Py_DECREF(fields);

Py_DECREF(names);

goto bail;

}

}

WRITE_CHAR_OR_BAIL(OBJECT_END);

Py_DECREF(fields);

Py_DECREF(names);

return 0;

}

Py_XDECREF(fields);

}

Py_XDECREF(names);

PyErr_Clear();

PyErr_SetString(PyExc_ValueError, "Unsupported void type in SOA schema");

goto bail;

}

/* String types - write NUMPY BYTE SIZE (not character count) */

if (_is_string_type(type_num)) {

WRITE_CHAR_OR_BAIL(TYPE_STRING);

BAIL_ON_NONZERO(_encode_longlong(DESCR_ELSIZE(fd), buffer));

return 0;

}

/* Boolean */

if (type_num == NPY_BOOL) {

WRITE_CHAR_OR_BAIL(TYPE_BOOL_TRUE);

return 0;

}

/* Numeric types */

int marker = _get_soa_type_marker(type_num);

if (marker < 0) {

PyErr_Format(PyExc_ValueError, "Unsupported SOA field type: %d", type_num);

goto bail;

}

WRITE_CHAR_OR_BAIL((char)marker);

return 0;

bail:

return 1;

}

static int _write_index_value(Py_ssize_t idx, int size, _bjdata_encoder_buffer_t* buffer) {

char buf[8];

int le = buffer->prefs.islittle;

if (size == 1) {

buf[0] = (char)idx;

WRITE_OR_BAIL(buf, 1);

} else if (size == 2) {

buf[le ? 0 : 1] = idx & 0xFF;

buf[le ? 1 : 0] = (idx >> 8) & 0xFF;

WRITE_OR_BAIL(buf, 2);

} else {

for (int i = 0; i < 4; i++) {

buf[le ? i : 3 - i] = (idx >> (8 * i)) & 0xFF;

}

WRITE_OR_BAIL(buf, 4);

}

return 0;

bail:

return 1;

}

static inline void _get_index_type(Py_ssize_t count, int* size, char* marker) {

if (count <= 255) {

*size = 1;

*marker = TYPE_UINT8;

} else if (count <= 65535) {

*size = 2;

*marker = TYPE_UINT16;

} else {

*size = 4;

*marker = TYPE_UINT32;

}

}

static int _analyze_string_field(PyArrayObject* flat, PyObject* field_name,

Py_ssize_t field_index, npy_intp count,

double threshold, _string_field_info_t* info) {

PyObject* unique_set = NULL;

PyObject* unique_list = NULL;

Py_ssize_t max_len = 0;

Py_ssize_t total_len = 0;

npy_intp j;

/* Initialize */

memset(info, 0, sizeof(_string_field_info_t));

info->encoding = SOA_STRING_FIXED;

info->fixed_len = 1;

if (count == 0) {

return 0;

}

BAIL_ON_NULL(unique_set = PySet_New(NULL));

/* First pass: collect unique values and compute lengths */

for (j = 0; j < count; j++) {

PyObject* rec = PyArray_GETITEM(flat, PyArray_GETPTR1(flat, j));

if (!rec) {

goto bail;

}

PyObject* val = NULL;

if (PyTuple_Check(rec)) {

val = PyTuple_GET_ITEM(rec, field_index);

Py_INCREF(val);

} else {

val = PyObject_GetItem(rec, field_name);

}

Py_DECREF(rec);

if (!val) {

goto bail;

}

/* Get UTF-8 length */

PyObject* utf8 = PyUnicode_AsEncodedString(val, "utf-8", NULL);

if (!utf8) {

Py_DECREF(val);

goto bail;

}

Py_ssize_t len = PyBytes_GET_SIZE(utf8);

max_len = (len > max_len) ? len : max_len;

total_len += len;

Py_DECREF(utf8);

PySet_Add(unique_set, val);

Py_DECREF(val);

}

Py_ssize_t num_unique = PySet_GET_SIZE(unique_set);

info->total_len = total_len;

info->fixed_len = max_len > 0 ? max_len : 1;

/* Force offset if threshold is exactly 0 */

if (threshold == 0.0) {

info->encoding = SOA_STRING_OFFSET;

_get_index_type(total_len, &info->index_size, &info->index_marker);

Py_DECREF(unique_set);

return 0;

}

/* Calculate costs */

double thresh = (threshold > 0) ? threshold : 0.3;

Py_ssize_t fixed_cost = max_len * count;

/* Dict cost: indices + dictionary overhead */

int idx_size;

char idx_marker;

_get_index_type(num_unique, &idx_size, &idx_marker);

/* Calculate dict strings total length */

Py_ssize_t dict_strings_total = 0;

unique_list = PySequence_List(unique_set);

if (unique_list) {

for (Py_ssize_t i = 0; i < num_unique; i++) {

PyObject* item = PyList_GET_ITEM(unique_list, i);

PyObject* enc = PyUnicode_AsEncodedString(item, "utf-8", NULL);

if (enc) {

dict_strings_total += PyBytes_GET_SIZE(enc);

Py_DECREF(enc);

}

}

}

Py_ssize_t dict_cost = idx_size * count + dict_strings_total + num_unique * 2;

/* Offset cost */

int off_size;

char off_marker;

_get_index_type(total_len, &off_size, &off_marker);

Py_ssize_t offset_cost = idx_size * count + (count + 1) * off_size + total_len;

/* Choose encoding */

if (num_unique <= (Py_ssize_t)(count * thresh) && dict_cost < fixed_cost && dict_cost < offset_cost) {

info->encoding = SOA_STRING_DICT;

info->dict_list = unique_list;

unique_list = NULL;

info->dict_count = num_unique;

info->index_size = idx_size;

info->index_marker = idx_marker;

} else if (max_len > 32 && offset_cost < fixed_cost) {

info->encoding = SOA_STRING_OFFSET;

info->index_size = off_size;

info->index_marker = off_marker;

} else {

info->encoding = SOA_STRING_FIXED;

}

Py_XDECREF(unique_list);

Py_DECREF(unique_set);

return 0;

bail:

Py_XDECREF(unique_list);

Py_XDECREF(unique_set);

return -1;

}

static int _write_string_schema(_string_field_info_t* info, _bjdata_encoder_buffer_t* buffer) {

Py_ssize_t i;

if (info->encoding == SOA_STRING_FIXED) {

WRITE_CHAR_OR_BAIL(TYPE_STRING);

BAIL_ON_NONZERO(_encode_longlong(info->fixed_len, buffer));

} else if (info->encoding == SOA_STRING_DICT) {

char hdr[] = {ARRAY_START, CONTAINER_TYPE, TYPE_STRING, CONTAINER_COUNT};

WRITE_OR_BAIL(hdr, 4);

BAIL_ON_NONZERO(_encode_longlong(info->dict_count, buffer));

for (i = 0; i < info->dict_count; i++) {

PyObject* item = PyList_GET_ITEM(info->dict_list, i);

PyObject* enc = PyUnicode_AsEncodedString(item, "utf-8", NULL);

if (!enc) {

goto bail;

}

BAIL_ON_NONZERO(_encode_longlong(PyBytes_GET_SIZE(enc), buffer));

WRITE_OR_BAIL(PyBytes_AS_STRING(enc), PyBytes_GET_SIZE(enc));

Py_DECREF(enc);

}

} else { /* SOA_STRING_OFFSET */

char hdr[] = {ARRAY_START, CONTAINER_TYPE, info->index_marker, ARRAY_END};

WRITE_OR_BAIL(hdr, 4);

}

return 0;

bail:

return 1;

}

static int _write_string_value(PyObject* str_val, _string_field_info_t* info,

Py_ssize_t record_index, _bjdata_encoder_buffer_t* buffer) {

PyObject* utf8 = PyUnicode_AsEncodedString(str_val, "utf-8", NULL);

if (!utf8) {

return 1;

}

int ret = 0;

if (info->encoding == SOA_STRING_FIXED) {

Py_ssize_t len = PyBytes_GET_SIZE(utf8);

Py_ssize_t write_len = (len < info->fixed_len) ? len : info->fixed_len;

if (_encoder_buffer_write(buffer, PyBytes_AS_STRING(utf8), write_len) != 0) {

ret = 1;

} else {

/* Pad with zeros */

for (Py_ssize_t p = len; p < info->fixed_len && ret == 0; p++) {

char zero = '\0';

if (_encoder_buffer_write(buffer, &zero, 1) != 0) {

ret = 1;

}

}

}

} else if (info->encoding == SOA_STRING_DICT) {

Py_ssize_t idx = PySequence_Index(info->dict_list, str_val);

if (idx < 0) {

PyErr_Clear();

idx = 0;

}

ret = _write_index_value(idx, info->index_size, buffer);

} else { /* SOA_STRING_OFFSET */

ret = _write_index_value(record_index, info->index_size, buffer);

}

Py_DECREF(utf8);

return ret;

}

static PyObject* _get_field_value(PyObject* rec, PyObject* field_name, Py_ssize_t field_index) {

if (PyTuple_Check(rec)) {

PyObject* val = PyTuple_GET_ITEM(rec, field_index);

Py_INCREF(val);

return val;

}

return PyObject_GetItem(rec, field_name);

}

static int _encode_soa(PyArrayObject* arr, _bjdata_encoder_buffer_t* buffer, int is_row_major) {

PyArray_Descr* dtype = PyArray_DESCR(arr);

PyObject* names = NULL, *fields = NULL;

PyArrayObject* flat = NULL;

npy_intp count, i, j, nf;

int ndim;

double threshold;

/* Per-field info */

Py_ssize_t* field_offset = NULL;

Py_ssize_t* field_itemsize = NULL;

int* field_type = NULL; /* 0=numeric, 1=bool, 2=string */

_string_field_info_t* str_info = NULL;

PyObject** str_values = NULL;

BAIL_ON_NULL(names = PyObject_GetAttrString((PyObject*)dtype, "names"));

BAIL_ON_NULL(fields = PyObject_GetAttrString((PyObject*)dtype, "fields"));

if (!PyMapping_Check(fields)) {

PyErr_SetString(PyExc_ValueError, "dtype.fields is not a mapping");

goto bail;

}

nf = PyTuple_GET_SIZE(names);

count = PyArray_SIZE(arr);

ndim = PyArray_NDIM(arr);

threshold = (buffer->prefs.soa_threshold >= 0) ? buffer->prefs.soa_threshold : 0.3;

BAIL_ON_NULL(flat = (PyArrayObject*)PyArray_Flatten(arr, NPY_CORDER));

/* Allocate per-field arrays */

BAIL_ON_NULL(field_offset = calloc(nf, sizeof(Py_ssize_t)));

BAIL_ON_NULL(field_itemsize = calloc(nf, sizeof(Py_ssize_t)));

BAIL_ON_NULL(field_type = calloc(nf, sizeof(int)));

BAIL_ON_NULL(str_info = calloc(nf, sizeof(_string_field_info_t)));

BAIL_ON_NULL(str_values = calloc(nf, sizeof(PyObject*)));

/* Analyze fields */

for (i = 0; i < nf; i++) {

PyObject* fname = PyTuple_GET_ITEM(names, i);

PyObject* info = PyObject_GetItem(fields, fname);

if (!info) {

goto bail;

}

PyArray_Descr* fd = (PyArray_Descr*)PyTuple_GET_ITEM(info, 0);

field_offset[i] = PyLong_AsSsize_t(PyTuple_GET_ITEM(info, 1));

field_itemsize[i] = DESCR_ELSIZE(fd);

int type_num = DESCR_TYPE_NUM(fd);

if (type_num == NPY_BOOL) {

field_type[i] = 1;

} else if (_is_string_type(type_num)) {

field_type[i] = 2;

if (_analyze_string_field(flat, fname, i, count, threshold, &str_info[i]) < 0) {

Py_DECREF(info);

goto bail;

}

/* Cache string values for offset encoding */

if (str_info[i].encoding == SOA_STRING_OFFSET) {

str_values[i] = PyList_New(count);

if (!str_values[i]) {

Py_DECREF(info);

goto bail;

}

for (j = 0; j < count; j++) {

PyObject* rec = PyArray_GETITEM(flat, PyArray_GETPTR1(flat, j));

if (!rec) {

Py_DECREF(info);

goto bail;

}

PyObject* val = _get_field_value(rec, fname, i);

Py_DECREF(rec);

if (!val) {

Py_DECREF(info);

goto bail;

}

PyList_SET_ITEM(str_values[i], j, val);

}

}

} else {

field_type[i] = 0;

}

Py_DECREF(info);

}

/* Write header */

WRITE_CHAR_OR_BAIL(is_row_major ? ARRAY_START : OBJECT_START);

WRITE_CHAR_OR_BAIL(CONTAINER_TYPE);

WRITE_CHAR_OR_BAIL(OBJECT_START);

/* Write schema */

for (i = 0; i < nf; i++) {

PyObject* nm = PyTuple_GET_ITEM(names, i);

BAIL_ON_NONZERO(_encode_field_name(nm, buffer));

if (field_type[i] == 2) {

BAIL_ON_NONZERO(_write_string_schema(&str_info[i], buffer));

} else {

PyObject* info = PyObject_GetItem(fields, nm);

if (!info) {

goto bail;

}

PyArray_Descr* fd = (PyArray_Descr*)PyTuple_GET_ITEM(info, 0);

int ret = _write_field_schema_recursive(fd, buffer);

Py_DECREF(info);

if (ret != 0) {

goto bail;

}

}

}

WRITE_CHAR_OR_BAIL(OBJECT_END);

/* Write count */

WRITE_CHAR_OR_BAIL(CONTAINER_COUNT);