From 0d94c9e092f24140f57b62f969c7c6f6f38a6be7 Mon Sep 17 00:00:00 2001 From: maniarathi Date: Wed, 5 Aug 2020 11:56:34 -0700 Subject: [PATCH] DRY-ing flatbuffer code (#1716) --- server/data_common/fbs/matrix.py | 124 +++++++++++-------------------- server/requirements.txt | 2 +- server/test/unit/decode_fbs.py | 43 ++--------- 3 files changed, 53 insertions(+), 116 deletions(-) diff --git a/server/data_common/fbs/matrix.py b/server/data_common/fbs/matrix.py index e10bc56a..213e3ce4 100644 --- a/server/data_common/fbs/matrix.py +++ b/server/data_common/fbs/matrix.py @@ -1,58 +1,24 @@ -import flatbuffers -import numpy as np -from scipy import sparse -import pandas as pd import json +import numpy as np +import pandas as pd +from flatbuffers import Builder +from scipy import sparse + import server.data_common.fbs.NetEncoding.Column as Column -import server.data_common.fbs.NetEncoding.TypedArray as TypedArray -import server.data_common.fbs.NetEncoding.Matrix as Matrix -import server.data_common.fbs.NetEncoding.Int32Array as Int32Array -import server.data_common.fbs.NetEncoding.Uint32Array as Uint32Array import server.data_common.fbs.NetEncoding.Float32Array as Float32Array import server.data_common.fbs.NetEncoding.Float64Array as Float64Array +import server.data_common.fbs.NetEncoding.Int32Array as Int32Array import server.data_common.fbs.NetEncoding.JSONEncodedArray as JSONEncodedArray - - -# Placeholder until recent enhancements to flatbuffers Python -# runtime are released, at which point we can use the default -# version. This code is a port of the head. See: -# -# https://github.com/google/flatbuffers/pull/4829 -# -def CreateNumpyVector(builder, x): - """CreateNumpyVector writes a numpy array into the buffer.""" - - if not isinstance(x, np.ndarray): - raise TypeError(f"non-numpy-ndarray passed to CreateNumpyVector ({type(x)}") - - 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") - - builder.StartVector(x.itemsize, x.size, x.dtype.alignment) - - # Ensure little endian byte ordering - if x.dtype.str[0] == "<": - x_little_endian = x - else: - x_little_endian = x.byteswap(inplace=False) - - # Calculate total length - length = int(x_little_endian.itemsize * x_little_endian.size) - builder.head = int(builder.Head() - length) - - # tobytes ensures c_contiguous ordering - builder.Bytes[builder.Head() : builder.Head() + length] = x_little_endian.tobytes(order="C") - - return builder.EndVector(x.size) +import server.data_common.fbs.NetEncoding.Matrix as Matrix +import server.data_common.fbs.NetEncoding.TypedArray as TypedArray +import server.data_common.fbs.NetEncoding.Uint32Array as Uint32Array # Serialization helper def serialize_column(builder, typed_arr): """ Serialize NetEncoding.Column """ + (u_type, u_value) = typed_arr Column.ColumnStart(builder) Column.ColumnAddUType(builder, u_type) @@ -63,6 +29,7 @@ def serialize_column(builder, typed_arr): # Serialization helper def serialize_matrix(builder, n_rows, n_cols, columns, col_idx): """ Serialize NetEncoding.Matrix """ + Matrix.MatrixStart(builder) Matrix.MatrixAddNRows(builder, n_rows) Matrix.MatrixAddNCols(builder, n_cols) @@ -77,9 +44,10 @@ def serialize_matrix(builder, n_rows, n_cols, columns, col_idx): # Serialization helper def serialize_typed_array(builder, source_array, encoding_info): """ - Serialize any of the various typed arrays, eg, Float32Array. Specific - means of serialization and type conversion are provided by type_info. + Serialize any of the various typed arrays, eg, Float32Array. Specific means of serialization and type conversion + are provided by type_info. """ + arr = source_array (array_type, as_type) = encoding_info(source_array) @@ -104,7 +72,8 @@ def serialize_typed_array(builder, source_array, encoding_info): arr = arr[0] elif arr.shape[1] == 1: arr = arr.T[0] - vec = CreateNumpyVector(builder, arr) + + vec = builder.CreateNumpyVector(arr) # serialize the typed array table builder.StartObject(1) @@ -113,38 +82,36 @@ def serialize_typed_array(builder, source_array, encoding_info): return (array_type, array_value) -column_encoding_type_map = { - # array protocol string: ( array_type, as_type ) - np.dtype(np.float64).str: (TypedArray.TypedArray.Float32Array, np.float32), - np.dtype(np.float32).str: (TypedArray.TypedArray.Float32Array, np.float32), - np.dtype(np.float16).str: (TypedArray.TypedArray.Float32Array, np.float32), - np.dtype(np.int8).str: (TypedArray.TypedArray.Int32Array, np.int32), - np.dtype(np.int16).str: (TypedArray.TypedArray.Int32Array, np.int32), - np.dtype(np.int32).str: (TypedArray.TypedArray.Int32Array, np.int32), - np.dtype(np.int64).str: (TypedArray.TypedArray.Int32Array, np.int32), - np.dtype(np.uint8).str: (TypedArray.TypedArray.Uint32Array, np.uint32), - np.dtype(np.uint16).str: (TypedArray.TypedArray.Uint32Array, np.uint32), - np.dtype(np.uint32).str: (TypedArray.TypedArray.Uint32Array, np.uint32), - np.dtype(np.uint64).str: (TypedArray.TypedArray.Uint32Array, np.uint32), -} -column_encoding_default = (TypedArray.TypedArray.JSONEncodedArray, "json") - - def column_encoding(arr): + column_encoding_type_map = { + # array protocol string: ( array_type, as_type ) + np.dtype(np.float64).str: (TypedArray.TypedArray.Float32Array, np.float32), + np.dtype(np.float32).str: (TypedArray.TypedArray.Float32Array, np.float32), + np.dtype(np.float16).str: (TypedArray.TypedArray.Float32Array, np.float32), + np.dtype(np.int8).str: (TypedArray.TypedArray.Int32Array, np.int32), + np.dtype(np.int16).str: (TypedArray.TypedArray.Int32Array, np.int32), + np.dtype(np.int32).str: (TypedArray.TypedArray.Int32Array, np.int32), + np.dtype(np.int64).str: (TypedArray.TypedArray.Int32Array, np.int32), + np.dtype(np.uint8).str: (TypedArray.TypedArray.Uint32Array, np.uint32), + np.dtype(np.uint16).str: (TypedArray.TypedArray.Uint32Array, np.uint32), + np.dtype(np.uint32).str: (TypedArray.TypedArray.Uint32Array, np.uint32), + np.dtype(np.uint64).str: (TypedArray.TypedArray.Uint32Array, np.uint32), + } + column_encoding_default = (TypedArray.TypedArray.JSONEncodedArray, "json") + return column_encoding_type_map.get(arr.dtype.str, column_encoding_default) -index_encoding_type_map = { - # array protocol string: ( array_type, as_type ) - np.dtype(np.int32).str: (TypedArray.TypedArray.Int32Array, np.int32), - np.dtype(np.int64).str: (TypedArray.TypedArray.Int32Array, np.int32), - np.dtype(np.uint32).str: (TypedArray.TypedArray.Uint32Array, np.uint32), - np.dtype(np.uint64).str: (TypedArray.TypedArray.Uint32Array, np.uint32), -} -index_encoding_default = (TypedArray.TypedArray.JSONEncodedArray, "json") - - def index_encoding(arr): + index_encoding_type_map = { + # array protocol string: ( array_type, as_type ) + np.dtype(np.int32).str: (TypedArray.TypedArray.Int32Array, np.int32), + np.dtype(np.int64).str: (TypedArray.TypedArray.Int32Array, np.int32), + np.dtype(np.uint32).str: (TypedArray.TypedArray.Uint32Array, np.uint32), + np.dtype(np.uint64).str: (TypedArray.TypedArray.Uint32Array, np.uint32), + } + index_encoding_default = (TypedArray.TypedArray.JSONEncodedArray, "json") + return index_encoding_type_map.get(arr.dtype.str, index_encoding_default) @@ -165,8 +132,7 @@ def guess_at_mem_needed(matrix): def encode_matrix_fbs(matrix, row_idx=None, col_idx=None): """ - Given a 2D DataFrame, ndarray or sparse equivalent, create and return a - Matrix flatbuffer. + Given a 2D DataFrame, ndarray or sparse equivalent, create and return a Matrix flatbuffer. :param matrix: 2D DataFrame, ndarray or sparse equivalent :param row_idx: index for row dimension, Index or ndarray @@ -183,7 +149,7 @@ def encode_matrix_fbs(matrix, row_idx=None, col_idx=None): (n_rows, n_cols) = matrix.shape # estimate size needed, so we don't unnecessarily realloc. - builder = flatbuffers.Builder(guess_at_mem_needed(matrix)) + builder = Builder(guess_at_mem_needed(matrix)) columns = [] for cidx in range(n_cols - 1, -1, -1): @@ -239,9 +205,9 @@ def deserialize_typed_array(tarr): def decode_matrix_fbs(fbs): """ - Given an FBS-encoded Matrix, return a Pandas DataFrame the contains the data - and indices. + Given an FBS-encoded Matrix, return a Pandas DataFrame the contains the data and indices. """ + matrix = Matrix.Matrix.GetRootAsMatrix(fbs, 0) n_rows = matrix.NRows() n_cols = matrix.NCols() diff --git a/server/requirements.txt b/server/requirements.txt index a4948f7d..cbac485c 100644 --- a/server/requirements.txt +++ b/server/requirements.txt @@ -8,7 +8,7 @@ Flask-Cors>=3.0.6 Flask-RESTful>=0.3.6 flask-server-timing>=0.1.2 flask-talisman>=0.7.0 -flatbuffers>=1.10.0 +flatbuffers>=1.11.0 flatten-dict>=0.2.0 fsspec>=0.4.4,<0.8.0 numba>=0.49.1 diff --git a/server/test/unit/decode_fbs.py b/server/test/unit/decode_fbs.py index 477ce189..6fb1fcec 100644 --- a/server/test/unit/decode_fbs.py +++ b/server/test/unit/decode_fbs.py @@ -1,46 +1,17 @@ """ Code to decode, for testing purposes, the flatbuffer encoded blobs. -This code will need to be updated if fbs/matrix.fbs changes. -For more information, see fbs/matrix.fbs and server/data_common/fbs/ + +This code will need to be updated if fbs/matrix.fbs changes. For more information, see fbs/matrix.fbs and +server/data_common/fbs/ """ -import json -import server.data_common.fbs.NetEncoding.Float32Array as Float32Array -import server.data_common.fbs.NetEncoding.Float64Array as Float64Array -import server.data_common.fbs.NetEncoding.Int32Array as Int32Array -import server.data_common.fbs.NetEncoding.JSONEncodedArray as JSONEncodedArray import server.data_common.fbs.NetEncoding.Matrix as Matrix -import server.data_common.fbs.NetEncoding.TypedArray as TypedArray -import server.data_common.fbs.NetEncoding.Uint32Array as Uint32Array - - -def decode_typed_array(tarr): - type_map = { - TypedArray.TypedArray.Uint32Array: Uint32Array.Uint32Array, - TypedArray.TypedArray.Int32Array: Int32Array.Int32Array, - TypedArray.TypedArray.Float32Array: Float32Array.Float32Array, - TypedArray.TypedArray.Float64Array: Float64Array.Float64Array, - TypedArray.TypedArray.JSONEncodedArray: JSONEncodedArray.JSONEncodedArray, - } - (u_type, u) = tarr - if u_type == TypedArray.TypedArray.NONE: - return None - - TarType = type_map.get(u_type, None) - assert TarType is not None - - arr = TarType() - arr.Init(u.Bytes, u.Pos) - narr = arr.DataAsNumpy() - if u_type == TypedArray.TypedArray.JSONEncodedArray: - narr = json.loads(narr.tostring().decode("utf-8")) - return narr +from server.data_common.fbs.matrix import deserialize_typed_array def decode_matrix_FBS(buf): """ - Given a FBS Matrix, return an decoded Python dict containing - same info in native format. + Given a FBS Matrix, return an decoded Python dict containing same info in native format. NOTE / TODO: row_idx not currently implemented """ df = Matrix.Matrix.GetRootAsMatrix(buf, 0) @@ -53,8 +24,8 @@ def decode_matrix_FBS(buf): for col_idx in range(0, columns_length): col = df.Columns(col_idx) tarr = (col.UType(), col.U()) - decoded_columns.append(decode_typed_array(tarr)) + decoded_columns.append(deserialize_typed_array(tarr)) - cidx = decode_typed_array((df.ColIndexType(), df.ColIndex())) + cidx = deserialize_typed_array((df.ColIndexType(), df.ColIndex())) return {"n_rows": n_rows, "n_cols": n_cols, "columns": decoded_columns, "col_idx": cidx, "row_idx": None}