Skip to content

F
flatbuffers
Commit ad8b1e5d authored by Felix Frank's avatar Felix Frank Committed by Robert
Browse files
Options

[Python] Fast serialization of numpy vectors (#4829) 

[Python] Fast serialization of numpy vectors (#4829)
parent 76d31e1b
Hide whitespace changes
Inline Side-by-side
Showing with 297 additions and 1 deletion
python/flatbuffers/builder.py
View file @ ad8b1e5d
......@@ -21,8 +21,9 @@ from . import packer
from . import compat
from .compat import range_func
from .compat import memoryview_type
from .compat import import_numpy, NumpyRequiredForThisFeature
np = import_numpy()
## @file
## @addtogroup flatbuffers_python_api
## @{
......@@ -441,6 +442,41 @@ class Builder(object):
return self.EndVector(len(x))
def CreateNumpyVector(self, x):
"""CreateNumpyVector writes a numpy array into the buffer."""
if np is None:
# Numpy is required for this feature
raise NumpyRequiredForThisFeature("Numpy was not found.")
if not isinstance(x, np.ndarray):
raise TypeError("non-numpy-ndarray passed to CreateNumpyVector")
if x.dtype.kind not in ['b', 'i', 'u', 'f']:
raise TypeError("numpy-ndarray holds elements of unsupported datatype")
if x.ndim > 1:
raise TypeError("multidimensional-ndarray passed to CreateNumpyVector")
self.StartVector(x.itemsize, x.size, x.dtype.alignment)
# Ensure little endian byte ordering
if x.dtype.str[0] == "<":
x_lend = x
else:
x_lend = x.byteswap(inplace=False)
# Calculate total length
l = UOffsetTFlags.py_type(x_lend.itemsize * x_lend.size)
## @cond FLATBUFFERS_INTERNAL
self.head = UOffsetTFlags.py_type(self.Head() - l)
## @endcond
# tobytes ensures c_contiguous ordering
self.Bytes[self.Head():self.Head()+l] = x_lend.tobytes(order='C')
return self.EndVector(x.size)
## @cond FLATBUFFERS_INTERNAL
def assertNested(self):
"""
......
tests/py_test.py
View file @ ad8b1e5d
......@@ -478,6 +478,266 @@ class TestByteLayout(unittest.TestCase):
# 1-byte pad:
self.assertBuilderEquals(b, [3, 0, 0, 0, 1, 2, 3, 0])
def test_create_numpy_vector_int8(self):
try:
imp.find_module('numpy')
# if numpy exists, then we should be able to get the
# vector as a numpy array
import numpy as np
# Systems endian:
b = flatbuffers.Builder(0)
x = np.array([1, 2, -3], dtype=np.int8)
b.CreateNumpyVector(x)
self.assertBuilderEquals(b, [
3, 0, 0, 0, # vector length
1, 2, 256 - 3, 0 # vector value + padding
])
# Reverse endian:
b = flatbuffers.Builder(0)
x_other_endian = x.byteswap().newbyteorder()
b.CreateNumpyVector(x_other_endian)
self.assertBuilderEquals(b, [
3, 0, 0, 0, # vector length
1, 2, 256 - 3, 0 # vector value + padding
])
except ImportError:
b = flatbuffers.Builder(0)
x = 0
assertRaises(
self,
lambda: b.CreateNumpyVector(x),
NumpyRequiredForThisFeature)
def test_create_numpy_vector_uint16(self):
try:
imp.find_module('numpy')
# if numpy exists, then we should be able to get the
# vector as a numpy array
import numpy as np
# Systems endian:
b = flatbuffers.Builder(0)
x = np.array([1, 2, 312], dtype=np.uint16)
b.CreateNumpyVector(x)
self.assertBuilderEquals(b, [
3, 0, 0, 0, # vector length
1, 0, # 1
2, 0, # 2
312 - 256, 1, # 312
0, 0 # padding
])
# Reverse endian:
b = flatbuffers.Builder(0)
x_other_endian = x.byteswap().newbyteorder()
b.CreateNumpyVector(x_other_endian)
self.assertBuilderEquals(b, [
3, 0, 0, 0, # vector length
1, 0, # 1
2, 0, # 2
312 - 256, 1, # 312
0, 0 # padding
])
except ImportError:
b = flatbuffers.Builder(0)
x = 0
assertRaises(
self,
lambda: b.CreateNumpyVector(x),
NumpyRequiredForThisFeature)
def test_create_numpy_vector_int64(self):
try:
imp.find_module('numpy')
# if numpy exists, then we should be able to get the
# vector as a numpy array
import numpy as np
# Systems endian:
b = flatbuffers.Builder(0)
x = np.array([1, 2, -12], dtype=np.int64)
b.CreateNumpyVector(x)
self.assertBuilderEquals(b, [
3, 0, 0, 0, # vector length
1, 0, 0, 0, 0, 0, 0, 0, # 1
2, 0, 0, 0, 0, 0, 0, 0, # 2
256 - 12, 255, 255, 255, 255, 255, 255, 255 # -12
])
# Reverse endian:
b = flatbuffers.Builder(0)
x_other_endian = x.byteswap().newbyteorder()
b.CreateNumpyVector(x_other_endian)
self.assertBuilderEquals(b, [
3, 0, 0, 0, # vector length
1, 0, 0, 0, 0, 0, 0, 0, # 1
2, 0, 0, 0, 0, 0, 0, 0, # 2
256 - 12, 255, 255, 255, 255, 255, 255, 255 # -12
])
except ImportError:
b = flatbuffers.Builder(0)
x = 0
assertRaises(
self,
lambda: b.CreateNumpyVector(x),
NumpyRequiredForThisFeature)
def test_create_numpy_vector_float32(self):
try:
imp.find_module('numpy')
# if numpy exists, then we should be able to get the
# vector as a numpy array
import numpy as np
# Systems endian:
b = flatbuffers.Builder(0)
x = np.array([1, 2, -12], dtype=np.float32)
b.CreateNumpyVector(x)
self.assertBuilderEquals(b, [
3, 0, 0, 0, # vector length
0, 0, 128, 63, # 1
0, 0, 0, 64, # 2
0, 0, 64, 193 # -12
])
# Reverse endian:
b = flatbuffers.Builder(0)
x_other_endian = x.byteswap().newbyteorder()
b.CreateNumpyVector(x_other_endian)
self.assertBuilderEquals(b, [
3, 0, 0, 0, # vector length
0, 0, 128, 63, # 1
0, 0, 0, 64, # 2
0, 0, 64, 193 # -12
])
except ImportError:
b = flatbuffers.Builder(0)
x = 0
assertRaises(
self,
lambda: b.CreateNumpyVector(x),
NumpyRequiredForThisFeature)
def test_create_numpy_vector_float64(self):
try:
imp.find_module('numpy')
# if numpy exists, then we should be able to get the
# vector as a numpy array
import numpy as np
# Systems endian:
b = flatbuffers.Builder(0)
x = np.array([1, 2, -12], dtype=np.float64)
b.CreateNumpyVector(x)
self.assertBuilderEquals(b, [
3, 0, 0, 0, # vector length
0, 0, 0, 0, 0, 0, 240, 63, # 1
0, 0, 0, 0, 0, 0, 0, 64, # 2
0, 0, 0, 0, 0, 0, 40, 192 # -12
])
# Reverse endian:
b = flatbuffers.Builder(0)
x_other_endian = x.byteswap().newbyteorder()
b.CreateNumpyVector(x_other_endian)
self.assertBuilderEquals(b, [
3, 0, 0, 0, # vector length
0, 0, 0, 0, 0, 0, 240, 63, # 1
0, 0, 0, 0, 0, 0, 0, 64, # 2
0, 0, 0, 0, 0, 0, 40, 192 # -12
])
except ImportError:
b = flatbuffers.Builder(0)
x = 0
assertRaises(
self,
lambda: b.CreateNumpyVector(x),
NumpyRequiredForThisFeature)
def test_create_numpy_vector_bool(self):
try:
imp.find_module('numpy')
# if numpy exists, then we should be able to get the
# vector as a numpy array
import numpy as np
# Systems endian:
b = flatbuffers.Builder(0)
x = np.array([True, False, True], dtype=np.bool)
b.CreateNumpyVector(x)
self.assertBuilderEquals(b, [
3, 0, 0, 0, # vector length
1, 0, 1, 0 # vector values + padding
])
# Reverse endian:
b = flatbuffers.Builder(0)
x_other_endian = x.byteswap().newbyteorder()
b.CreateNumpyVector(x_other_endian)
self.assertBuilderEquals(b, [
3, 0, 0, 0, # vector length
1, 0, 1, 0 # vector values + padding
])
except ImportError:
b = flatbuffers.Builder(0)
x = 0
assertRaises(
self,
lambda: b.CreateNumpyVector(x),
NumpyRequiredForThisFeature)
def test_create_numpy_vector_reject_strings(self):
try:
imp.find_module('numpy')
# if numpy exists, then we should be able to get the
# vector as a numpy array
import numpy as np
# Create String array
b = flatbuffers.Builder(0)
x = np.array(["hello", "fb", "testing"])
assertRaises(
self,
lambda: b.CreateNumpyVector(x),
TypeError)
except ImportError:
b = flatbuffers.Builder(0)
x = 0
assertRaises(
self,
lambda: b.CreateNumpyVector(x),
NumpyRequiredForThisFeature)
def test_create_numpy_vector_reject_object(self):
try:
imp.find_module('numpy')
# if numpy exists, then we should be able to get the
# vector as a numpy array
import numpy as np
# Create String array
b = flatbuffers.Builder(0)
x = np.array([{"m": 0}, {"as": -2.1, 'c': 'c'}])
assertRaises(
self,
lambda: b.CreateNumpyVector(x),
TypeError)
except ImportError:
b = flatbuffers.Builder(0)
x = 0
assertRaises(
self,
lambda: b.CreateNumpyVector(x),
NumpyRequiredForThisFeature)
def test_empty_vtable(self):
b = flatbuffers.Builder(0)
b.StartObject(0)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment