mirror of
https://github.com/chanzuckerberg/cellxgene.git
synced 2026-09-15 20:57:56 +08:00
252 lines
8.8 KiB
Python
252 lines
8.8 KiB
Python
import json
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
from flatbuffers import Builder
|
|
from scipy import sparse
|
|
|
|
from server.common.utils.type_conversion_utils import get_encoding_dtype_of_array
|
|
|
|
import server.common.fbs.NetEncoding.Column as Column
|
|
import server.common.fbs.NetEncoding.Float32Array as Float32Array
|
|
import server.common.fbs.NetEncoding.Float64Array as Float64Array
|
|
import server.common.fbs.NetEncoding.Int32Array as Int32Array
|
|
import server.common.fbs.NetEncoding.JSONEncodedArray as JSONEncodedArray
|
|
import server.common.fbs.NetEncoding.Matrix as Matrix
|
|
import server.common.fbs.NetEncoding.TypedArray as TypedArray
|
|
import server.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)
|
|
Column.ColumnAddU(builder, u_value)
|
|
return Column.ColumnEnd(builder)
|
|
|
|
|
|
# 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)
|
|
Matrix.MatrixAddColumns(builder, columns)
|
|
if col_idx is not None:
|
|
(u_type, u_val) = col_idx
|
|
Matrix.MatrixAddColIndexType(builder, u_type)
|
|
Matrix.MatrixAddColIndex(builder, u_val)
|
|
return Matrix.MatrixEnd(builder)
|
|
|
|
|
|
# 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.
|
|
"""
|
|
|
|
arr = source_array
|
|
(array_type, as_type) = encoding_info(source_array)
|
|
|
|
if isinstance(arr, pd.Index):
|
|
arr = arr.to_series()
|
|
|
|
# convert to a simple ndarray
|
|
if as_type == "json":
|
|
as_json = arr.to_json(orient="records")
|
|
arr = np.array(bytearray(as_json, "utf-8"))
|
|
else:
|
|
if sparse.issparse(arr):
|
|
arr = arr.toarray()
|
|
elif isinstance(arr, pd.Series):
|
|
arr = arr.to_numpy()
|
|
if arr.dtype != as_type:
|
|
arr = arr.astype(as_type)
|
|
|
|
# serialize the ndarray into a vector
|
|
if arr.ndim == 2:
|
|
if arr.shape[0] == 1:
|
|
arr = arr[0]
|
|
elif arr.shape[1] == 1:
|
|
arr = arr.T[0]
|
|
|
|
vec = builder.CreateNumpyVector(arr)
|
|
|
|
# serialize the typed array table
|
|
builder.StartObject(1)
|
|
builder.PrependUOffsetTRelativeSlot(0, vec, 0)
|
|
array_value = builder.EndObject()
|
|
return (array_type, array_value)
|
|
|
|
|
|
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")
|
|
|
|
encoding_dtype = np.dtype(get_encoding_dtype_of_array(arr))
|
|
return column_encoding_type_map.get(encoding_dtype.str, column_encoding_default)
|
|
|
|
|
|
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)
|
|
|
|
|
|
def guess_at_mem_needed(matrix):
|
|
(n_rows, n_cols) = matrix.shape
|
|
if isinstance(matrix, np.ndarray) or sparse.issparse(matrix):
|
|
guess = (n_rows * n_cols * matrix.dtype.itemsize) + 1024
|
|
elif isinstance(matrix, pd.DataFrame):
|
|
# XXX TODO - DataFrame type estimate
|
|
guess = 1
|
|
else:
|
|
guess = 1
|
|
|
|
# round up to nearest 1024 bytes
|
|
guess = (guess + 0x400) & (~0x3FF)
|
|
return guess
|
|
|
|
|
|
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.
|
|
|
|
:param matrix: 2D DataFrame, ndarray or sparse equivalent
|
|
:param row_idx: index for row dimension, Index or ndarray
|
|
:param col_idx: index for col dimension, Index or ndarray
|
|
|
|
NOTE: row indices are (currently) unsupported and must be None
|
|
"""
|
|
|
|
if row_idx is not None:
|
|
raise ValueError("row indexing not supported for FBS Matrix")
|
|
if matrix.ndim != 2:
|
|
raise ValueError("FBS Matrix must be 2D")
|
|
|
|
(n_rows, n_cols) = matrix.shape
|
|
|
|
# estimate size needed, so we don't unnecessarily realloc.
|
|
builder = Builder(guess_at_mem_needed(matrix))
|
|
|
|
columns = []
|
|
for cidx in range(n_cols - 1, -1, -1):
|
|
# serialize the typed array
|
|
col = matrix.iloc[:, cidx] if isinstance(matrix, pd.DataFrame) else matrix[:, cidx]
|
|
typed_arr = serialize_typed_array(builder, col, column_encoding)
|
|
|
|
# serialize the Column union
|
|
columns.append(serialize_column(builder, typed_arr))
|
|
|
|
# Serialize Matrix.columns[]
|
|
Matrix.MatrixStartColumnsVector(builder, n_cols)
|
|
for c in columns:
|
|
builder.PrependUOffsetTRelative(c)
|
|
matrix_column_vec = builder.EndVector(n_cols)
|
|
|
|
# serialize the colIndex if provided
|
|
cidx = None
|
|
if col_idx is not None:
|
|
cidx = serialize_typed_array(builder, col_idx, index_encoding)
|
|
|
|
# Serialize Matrix
|
|
matrix = serialize_matrix(builder, n_rows, n_cols, matrix_column_vec, cidx)
|
|
|
|
builder.Finish(matrix)
|
|
return builder.Output()
|
|
|
|
|
|
def deserialize_typed_array(tarr):
|
|
type_map = {
|
|
TypedArray.TypedArray.NONE: None,
|
|
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 is TypedArray.TypedArray.NONE:
|
|
return None
|
|
|
|
TarType = type_map.get(u_type, None)
|
|
if TarType is None:
|
|
raise TypeError(f"FBS contains unknown data type: {u_type}")
|
|
|
|
arr = TarType()
|
|
arr.Init(u.Bytes, u.Pos)
|
|
narr = arr.DataAsNumpy()
|
|
if u_type == TypedArray.TypedArray.JSONEncodedArray:
|
|
narr = json.loads(narr.tobytes().decode("utf-8"))
|
|
return narr
|
|
|
|
|
|
def decode_matrix_fbs(fbs):
|
|
"""
|
|
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()
|
|
if n_rows == 0 or n_cols == 0:
|
|
return pd.DataFrame()
|
|
|
|
if matrix.RowIndexType() is not TypedArray.TypedArray.NONE:
|
|
raise ValueError("row indexing not supported for FBS Matrix")
|
|
|
|
columns_length = matrix.ColumnsLength()
|
|
|
|
columns_index = deserialize_typed_array((matrix.ColIndexType(), matrix.ColIndex()))
|
|
if columns_index is None:
|
|
columns_index = range(0, n_cols)
|
|
|
|
# sanity checks
|
|
if len(columns_index) != n_cols or columns_length != n_cols:
|
|
raise ValueError("FBS column count does not match number of columns in underlying matrix")
|
|
|
|
columns_data = {}
|
|
columns_type = {}
|
|
for col_idx in range(0, columns_length):
|
|
col = matrix.Columns(col_idx)
|
|
tarr = (col.UType(), col.U())
|
|
data = deserialize_typed_array(tarr)
|
|
columns_data[columns_index[col_idx]] = data
|
|
if len(data) != n_rows:
|
|
raise ValueError("FBS column length does not match number of rows")
|
|
if col.UType() is TypedArray.TypedArray.JSONEncodedArray:
|
|
columns_type[columns_index[col_idx]] = "category"
|
|
|
|
df = pd.DataFrame.from_dict(data=columns_data).astype(columns_type, copy=False)
|
|
|
|
# more sanity checks
|
|
if not df.columns.is_unique or len(df.columns) != n_cols:
|
|
raise KeyError("FBS column indices are not unique")
|
|
|
|
return df
|