[Python] Fast serialization of numpy vectors (#4829)
[Python] Fast serialization of numpy vectors (#4829)
This commit is contained in:
parent
76d31e1b5e
commit
ad8b1e5dbd
|
@ -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):
|
||||
"""
|
||||
|
|
260
tests/py_test.py
260
tests/py_test.py
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue