cellxgene/server/app/util/fbs/matrix.py

import flatbuffers
import numpy as np
from scipy import sparse
import pandas as pd
import json

import server.app.util.fbs.NetEncoding.Column as Column
import server.app.util.fbs.NetEncoding.TypedArray as TypedArray
import server.app.util.fbs.NetEncoding.Matrix as Matrix
import server.app.util.fbs.NetEncoding.Int32Array as Int32Array
import server.app.util.fbs.NetEncoding.Uint32Array as Uint32Array
import server.app.util.fbs.NetEncoding.Float32Array as Float32Array
import server.app.util.fbs.NetEncoding.Float64Array as Float64Array
import server.app.util.fbs.NetEncoding.JSONEncodedArray as JSONEncodedArray
from server.app.util.matrix_proxy import MatrixProxy


# 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)


# 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 MatrixProxy.ismatrixproxy(arr) or 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 = CreateNumpyVector(builder, arr)

    # serialize the typed array table
    builder.StartObject(1)
    builder.PrependUOffsetTRelativeSlot(0, vec, 0)
    array_value = builder.EndObject()
    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):
    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):
    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 = flatbuffers.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.tostring().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