Fast Numpy/Protobuf deserialization example
In my last post I wrote about how I was able to improve protobuf deserialization using a C++ extension. I’ve boiled down the code to its essence to show how I did it. Rather than zip everything up in a file, the code is short enough to show in its entirety.
Here’s the simplified protobuf message which is used to represent a time series as 2 arrays:
package timeseries;
message TimeSeries {
repeated double times = 2;
repeated double values = 3;
}
I then wrote a test app in Python and C++ to provide a benchmark. Here is the Python version:
import numpy
import time_series_pb2
def write_test():
ts = time_series_pb2.TimeSeries()
for i in range(10000000):
ts.times.append(i)
ts.values.append(i*10.0)
import time
start = time.time();
f = open("ts.bin", "wb")
f.write(ts.SerializeToString())
f.close()
print time.time() - start
def read_test():
ts = time_series_pb2.TimeSeries()
import time
start = time.time();
f = open("ts.bin", "rb")
ts.ParseFromString(f.read())
f.close()
t = numpy.array(ts.times._values)
v = numpy.array(ts.values._values)
print 'Read time:', time.time() - start
print "Read time series of length %d" % len(ts.times)
if __name__ == "__main__":
import sys
if len(sys.argv) < 2:
print "usage: %s <--read> | <--write>" % sys.argv[0]
if sys.argv[1] == "--read":
read_test()
else:
write_test()
I will spare you the C++ standalone code, since it was only a stepping stone. Instead here is the C++ extension, with 2 exposed methods, one which deserializes a string and the other which operates on a file.
#include <fcntl.h>
#include <Python.h>
#include <numpy/arrayobject.h>
#include <google/protobuf/io/coded_stream.h>
#include <google/protobuf/io/zero_copy_stream_impl_lite.h>
#include <google/protobuf/io/zero_copy_stream_impl.h>
#include "time_series.pb.h"
static PyObject* construct_numpy_arrays(timeseries::TimeSeries* ts)
{
// returns a tuple (t,v) where t and v are double arrays of the same length
PyObject* data_tuple = PyTuple_New(2);
long array_size = ts->times_size();
double* times = new double[array_size];
double* values = new double[array_size];
// the data must be copied because the tsid will go away and its mutable data
// will too
memcpy(times, ts->times().data(), ts->times_size()*sizeof(double));
memcpy(values, ts->values().data(), ts->values_size()*sizeof(double));
// put the arrays into numpy array objects
npy_intp dims[1] = {array_size};
PyObject* time_array = PyArray_SimpleNewFromData(1, dims, PyArray_DOUBLE, times);
PyObject* value_array = PyArray_SimpleNewFromData(1, dims, PyArray_DOUBLE, values);
PyTuple_SetItem(data_tuple, 0, time_array);
PyTuple_SetItem(data_tuple, 1, value_array);
return data_tuple;
}
static PyObject* TimeSeries_load(PyObject* self, PyObject* args)
{
char* filename = NULL;
if (! PyArg_ParseTuple(args, "s", &filename))
{
return NULL;
}
timeseries::TimeSeries ts;
int fd = open(filename, O_RDONLY);
google::protobuf::io::FileInputStream fs(fd);
google::protobuf::io::CodedInputStream coded_fs(&fs);
coded_fs.SetTotalBytesLimit(500*1024*1024, -1);
ts.ParseFromCodedStream(&coded_fs);
fs.Close();
close(fd);
return construct_numpy_arrays(&ts);
}
static PyObject* TimeSeries_deserialize(PyObject* self, PyObject* args)
{
int buffer_length;
char* serialization = NULL;
if (! PyArg_ParseTuple(args, "t#", &serialization, &buffer_length))
{
return NULL;
}
google::protobuf::io::ArrayInputStream input(serialization, buffer_length);
google::protobuf::io::CodedInputStream coded_fs(&input);
coded_fs.SetTotalBytesLimit(500*1024*1024, -1);
timeseries::TimeSeries ts;
ts.ParseFromCodedStream(&coded_fs);
return construct_numpy_arrays(&ts);
}
static PyMethodDef TSMethods[] = {
{"load", TimeSeries_load, METH_VARARGS, "loads a TimeSeries from a file"},
{"deserialize", TimeSeries_deserialize, METH_VARARGS, "loads a TimeSeries from a string"}
};
#ifndef PyMODINIT_FUNC /* declarations for DLL import/export */
#define PyMODINIT_FUNC void
#endif
PyMODINIT_FUNC inittimeseries(void)
{
import_array();
(void) Py_InitModule("timeseries", TSMethods);
}
Calling the exension from python is trivial:
import time
import timeseries
start = time.time()
t, v = timeseries.load('ts.bin')
print "read and converted to numpy array in %f" % (time.time()-start)
print "timeseries contained %d values" % len(v)
Posted by Bryce Hendrix @ 2:17 pm under Uncategorized | Comments (1)
