Starting the C version of rocket

Currently, only append_list is written (and hardly tested)

nilmdb/server/bulkdata.py

@@ -11,7 +11,8 @@ import os
 import cPickle as pickle
 import re
 
-from . import pyrocket as rocket
+#from . import pyrocket as rocket
+from . import rocket
 
 # Up to 256 open file descriptors at any given time.
 # These variables are global so they can be used in the decorator arguments.

nilmdb/server/rocket.c

@@ -0,0 +1,332 @@
+#include <Python.h>
+#include <structmember.h>
+#include <endian.h>
+
+typedef enum {
+	/* It's safe to treat signed as unsigned in this code */
+	LAYOUT_TYPE_NONE,
+	LAYOUT_TYPE_UINT8,
+	LAYOUT_TYPE_UINT16,
+	LAYOUT_TYPE_UINT32,
+	LAYOUT_TYPE_UINT64,
+	LAYOUT_TYPE_FLOAT32,
+	LAYOUT_TYPE_FLOAT64,
+} layout_type_t;
+
+const int layout_size[] = {
+	[LAYOUT_TYPE_NONE] = 0,
+	[LAYOUT_TYPE_UINT8] = 1,
+	[LAYOUT_TYPE_UINT16] = 2,
+	[LAYOUT_TYPE_UINT32] = 4,
+	[LAYOUT_TYPE_UINT64] = 8,
+	[LAYOUT_TYPE_FLOAT32] = 4,
+	[LAYOUT_TYPE_FLOAT64] = 8
+};
+
+struct {
+	char *string;
+	layout_type_t layout;
+	int size;
+} type_lookup[] = {
+	{ "int8",    LAYOUT_TYPE_UINT8,   1 },
+	{ "uint8",   LAYOUT_TYPE_UINT8,   1 },
+	{ "int16",   LAYOUT_TYPE_UINT16,  2 },
+	{ "uint16",  LAYOUT_TYPE_UINT16,  2 },
+	{ "int32",   LAYOUT_TYPE_UINT32,  4 },
+	{ "uint32",  LAYOUT_TYPE_UINT32,  4 },
+	{ "int64",   LAYOUT_TYPE_UINT64,  8 },
+	{ "uint64",  LAYOUT_TYPE_UINT64,  8 },
+	{ "float32", LAYOUT_TYPE_FLOAT32, 4 },
+	{ "float64", LAYOUT_TYPE_FLOAT64, 8 },
+	{ NULL }
+};
+
+typedef struct {
+	PyObject_HEAD
+	layout_type_t layout_type;
+	int layout_count;
+	int binary_size;
+} Rocket;
+
+static void Rocket_dealloc(Rocket *self)
+{
+	self->ob_type->tp_free((PyObject *)self);
+}
+
+static PyObject *Rocket_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
+{
+	Rocket *self;
+
+	self = (Rocket *)type->tp_alloc(type, 0);
+	if (!self)
+		return NULL;
+	self->layout_type = LAYOUT_TYPE_NONE;
+	self->layout_count = 0;
+	self->binary_size = 0;
+	return (PyObject *)self;
+}
+
+static int Rocket_init(Rocket *self, PyObject *args, PyObject *kwds)
+{
+	const char *layout;
+	static char *kwlist[] = { "layout", NULL };
+	if (!PyArg_ParseTupleAndKeywords(args, kwds, "s", kwlist, &layout))
+		return -1;
+	if (!layout)
+		return -1;
+
+	const char *under;
+	char *tmp;
+	under = strchr(layout, '_');
+	if (!under) {
+		PyErr_SetString(PyExc_ValueError, "no such layout: "
+				"badly formatted string");
+		return -1;
+	}
+	self->layout_count = strtoul(under+1, &tmp, 10);
+	if (self->layout_count < 1 || *tmp != '\0') {
+		PyErr_SetString(PyExc_ValueError, "no such layout: "
+				"bad count");
+		return -1;
+	}
+
+	int i;
+	for (i = 0; type_lookup[i].string; i++)
+		if (strncmp(layout, type_lookup[i].string, under-layout) == 0)
+			break;
+	if (!type_lookup[i].string) {
+		PyErr_SetString(PyExc_ValueError, "no such layout: "
+				"bad data type");
+		return -1;
+	}
+	self->layout_type = type_lookup[i].layout;
+	self->binary_size = 8 + (type_lookup[i].size * self->layout_count);
+
+	return 0;
+}
+
+FILE *PyFile_AsFile(PyObject *file)
+{
+	PyObject *result;
+	int fd;
+	result = PyObject_CallMethod(file, "fileno", NULL);
+	if (result == NULL)
+		return NULL;
+	fd = PyInt_AsLong(result);
+	if (fd < 0)
+		return NULL;
+	return fdopen(fd, "a+b");
+}
+
+/* Convert numbers, being careful with aliasing rules */
+static inline uint32_t float_to_uint32(float x)
+{
+	return *(uint32_t *)(char *)&x;
+}
+static inline uint64_t double_to_uint64(double x)
+{
+	return *(uint64_t *)(char *)&x;
+}
+
+static inline void write_binary(FILE *out, uint64_t val, layout_type_t type)
+{
+	int ret = 0;
+	uint8_t x8;
+	uint16_t x16;
+	uint32_t x32;
+	uint64_t x64;
+	switch (type) {
+	case LAYOUT_TYPE_UINT8:
+		x8 = val;
+		ret = fwrite(&x8, 1, 1, out);
+		break;
+	case LAYOUT_TYPE_UINT16:
+		x16 = htole16(val);
+		ret = fwrite(&x16, 2, 1, out);
+		break;
+	case LAYOUT_TYPE_UINT32:
+	case LAYOUT_TYPE_FLOAT32:
+		x32 = htole32(val);
+		ret = fwrite(&x32, 4, 1, out);
+		break;
+	case LAYOUT_TYPE_UINT64:
+	case LAYOUT_TYPE_FLOAT64:
+		x64 = htole64(val);
+		ret = fwrite(&x64, 8, 1, out);
+		break;
+	default:
+		PyErr_SetString(PyExc_TypeError, "unknown type");
+		return;
+	}
+	if (ret == 0)
+		PyErr_SetFromErrno(PyExc_OSError);
+}
+
+static inline void write_pyobject(FILE *out, PyObject *val, layout_type_t type)
+{
+	uint64_t u = 0;
+	double d = 0;
+
+	switch (type) {
+	case LAYOUT_TYPE_UINT8:
+	case LAYOUT_TYPE_UINT16:
+	case LAYOUT_TYPE_UINT32:
+	case LAYOUT_TYPE_UINT64:
+		u = PyInt_AsLong(val);
+		if (PyErr_Occurred())
+			return;
+		write_binary(out, u, type);
+		break;
+	case LAYOUT_TYPE_FLOAT32:
+		d = PyFloat_AsDouble(val);
+		if (PyErr_Occurred())
+			return;
+		write_binary(out, float_to_uint32(d), type);
+	case LAYOUT_TYPE_FLOAT64:
+		d = PyFloat_AsDouble(val);
+		if (PyErr_Occurred())
+			return;
+		write_binary(out, double_to_uint64(d), type);
+		break;
+	default:
+		PyErr_SetString(PyExc_TypeError, "unknown type");
+		return;
+	}
+}
+
+static PyObject *Rocket_append_list(Rocket *self, PyObject *args)
+{
+	PyObject *file, *list;
+	FILE *out;
+	if (!PyArg_ParseTuple(args, "OO:append_list", &file, &list))
+		return NULL;
+	if ((out = PyFile_AsFile(file)) == NULL)
+		return NULL;
+	if (!PyList_Check(list)) {
+		PyErr_SetString(PyExc_TypeError, "need a list");
+		return NULL;
+	}
+
+	Py_ssize_t count = PyList_Size(list);
+	Py_ssize_t row;
+	for (row = 0; row < count; row++) {
+		PyObject *rowlist = PyList_GetItem(list, row);
+		if (!PyList_Check(list)) {
+			PyErr_SetString(PyExc_TypeError, "rows must be lists");
+			return NULL;
+		}
+		if (PyList_Size(rowlist) != self->layout_count + 1) {
+			PyErr_SetString(PyExc_TypeError, "short row list");
+			return NULL;
+		}
+
+		/* Extract and write timestamp */
+		write_pyobject(out, PyList_GetItem(rowlist, 0),
+			       LAYOUT_TYPE_FLOAT64);
+		if (PyErr_Occurred())
+			return NULL;
+
+		/* Extract and write values */
+		int i;
+		for (i = 0; i < self->layout_count; i++) {
+			write_pyobject(out, PyList_GetItem(rowlist, i+1),
+				self->layout_type);
+			if (PyErr_Occurred())
+				return NULL;
+		}
+	}
+
+	/* All done */
+	Py_INCREF(Py_None);
+	return Py_None;
+}
+
+static PyObject *Rocket_extract_list(Rocket *self, PyObject *args)
+{
+	PyObject *file, *offset, *count;
+	FILE *in;
+	if (!PyArg_ParseTuple(args, "OOO:extract_list",
+			      &file, &offset, &count))
+		return NULL;
+	if ((in = PyFile_AsFile(file)) == NULL)
+		return NULL;
+
+	PyErr_SetString(PyExc_NotImplementedError, "uh oh");
+	return NULL;
+	Py_INCREF(Py_None);
+	return Py_None;
+}
+
+static PyObject *Rocket_extract_string(Rocket *self, PyObject *args)
+{
+	PyObject *file, *offset, *count;
+	FILE *in;
+	if (!PyArg_ParseTuple(args, "OOO:extract_list",
+			      &file, &offset, &count))
+		return NULL;
+	if ((in = PyFile_AsFile(file)) == NULL)
+		return NULL;
+
+	PyErr_SetString(PyExc_NotImplementedError, "uh oh");
+	return NULL;
+	Py_INCREF(Py_None);
+	return Py_None;
+}
+
+static PyMemberDef Rocket_members[] = {
+	{ "binary_size", T_INT, offsetof(Rocket, binary_size), 0,
+	  "binary size per row" },
+	{ NULL },
+};
+
+static PyMethodDef Rocket_methods[] = {
+	{ "append_list", (PyCFunction)Rocket_append_list, METH_VARARGS,
+	  "Append the list data to the file" },
+	{ "extract_list", (PyCFunction)Rocket_extract_list, METH_VARARGS,
+	  "Extract count rows of data from the file at offset offset. "
+	  "Return a list of lists [[row],[row],...]" },
+	{ "extract_string", (PyCFunction)Rocket_extract_string, METH_VARARGS,
+	  "Extract count rows of data from the file at offset offset. "
+	  "Return an ascii formatted string according to the layout" },
+	{ NULL },
+};
+
+static PyTypeObject RocketType = {
+	PyObject_HEAD_INIT(NULL)
+
+	.tp_name	= "rocket.Rocket",
+	.tp_basicsize	= sizeof(Rocket),
+	.tp_flags	= Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
+
+	.tp_new		= Rocket_new,
+	.tp_dealloc	= (destructor)Rocket_dealloc,
+	.tp_init	= (initproc)Rocket_init,
+	.tp_methods	= Rocket_methods,
+	.tp_members	= Rocket_members,
+
+	.tp_doc		= ("C implementation of the \"rocket\" data parsing "
+			   "interface, which translates between the binary "
+			   "format on disk and the ASCII or Python list "
+			   "format used when communicating with the rest of "
+			   "the system.")
+};
+
+static PyMethodDef module_methods[] = {
+	{ NULL },
+};
+
+PyMODINIT_FUNC
+initrocket(void)
+{
+	PyObject *module;
+
+	RocketType.tp_new = PyType_GenericNew;
+	if (PyType_Ready(&RocketType) < 0)
+		return;
+
+	module = Py_InitModule3("rocket", module_methods,
+				"Rocket data parsing and formatting module");
+	Py_INCREF(&RocketType);
+	PyModule_AddObject(module, "Rocket", (PyObject *)&RocketType);
+	return;
+}

setup.py

@@ -56,7 +56,7 @@ try:
 except ImportError:
     use_cython = False
 
-ext_modules = []
+ext_modules = [ Extension('nilmdb.server.rocket', ['nilmdb/server/rocket.c' ]) ]
 for modulename in cython_modules:
     filename = modulename.replace('.','/')
     if use_cython: