mirror of
https://github.com/chanzuckerberg/cellxgene.git
synced 2026-09-28 23:58:11 +08:00
Refactor czi_hosted and server into backend directory, pull common code into backend/common, refactor tests (#2102)
* move local_server -> backend/server server-> backend/czi_hosted, pull common code into backend/common update imports, tests and make commands
This commit is contained in:
@@ -1,396 +0,0 @@
|
||||
from abc import ABCMeta, abstractmethod
|
||||
from os.path import basename, splitext
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from server_timing import Timing as ServerTiming
|
||||
|
||||
from server.common.config.app_config import AppConfig
|
||||
from server.common.constants import Axis
|
||||
from server.common.errors import FilterError, JSONEncodingValueError, ExceedsLimitError
|
||||
from server.common.utils.utils import jsonify_numpy
|
||||
from server.data_common.fbs.matrix import encode_matrix_fbs
|
||||
|
||||
|
||||
class DataAdaptor(metaclass=ABCMeta):
|
||||
"""Base class for loading and accessing matrix data"""
|
||||
|
||||
def __init__(self, data_locator, app_config, dataset_config=None):
|
||||
if type(app_config) != AppConfig:
|
||||
raise TypeError("config expected to be of type AppConfig")
|
||||
|
||||
# location to the dataset
|
||||
self.data_locator = data_locator
|
||||
|
||||
# config is the application configuration
|
||||
self.app_config = app_config
|
||||
self.server_config = self.app_config.server_config
|
||||
self.dataset_config = dataset_config or app_config.default_dataset_config
|
||||
|
||||
# parameters set by this data adaptor based on the data.
|
||||
self.parameters = {}
|
||||
self.uri_path = None
|
||||
|
||||
def set_uri_path(self, path):
|
||||
# uri path to the dataset, e.g. /d/<datasetname>
|
||||
self.uri_path = path
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def pre_load_validation(data_locator):
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def open(data_locator, app_config, dataset_config):
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def file_size(data_locator):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_name(self):
|
||||
"""return a string name for this data adaptor"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_library_versions(self):
|
||||
"""return a dictionary of library name to library versions"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_embedding_names(self):
|
||||
"""return a list of pre-computed embedding names"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_embedding_array(self, ename, dims=2):
|
||||
"""return an numpy array for the given pre-computed embedding name."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def compute_embedding(self, method, filter):
|
||||
"""compute a new embedding on the specified obs subset, and return the embedding schema. """
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_X_array(self, obs_mask=None, var_mask=None):
|
||||
"""return the X array, possibly filtered by obs_mask or var_mask.
|
||||
the return type is either ndarray or scipy.sparse.spmatrix."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_shape(self):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def query_var_array(self, term_var):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def query_obs_array(self, term_var):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_colors(self):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_obs_index(self):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_obs_columns(self):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_obs_keys(self):
|
||||
# return list of keys
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_var_keys(self):
|
||||
# return list of keys
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def cleanup(self):
|
||||
pass
|
||||
|
||||
def get_data_locator(self):
|
||||
return self.data_locator
|
||||
|
||||
def get_location(self):
|
||||
return self.data_locator.uri_or_path
|
||||
|
||||
def get_about(self):
|
||||
return None
|
||||
|
||||
def get_title(self):
|
||||
# default to file name
|
||||
location = self.get_location()
|
||||
if location.endswith("/"):
|
||||
location = location[:-1]
|
||||
return splitext(basename(location))[0]
|
||||
|
||||
def get_corpora_props(self):
|
||||
return None
|
||||
|
||||
@abstractmethod
|
||||
def get_schema(self):
|
||||
"""
|
||||
Return current schema
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def annotation_to_fbs_matrix(self, axis, field=None, uid=None):
|
||||
"""
|
||||
Gets annotation value for each observation
|
||||
:param axis: string obs or var
|
||||
:param fields: list of keys for annotation to return, returns all annotation values if not set.
|
||||
:return: flatbuffer: in fbs/matrix.fbs encoding
|
||||
"""
|
||||
pass
|
||||
|
||||
def update_parameters(self, parameters):
|
||||
parameters.update(self.parameters)
|
||||
|
||||
def _index_filter_to_mask(self, filter, count):
|
||||
mask = np.zeros((count,), dtype=np.bool)
|
||||
for i in filter:
|
||||
if type(i) == list:
|
||||
mask[i[0] : i[1]] = True
|
||||
else:
|
||||
mask[i] = True
|
||||
return mask
|
||||
|
||||
def _axis_filter_to_mask(self, axis, filter, count):
|
||||
mask = np.ones((count,), dtype=np.bool)
|
||||
if "index" in filter:
|
||||
mask = np.logical_and(mask, self._index_filter_to_mask(filter["index"], count))
|
||||
if "annotation_value" in filter:
|
||||
mask = np.logical_and(mask, self._annotation_filter_to_mask(axis, filter["annotation_value"], count))
|
||||
|
||||
return mask
|
||||
|
||||
def _annotation_filter_to_mask(self, axis, filter, count):
|
||||
mask = np.ones((count,), dtype=np.bool)
|
||||
for v in filter:
|
||||
name = v["name"]
|
||||
if axis == Axis.VAR:
|
||||
anno_data = self.query_var_array(name)
|
||||
elif axis == Axis.OBS:
|
||||
anno_data = self.query_obs_array(name)
|
||||
|
||||
if anno_data.dtype.name in ["boolean", "category", "object"]:
|
||||
values = v.get("values", [])
|
||||
key_idx = np.in1d(anno_data, values)
|
||||
mask = np.logical_and(mask, key_idx)
|
||||
|
||||
else:
|
||||
min_ = v.get("min", None)
|
||||
max_ = v.get("max", None)
|
||||
if min_ is not None:
|
||||
key_idx = (anno_data >= min_).ravel()
|
||||
mask = np.logical_and(mask, key_idx)
|
||||
if max_ is not None:
|
||||
key_idx = (anno_data <= max_).ravel()
|
||||
mask = np.logical_and(mask, key_idx)
|
||||
|
||||
return mask
|
||||
|
||||
def _filter_to_mask(self, filter):
|
||||
"""
|
||||
Return the filter as a row and column selection list.
|
||||
No filter on a dimension means 'all'
|
||||
"""
|
||||
shape = self.get_shape()
|
||||
var_selector = None
|
||||
obs_selector = None
|
||||
if filter is not None:
|
||||
if Axis.OBS in filter:
|
||||
obs_selector = self._axis_filter_to_mask(Axis.OBS, filter["obs"], shape[0])
|
||||
|
||||
if Axis.VAR in filter:
|
||||
var_selector = self._axis_filter_to_mask(Axis.VAR, filter["var"], shape[1])
|
||||
|
||||
return (obs_selector, var_selector)
|
||||
|
||||
def check_new_labels(self, labels_df):
|
||||
"""Check the new annotations labels, then set the labels_df index"""
|
||||
if labels_df is None or labels_df.empty:
|
||||
return
|
||||
|
||||
labels_df.index = self.get_obs_index()
|
||||
if labels_df.index.name is None:
|
||||
labels_df.index.name = "index"
|
||||
|
||||
# all labels must have a name, which must be unique and not used in obs column names
|
||||
if not labels_df.columns.is_unique:
|
||||
raise KeyError("All column names specified in user annotations must be unique.")
|
||||
|
||||
# the label index must be unique, and must have same values the anndata obs index
|
||||
if not labels_df.index.is_unique:
|
||||
raise KeyError("All row index values specified in user annotations must be unique.")
|
||||
|
||||
obs_columns = self.get_obs_columns()
|
||||
|
||||
duplicate_columns = list(set(labels_df.columns) & set(obs_columns))
|
||||
if len(duplicate_columns) > 0:
|
||||
raise KeyError(
|
||||
"Labels file may not contain column names which overlap " f"with h5ad obs columns {duplicate_columns}"
|
||||
)
|
||||
|
||||
# labels must have same count as obs annotations
|
||||
shape = self.get_shape()
|
||||
if labels_df.shape[0] != shape[0]:
|
||||
raise ValueError("Labels file must have same number of rows as data file.")
|
||||
|
||||
# This will convert a float column that contains integer data into an integer type.
|
||||
# This case can occur when a user makes a copy of a category that originally contained integer data.
|
||||
# The client always copies array data to floats, therefore the copy will contain floats instead of integers.
|
||||
# float data is not allowed as a categorical type.
|
||||
if any([np.issubdtype(coltype.type, np.floating) for coltype in labels_df.dtypes]):
|
||||
labels_df = labels_df.convert_dtypes()
|
||||
for col, dtype in zip(labels_df, labels_df.dtypes):
|
||||
if isinstance(dtype, pd.Int32Dtype):
|
||||
labels_df[col] = labels_df[col].astype("int32")
|
||||
if isinstance(dtype, pd.Int64Dtype):
|
||||
labels_df[col] = labels_df[col].astype("int64")
|
||||
|
||||
if any([np.issubdtype(coltype.type, np.floating) for coltype in labels_df.dtypes]):
|
||||
raise ValueError("Columns may not have floating point types")
|
||||
|
||||
return labels_df
|
||||
|
||||
def data_frame_to_fbs_matrix(self, filter, axis):
|
||||
"""
|
||||
Retrieves data 'X' and returns in a flatbuffer Matrix.
|
||||
:param filter: filter: dictionary with filter params
|
||||
:param axis: string obs or var
|
||||
:return: flatbuffer Matrix
|
||||
|
||||
Caveats:
|
||||
* currently only supports access on VAR axis
|
||||
* currently only supports filtering on VAR axis
|
||||
"""
|
||||
if axis != Axis.VAR:
|
||||
raise ValueError("Only VAR dimension access is supported")
|
||||
|
||||
try:
|
||||
obs_selector, var_selector = self._filter_to_mask(filter)
|
||||
except (KeyError, IndexError, TypeError, AttributeError):
|
||||
raise FilterError("Error parsing filter")
|
||||
|
||||
if obs_selector is not None:
|
||||
raise FilterError("filtering on obs unsupported")
|
||||
|
||||
num_columns = self.get_shape()[1] if var_selector is None else np.count_nonzero(var_selector)
|
||||
if self.server_config.exceeds_limit("column_request_max", num_columns):
|
||||
raise ExceedsLimitError("Requested dataframe columns exceed column request limit")
|
||||
|
||||
X = self.get_X_array(obs_selector, var_selector)
|
||||
col_idx = np.nonzero([] if var_selector is None else var_selector)[0]
|
||||
return encode_matrix_fbs(X, col_idx=col_idx, row_idx=None)
|
||||
|
||||
def diffexp_topN(self, obsFilterA, obsFilterB, top_n=None):
|
||||
"""
|
||||
Computes the top N differentially expressed variables between two observation sets. If mode
|
||||
is "TOP_N", then stats for the top N
|
||||
dataframes
|
||||
contain a subset of variables, then statistics for all variables will be returned, otherwise
|
||||
only the top N vars will be returned.
|
||||
:param obsFilterA: filter: dictionary with filter params for first set of observations
|
||||
:param obsFilterB: filter: dictionary with filter params for second set of observations
|
||||
:param top_n: Limit results to top N (Top var mode only)
|
||||
:return: top N genes and corresponding stats
|
||||
"""
|
||||
if Axis.VAR in obsFilterA or Axis.VAR in obsFilterB:
|
||||
raise FilterError("Observation filters may not contain variable conditions")
|
||||
try:
|
||||
shape = self.get_shape()
|
||||
obs_mask_A = self._axis_filter_to_mask(Axis.OBS, obsFilterA["obs"], shape[0])
|
||||
obs_mask_B = self._axis_filter_to_mask(Axis.OBS, obsFilterB["obs"], shape[0])
|
||||
except (KeyError, IndexError):
|
||||
raise FilterError("Error parsing filter")
|
||||
if top_n is None:
|
||||
top_n = self.dataset_config.diffexp__top_n
|
||||
|
||||
if self.server_config.exceeds_limit(
|
||||
"diffexp_cellcount_max", np.count_nonzero(obs_mask_A) + np.count_nonzero(obs_mask_B)
|
||||
):
|
||||
raise ExceedsLimitError("Diffexp request exceeds max cell count limit")
|
||||
|
||||
result = self.compute_diffexp_ttest(obs_mask_A, obs_mask_B, top_n, self.dataset_config.diffexp__lfc_cutoff)
|
||||
|
||||
try:
|
||||
return jsonify_numpy(result)
|
||||
except ValueError:
|
||||
raise JSONEncodingValueError("Error encoding differential expression to JSON")
|
||||
|
||||
@abstractmethod
|
||||
def compute_diffexp_ttest(self, maskA, maskB, top_n, lfc_cutoff):
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def normalize_embedding(embedding):
|
||||
"""Normalize embedding layout to meet client assumptions.
|
||||
Embedding is an ndarray, shape (n_obs, n)., where n is normally 2
|
||||
"""
|
||||
|
||||
# scale isotropically
|
||||
try:
|
||||
min = np.nanmin(embedding, axis=0)
|
||||
max = np.nanmax(embedding, axis=0)
|
||||
except RuntimeError:
|
||||
# indicates entire array was NaN, which should propagate
|
||||
min = np.NaN
|
||||
max = np.NaN
|
||||
|
||||
scale = np.amax(max - min)
|
||||
normalized_layout = (embedding - min) / scale
|
||||
|
||||
# translate to center on both axis
|
||||
translate = 0.5 - ((max - min) / scale / 2)
|
||||
normalized_layout = normalized_layout + translate
|
||||
|
||||
normalized_layout = normalized_layout.astype(dtype=np.float32)
|
||||
return normalized_layout
|
||||
|
||||
def layout_to_fbs_matrix(self, fields):
|
||||
"""
|
||||
return specified embeddings as a flatbuffer, using the cellxgene matrix fbs encoding.
|
||||
|
||||
* returns only first two dimensions, with name {ename}_0 and {ename}_1,
|
||||
where {ename} is the embedding name.
|
||||
* client assumes each will be individually centered & scaled (isotropically)
|
||||
to a [0, 1] range.
|
||||
* does not support filtering
|
||||
|
||||
"""
|
||||
embeddings = self.get_embedding_names() if fields is None or len(fields) == 0 else fields
|
||||
layout_data = []
|
||||
with ServerTiming.time("layout.query"):
|
||||
for ename in embeddings:
|
||||
embedding = self.get_embedding_array(ename, 2)
|
||||
normalized_layout = DataAdaptor.normalize_embedding(embedding)
|
||||
layout_data.append(pd.DataFrame(normalized_layout, columns=[f"{ename}_0", f"{ename}_1"]))
|
||||
|
||||
with ServerTiming.time("layout.encode"):
|
||||
if layout_data:
|
||||
df = pd.concat(layout_data, axis=1, copy=False)
|
||||
else:
|
||||
df = pd.DataFrame()
|
||||
fbs = encode_matrix_fbs(df, col_idx=df.columns, row_idx=None)
|
||||
|
||||
return fbs
|
||||
|
||||
def get_last_mod_time(self):
|
||||
try:
|
||||
lastmod = self.get_data_locator().lastmodtime()
|
||||
except RuntimeError:
|
||||
lastmod = None
|
||||
return lastmod
|
||||
@@ -1,41 +0,0 @@
|
||||
# automatically generated by the FlatBuffers compiler, do not modify
|
||||
|
||||
# namespace: NetEncoding
|
||||
|
||||
import flatbuffers
|
||||
|
||||
class Column(object):
|
||||
__slots__ = ['_tab']
|
||||
|
||||
@classmethod
|
||||
def GetRootAsColumn(cls, buf, offset):
|
||||
n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
|
||||
x = Column()
|
||||
x.Init(buf, n + offset)
|
||||
return x
|
||||
|
||||
# Column
|
||||
def Init(self, buf, pos):
|
||||
self._tab = flatbuffers.table.Table(buf, pos)
|
||||
|
||||
# Column
|
||||
def UType(self):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
|
||||
if o != 0:
|
||||
return self._tab.Get(flatbuffers.number_types.Uint8Flags, o + self._tab.Pos)
|
||||
return 0
|
||||
|
||||
# Column
|
||||
def U(self):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(6))
|
||||
if o != 0:
|
||||
from flatbuffers.table import Table
|
||||
obj = Table(bytearray(), 0)
|
||||
self._tab.Union(obj, o)
|
||||
return obj
|
||||
return None
|
||||
|
||||
def ColumnStart(builder): builder.StartObject(2)
|
||||
def ColumnAddUType(builder, uType): builder.PrependUint8Slot(0, uType, 0)
|
||||
def ColumnAddU(builder, u): builder.PrependUOffsetTRelativeSlot(1, flatbuffers.number_types.UOffsetTFlags.py_type(u), 0)
|
||||
def ColumnEnd(builder): return builder.EndObject()
|
||||
@@ -1,46 +0,0 @@
|
||||
# automatically generated by the FlatBuffers compiler, do not modify
|
||||
|
||||
# namespace: NetEncoding
|
||||
|
||||
import flatbuffers
|
||||
|
||||
class Float32Array(object):
|
||||
__slots__ = ['_tab']
|
||||
|
||||
@classmethod
|
||||
def GetRootAsFloat32Array(cls, buf, offset):
|
||||
n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
|
||||
x = Float32Array()
|
||||
x.Init(buf, n + offset)
|
||||
return x
|
||||
|
||||
# Float32Array
|
||||
def Init(self, buf, pos):
|
||||
self._tab = flatbuffers.table.Table(buf, pos)
|
||||
|
||||
# Float32Array
|
||||
def Data(self, j):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
|
||||
if o != 0:
|
||||
a = self._tab.Vector(o)
|
||||
return self._tab.Get(flatbuffers.number_types.Float32Flags, a + flatbuffers.number_types.UOffsetTFlags.py_type(j * 4))
|
||||
return 0
|
||||
|
||||
# Float32Array
|
||||
def DataAsNumpy(self):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
|
||||
if o != 0:
|
||||
return self._tab.GetVectorAsNumpy(flatbuffers.number_types.Float32Flags, o)
|
||||
return 0
|
||||
|
||||
# Float32Array
|
||||
def DataLength(self):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
|
||||
if o != 0:
|
||||
return self._tab.VectorLen(o)
|
||||
return 0
|
||||
|
||||
def Float32ArrayStart(builder): builder.StartObject(1)
|
||||
def Float32ArrayAddData(builder, data): builder.PrependUOffsetTRelativeSlot(0, flatbuffers.number_types.UOffsetTFlags.py_type(data), 0)
|
||||
def Float32ArrayStartDataVector(builder, numElems): return builder.StartVector(4, numElems, 4)
|
||||
def Float32ArrayEnd(builder): return builder.EndObject()
|
||||
@@ -1,46 +0,0 @@
|
||||
# automatically generated by the FlatBuffers compiler, do not modify
|
||||
|
||||
# namespace: NetEncoding
|
||||
|
||||
import flatbuffers
|
||||
|
||||
class Float64Array(object):
|
||||
__slots__ = ['_tab']
|
||||
|
||||
@classmethod
|
||||
def GetRootAsFloat64Array(cls, buf, offset):
|
||||
n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
|
||||
x = Float64Array()
|
||||
x.Init(buf, n + offset)
|
||||
return x
|
||||
|
||||
# Float64Array
|
||||
def Init(self, buf, pos):
|
||||
self._tab = flatbuffers.table.Table(buf, pos)
|
||||
|
||||
# Float64Array
|
||||
def Data(self, j):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
|
||||
if o != 0:
|
||||
a = self._tab.Vector(o)
|
||||
return self._tab.Get(flatbuffers.number_types.Float64Flags, a + flatbuffers.number_types.UOffsetTFlags.py_type(j * 8))
|
||||
return 0
|
||||
|
||||
# Float64Array
|
||||
def DataAsNumpy(self):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
|
||||
if o != 0:
|
||||
return self._tab.GetVectorAsNumpy(flatbuffers.number_types.Float64Flags, o)
|
||||
return 0
|
||||
|
||||
# Float64Array
|
||||
def DataLength(self):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
|
||||
if o != 0:
|
||||
return self._tab.VectorLen(o)
|
||||
return 0
|
||||
|
||||
def Float64ArrayStart(builder): builder.StartObject(1)
|
||||
def Float64ArrayAddData(builder, data): builder.PrependUOffsetTRelativeSlot(0, flatbuffers.number_types.UOffsetTFlags.py_type(data), 0)
|
||||
def Float64ArrayStartDataVector(builder, numElems): return builder.StartVector(8, numElems, 8)
|
||||
def Float64ArrayEnd(builder): return builder.EndObject()
|
||||
@@ -1,46 +0,0 @@
|
||||
# automatically generated by the FlatBuffers compiler, do not modify
|
||||
|
||||
# namespace: NetEncoding
|
||||
|
||||
import flatbuffers
|
||||
|
||||
class Int32Array(object):
|
||||
__slots__ = ['_tab']
|
||||
|
||||
@classmethod
|
||||
def GetRootAsInt32Array(cls, buf, offset):
|
||||
n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
|
||||
x = Int32Array()
|
||||
x.Init(buf, n + offset)
|
||||
return x
|
||||
|
||||
# Int32Array
|
||||
def Init(self, buf, pos):
|
||||
self._tab = flatbuffers.table.Table(buf, pos)
|
||||
|
||||
# Int32Array
|
||||
def Data(self, j):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
|
||||
if o != 0:
|
||||
a = self._tab.Vector(o)
|
||||
return self._tab.Get(flatbuffers.number_types.Int32Flags, a + flatbuffers.number_types.UOffsetTFlags.py_type(j * 4))
|
||||
return 0
|
||||
|
||||
# Int32Array
|
||||
def DataAsNumpy(self):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
|
||||
if o != 0:
|
||||
return self._tab.GetVectorAsNumpy(flatbuffers.number_types.Int32Flags, o)
|
||||
return 0
|
||||
|
||||
# Int32Array
|
||||
def DataLength(self):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
|
||||
if o != 0:
|
||||
return self._tab.VectorLen(o)
|
||||
return 0
|
||||
|
||||
def Int32ArrayStart(builder): builder.StartObject(1)
|
||||
def Int32ArrayAddData(builder, data): builder.PrependUOffsetTRelativeSlot(0, flatbuffers.number_types.UOffsetTFlags.py_type(data), 0)
|
||||
def Int32ArrayStartDataVector(builder, numElems): return builder.StartVector(4, numElems, 4)
|
||||
def Int32ArrayEnd(builder): return builder.EndObject()
|
||||
@@ -1,46 +0,0 @@
|
||||
# automatically generated by the FlatBuffers compiler, do not modify
|
||||
|
||||
# namespace: NetEncoding
|
||||
|
||||
import flatbuffers
|
||||
|
||||
class JSONEncodedArray(object):
|
||||
__slots__ = ['_tab']
|
||||
|
||||
@classmethod
|
||||
def GetRootAsJSONEncodedArray(cls, buf, offset):
|
||||
n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
|
||||
x = JSONEncodedArray()
|
||||
x.Init(buf, n + offset)
|
||||
return x
|
||||
|
||||
# JSONEncodedArray
|
||||
def Init(self, buf, pos):
|
||||
self._tab = flatbuffers.table.Table(buf, pos)
|
||||
|
||||
# JSONEncodedArray
|
||||
def Data(self, j):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
|
||||
if o != 0:
|
||||
a = self._tab.Vector(o)
|
||||
return self._tab.Get(flatbuffers.number_types.Uint8Flags, a + flatbuffers.number_types.UOffsetTFlags.py_type(j * 1))
|
||||
return 0
|
||||
|
||||
# JSONEncodedArray
|
||||
def DataAsNumpy(self):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
|
||||
if o != 0:
|
||||
return self._tab.GetVectorAsNumpy(flatbuffers.number_types.Uint8Flags, o)
|
||||
return 0
|
||||
|
||||
# JSONEncodedArray
|
||||
def DataLength(self):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
|
||||
if o != 0:
|
||||
return self._tab.VectorLen(o)
|
||||
return 0
|
||||
|
||||
def JSONEncodedArrayStart(builder): builder.StartObject(1)
|
||||
def JSONEncodedArrayAddData(builder, data): builder.PrependUOffsetTRelativeSlot(0, flatbuffers.number_types.UOffsetTFlags.py_type(data), 0)
|
||||
def JSONEncodedArrayStartDataVector(builder, numElems): return builder.StartVector(1, numElems, 1)
|
||||
def JSONEncodedArrayEnd(builder): return builder.EndObject()
|
||||
@@ -1,98 +0,0 @@
|
||||
# automatically generated by the FlatBuffers compiler, do not modify
|
||||
|
||||
# namespace: NetEncoding
|
||||
|
||||
import flatbuffers
|
||||
|
||||
class Matrix(object):
|
||||
__slots__ = ['_tab']
|
||||
|
||||
@classmethod
|
||||
def GetRootAsMatrix(cls, buf, offset):
|
||||
n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
|
||||
x = Matrix()
|
||||
x.Init(buf, n + offset)
|
||||
return x
|
||||
|
||||
# Matrix
|
||||
def Init(self, buf, pos):
|
||||
self._tab = flatbuffers.table.Table(buf, pos)
|
||||
|
||||
# Matrix
|
||||
def NRows(self):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
|
||||
if o != 0:
|
||||
return self._tab.Get(flatbuffers.number_types.Uint32Flags, o + self._tab.Pos)
|
||||
return 0
|
||||
|
||||
# Matrix
|
||||
def NCols(self):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(6))
|
||||
if o != 0:
|
||||
return self._tab.Get(flatbuffers.number_types.Uint32Flags, o + self._tab.Pos)
|
||||
return 0
|
||||
|
||||
# Matrix
|
||||
def Columns(self, j):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(8))
|
||||
if o != 0:
|
||||
x = self._tab.Vector(o)
|
||||
x += flatbuffers.number_types.UOffsetTFlags.py_type(j) * 4
|
||||
x = self._tab.Indirect(x)
|
||||
from .Column import Column
|
||||
obj = Column()
|
||||
obj.Init(self._tab.Bytes, x)
|
||||
return obj
|
||||
return None
|
||||
|
||||
# Matrix
|
||||
def ColumnsLength(self):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(8))
|
||||
if o != 0:
|
||||
return self._tab.VectorLen(o)
|
||||
return 0
|
||||
|
||||
# Matrix
|
||||
def ColIndexType(self):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(10))
|
||||
if o != 0:
|
||||
return self._tab.Get(flatbuffers.number_types.Uint8Flags, o + self._tab.Pos)
|
||||
return 0
|
||||
|
||||
# Matrix
|
||||
def ColIndex(self):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(12))
|
||||
if o != 0:
|
||||
from flatbuffers.table import Table
|
||||
obj = Table(bytearray(), 0)
|
||||
self._tab.Union(obj, o)
|
||||
return obj
|
||||
return None
|
||||
|
||||
# Matrix
|
||||
def RowIndexType(self):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(14))
|
||||
if o != 0:
|
||||
return self._tab.Get(flatbuffers.number_types.Uint8Flags, o + self._tab.Pos)
|
||||
return 0
|
||||
|
||||
# Matrix
|
||||
def RowIndex(self):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(16))
|
||||
if o != 0:
|
||||
from flatbuffers.table import Table
|
||||
obj = Table(bytearray(), 0)
|
||||
self._tab.Union(obj, o)
|
||||
return obj
|
||||
return None
|
||||
|
||||
def MatrixStart(builder): builder.StartObject(7)
|
||||
def MatrixAddNRows(builder, nRows): builder.PrependUint32Slot(0, nRows, 0)
|
||||
def MatrixAddNCols(builder, nCols): builder.PrependUint32Slot(1, nCols, 0)
|
||||
def MatrixAddColumns(builder, columns): builder.PrependUOffsetTRelativeSlot(2, flatbuffers.number_types.UOffsetTFlags.py_type(columns), 0)
|
||||
def MatrixStartColumnsVector(builder, numElems): return builder.StartVector(4, numElems, 4)
|
||||
def MatrixAddColIndexType(builder, colIndexType): builder.PrependUint8Slot(3, colIndexType, 0)
|
||||
def MatrixAddColIndex(builder, colIndex): builder.PrependUOffsetTRelativeSlot(4, flatbuffers.number_types.UOffsetTFlags.py_type(colIndex), 0)
|
||||
def MatrixAddRowIndexType(builder, rowIndexType): builder.PrependUint8Slot(5, rowIndexType, 0)
|
||||
def MatrixAddRowIndex(builder, rowIndex): builder.PrependUOffsetTRelativeSlot(6, flatbuffers.number_types.UOffsetTFlags.py_type(rowIndex), 0)
|
||||
def MatrixEnd(builder): return builder.EndObject()
|
||||
@@ -1,12 +0,0 @@
|
||||
# automatically generated by the FlatBuffers compiler, do not modify
|
||||
|
||||
# namespace: NetEncoding
|
||||
|
||||
class TypedArray(object):
|
||||
NONE = 0
|
||||
Float32Array = 1
|
||||
Int32Array = 2
|
||||
Uint32Array = 3
|
||||
Float64Array = 4
|
||||
JSONEncodedArray = 5
|
||||
|
||||
@@ -1,46 +0,0 @@
|
||||
# automatically generated by the FlatBuffers compiler, do not modify
|
||||
|
||||
# namespace: NetEncoding
|
||||
|
||||
import flatbuffers
|
||||
|
||||
class Uint32Array(object):
|
||||
__slots__ = ['_tab']
|
||||
|
||||
@classmethod
|
||||
def GetRootAsUint32Array(cls, buf, offset):
|
||||
n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
|
||||
x = Uint32Array()
|
||||
x.Init(buf, n + offset)
|
||||
return x
|
||||
|
||||
# Uint32Array
|
||||
def Init(self, buf, pos):
|
||||
self._tab = flatbuffers.table.Table(buf, pos)
|
||||
|
||||
# Uint32Array
|
||||
def Data(self, j):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
|
||||
if o != 0:
|
||||
a = self._tab.Vector(o)
|
||||
return self._tab.Get(flatbuffers.number_types.Uint32Flags, a + flatbuffers.number_types.UOffsetTFlags.py_type(j * 4))
|
||||
return 0
|
||||
|
||||
# Uint32Array
|
||||
def DataAsNumpy(self):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
|
||||
if o != 0:
|
||||
return self._tab.GetVectorAsNumpy(flatbuffers.number_types.Uint32Flags, o)
|
||||
return 0
|
||||
|
||||
# Uint32Array
|
||||
def DataLength(self):
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
|
||||
if o != 0:
|
||||
return self._tab.VectorLen(o)
|
||||
return 0
|
||||
|
||||
def Uint32ArrayStart(builder): builder.StartObject(1)
|
||||
def Uint32ArrayAddData(builder, data): builder.PrependUOffsetTRelativeSlot(0, flatbuffers.number_types.UOffsetTFlags.py_type(data), 0)
|
||||
def Uint32ArrayStartDataVector(builder, numElems): return builder.StartVector(4, numElems, 4)
|
||||
def Uint32ArrayEnd(builder): return builder.EndObject()
|
||||
@@ -1,248 +0,0 @@
|
||||
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.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
|
||||
|
||||
|
||||
# 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.Float64Array, np.float64),
|
||||
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)
|
||||
|
||||
|
||||
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.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
|
||||
@@ -1,286 +0,0 @@
|
||||
from enum import Enum
|
||||
import threading
|
||||
import time
|
||||
from server.data_common.rwlock import RWLock
|
||||
from server.common.errors import DatasetAccessError
|
||||
from server.common.data_locator import DataLocator
|
||||
from contextlib import contextmanager
|
||||
from http import HTTPStatus
|
||||
|
||||
|
||||
class MatrixDataCacheItem(object):
|
||||
"""This class provides access and caching for a dataset. The first time a dataset is accessed, it is
|
||||
opened and cached. Later accesses use the cached version. It may also be deleted by the
|
||||
MatrixDataCacheManager to make room for another dataset. While a dataset is actively being used
|
||||
(during the lifetime of a api request), a reader lock is locked. During that time, the dataset cannot
|
||||
be removed."""
|
||||
|
||||
def __init__(self, loader):
|
||||
self.loader = loader
|
||||
self.data_adaptor = None
|
||||
self.data_lock = RWLock()
|
||||
|
||||
def acquire_existing(self):
|
||||
"""If the data_adaptor exists, take a read lock and return it, else return None"""
|
||||
self.data_lock.r_acquire()
|
||||
if self.data_adaptor:
|
||||
return self.data_adaptor
|
||||
|
||||
self.data_lock.r_release()
|
||||
return None
|
||||
|
||||
def acquire_and_open(self, app_config, dataset_config=None):
|
||||
"""returns the data_adaptor if cached. opens the data_adaptor if not.
|
||||
In either case, the a reader lock is taken. Must call release when
|
||||
the data_adaptor is no longer needed"""
|
||||
self.data_lock.r_acquire()
|
||||
if self.data_adaptor:
|
||||
return self.data_adaptor
|
||||
self.data_lock.r_release()
|
||||
|
||||
self.data_lock.w_acquire()
|
||||
# the data may have been loaded while waiting on the lock
|
||||
if not self.data_adaptor:
|
||||
try:
|
||||
self.loader.pre_load_validation()
|
||||
self.data_adaptor = self.loader.open(app_config, dataset_config)
|
||||
except Exception as e:
|
||||
# necessary to hold the reader lock after an exception, since
|
||||
# the release will occur when the context exits.
|
||||
self.data_lock.w_demote()
|
||||
raise DatasetAccessError(str(e))
|
||||
|
||||
# demote the write lock to a read lock.
|
||||
self.data_lock.w_demote()
|
||||
return self.data_adaptor
|
||||
|
||||
def release(self):
|
||||
"""Release the reader lock"""
|
||||
self.data_lock.r_release()
|
||||
|
||||
def delete(self):
|
||||
"""Clear resources used by this dataset"""
|
||||
with self.data_lock.w_locked():
|
||||
if self.data_adaptor:
|
||||
self.data_adaptor.cleanup()
|
||||
self.data_adaptor = None
|
||||
|
||||
def attempt_delete(self):
|
||||
"""Delete, but only if the write lock can be immediately locked. Return True if the delete happened"""
|
||||
if self.data_lock.w_acquire_non_blocking():
|
||||
if self.data_adaptor:
|
||||
try:
|
||||
self.data_adaptor.cleanup()
|
||||
self.data_adaptor = None
|
||||
except Exception:
|
||||
# catch all exceptions to ensure the lock is released
|
||||
pass
|
||||
|
||||
self.data_lock.w_release()
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
|
||||
|
||||
class MatrixDataCacheInfo(object):
|
||||
def __init__(self, cache_item, timestamp):
|
||||
# The MatrixDataCacheItem in the cache
|
||||
self.cache_item = cache_item
|
||||
# The last time the cache_item was accessed
|
||||
self.last_access = timestamp
|
||||
# The number of times the cache_item was accessed (used for testing)
|
||||
self.num_access = 1
|
||||
|
||||
|
||||
class MatrixDataCacheManager(object):
|
||||
"""A class to manage the cached datasets. This is intended to be used as a context manager
|
||||
for handling api requests. When the context is created, the data_adator is either loaded or
|
||||
retrieved from a cache. In either case, the reader lock is taken during this time, and release
|
||||
when the context ends. This class currently implements a simple least recently used cache,
|
||||
which can delete a dataset from the cache to make room for a new one.
|
||||
|
||||
This is the intended usage pattern:
|
||||
|
||||
m = MatrixDataCacheManager(max_cached=..., timelimmit_s = ...)
|
||||
with m.data_adaptor(location, app_config) as data_adaptor:
|
||||
# use the data_adaptor for some operation
|
||||
"""
|
||||
|
||||
# FIXME: If the number of active datasets exceeds the max_cached, then each request could
|
||||
# lead to a dataset being deleted and a new only being opened: the cache will get thrashed.
|
||||
# In this case, we may need to send back a 503 (Server Unavailable), or some other error message.
|
||||
|
||||
# NOTE: If the actual dataset is changed. E.g. a new set of datafiles replaces an existing set,
|
||||
# then the cache will not react to this, however once the cache time limit is reached, the dataset
|
||||
# will automatically be refreshed.
|
||||
|
||||
def __init__(self, max_cached, timelimit_s=None):
|
||||
# key is tuple(url_dataroot, location), value is a MatrixDataCacheInfo
|
||||
self.datasets = {}
|
||||
|
||||
# lock to protect the datasets
|
||||
self.lock = threading.Lock()
|
||||
|
||||
# The number of datasets to cache. When max_cached is reached, the least recently used
|
||||
# cache is replaced with the newly requested one.
|
||||
# TODO: This is very simple. This can be improved by taking into account how much space is actually
|
||||
# taken by each dataset, instead of arbitrarily picking a max datasets to cache.
|
||||
self.max_cached = max_cached
|
||||
|
||||
# items are automatically removed from the cache once this time limit is reached
|
||||
self.timelimit_s = timelimit_s
|
||||
|
||||
@contextmanager
|
||||
def data_adaptor(self, url_dataroot, location, app_config):
|
||||
# create a loader for to this location if it does not already exist
|
||||
|
||||
delete_adaptor = None
|
||||
data_adaptor = None
|
||||
cache_item = None
|
||||
|
||||
key = (url_dataroot, location)
|
||||
with self.lock:
|
||||
self.evict_old_datasets()
|
||||
info = self.datasets.get(key)
|
||||
if info is not None:
|
||||
info.last_access = time.time()
|
||||
info.num_access += 1
|
||||
self.datasets[key] = info
|
||||
data_adaptor = info.cache_item.acquire_existing()
|
||||
cache_item = info.cache_item
|
||||
|
||||
if data_adaptor is None:
|
||||
while True:
|
||||
if len(self.datasets) < self.max_cached:
|
||||
break
|
||||
|
||||
items = list(self.datasets.items())
|
||||
items = sorted(items, key=lambda x: x[1].last_access)
|
||||
# close the least recently used loader
|
||||
oldest = items[0]
|
||||
oldest_cache = oldest[1].cache_item
|
||||
oldest_key = oldest[0]
|
||||
del self.datasets[oldest_key]
|
||||
delete_adaptor = oldest_cache
|
||||
|
||||
loader = MatrixDataLoader(location, app_config=app_config)
|
||||
cache_item = MatrixDataCacheItem(loader)
|
||||
item = MatrixDataCacheInfo(cache_item, time.time())
|
||||
self.datasets[key] = item
|
||||
|
||||
try:
|
||||
assert cache_item
|
||||
if delete_adaptor:
|
||||
delete_adaptor.delete()
|
||||
if data_adaptor is None:
|
||||
dataset_config = app_config.get_dataset_config(url_dataroot)
|
||||
data_adaptor = cache_item.acquire_and_open(app_config, dataset_config)
|
||||
yield data_adaptor
|
||||
except DatasetAccessError:
|
||||
cache_item.release()
|
||||
with self.lock:
|
||||
del self.datasets[key]
|
||||
cache_item.delete()
|
||||
cache_item = None
|
||||
raise
|
||||
|
||||
finally:
|
||||
if cache_item:
|
||||
cache_item.release()
|
||||
|
||||
def evict_old_datasets(self):
|
||||
# must be called with the lock held
|
||||
if self.timelimit_s is None:
|
||||
return
|
||||
|
||||
now = time.time()
|
||||
to_del = []
|
||||
for key, info in self.datasets.items():
|
||||
if (now - info.last_access) > self.timelimit_s:
|
||||
# remove the data_cache when if it has been in the cache too long
|
||||
to_del.append((key, info))
|
||||
|
||||
for key, info in to_del:
|
||||
# try and get the write_lock for the dataset.
|
||||
# if this returns false, it means the dataset is being used, and should
|
||||
# not be removed.
|
||||
if info.cache_item.attempt_delete():
|
||||
del self.datasets[key]
|
||||
|
||||
|
||||
class MatrixDataType(Enum):
|
||||
H5AD = "h5ad"
|
||||
CXG = "cxg"
|
||||
UNKNOWN = "unknown"
|
||||
|
||||
|
||||
class MatrixDataLoader(object):
|
||||
def __init__(self, location, matrix_data_type=None, app_config=None):
|
||||
""" location can be a string or DataLocator """
|
||||
region_name = None if app_config is None else app_config.server_config.data_locator__s3__region_name
|
||||
self.location = DataLocator(location, region_name=region_name)
|
||||
if not self.location.exists():
|
||||
raise DatasetAccessError("Dataset does not exist.", HTTPStatus.NOT_FOUND)
|
||||
|
||||
# matrix_data_type is an enum value of type MatrixDataType
|
||||
self.matrix_data_type = matrix_data_type
|
||||
# matrix_type is a DataAdaptor type, which corresonds to the matrix_data_type
|
||||
self.matrix_type = None
|
||||
|
||||
if matrix_data_type is None:
|
||||
self.matrix_data_type = self.__matrix_data_type()
|
||||
|
||||
if not self.__matrix_data_type_allowed(app_config):
|
||||
raise DatasetAccessError("Dataset does not have an allowed type.")
|
||||
|
||||
if self.matrix_data_type == MatrixDataType.H5AD:
|
||||
from server.data_anndata.anndata_adaptor import AnndataAdaptor
|
||||
|
||||
self.matrix_type = AnndataAdaptor
|
||||
elif self.matrix_data_type == MatrixDataType.CXG:
|
||||
from server.data_cxg.cxg_adaptor import CxgAdaptor
|
||||
|
||||
self.matrix_type = CxgAdaptor
|
||||
|
||||
def __matrix_data_type(self):
|
||||
if self.location.path.endswith(".h5ad"):
|
||||
return MatrixDataType.H5AD
|
||||
elif ".cxg" in self.location.path:
|
||||
return MatrixDataType.CXG
|
||||
else:
|
||||
return MatrixDataType.UNKNOWN
|
||||
|
||||
def __matrix_data_type_allowed(self, app_config):
|
||||
if self.matrix_data_type == MatrixDataType.UNKNOWN:
|
||||
return False
|
||||
|
||||
if not app_config:
|
||||
return True
|
||||
if not app_config.is_multi_dataset():
|
||||
return True
|
||||
if len(app_config.server_config.multi_dataset__allowed_matrix_types) == 0:
|
||||
return True
|
||||
|
||||
for val in app_config.server_config.multi_dataset__allowed_matrix_types:
|
||||
try:
|
||||
if self.matrix_data_type == MatrixDataType(val):
|
||||
return True
|
||||
except ValueError:
|
||||
# Check case where multi_dataset_allowed_matrix_type does not have a
|
||||
# valid MatrixDataType value. TODO: Add a feature to check
|
||||
# the AppConfig for errors on startup
|
||||
return False
|
||||
|
||||
return False
|
||||
|
||||
def pre_load_validation(self):
|
||||
if self.matrix_data_type == MatrixDataType.UNKNOWN:
|
||||
raise DatasetAccessError("Dataset does not have a recognized type: .h5ad or .cxg")
|
||||
self.matrix_type.pre_load_validation(self.location)
|
||||
|
||||
def file_size(self):
|
||||
return self.matrix_type.file_size(self.location)
|
||||
|
||||
def open(self, app_config, dataset_config=None):
|
||||
# create and return a DataAdaptor object
|
||||
return self.matrix_type.open(self.location, app_config, dataset_config)
|
||||
@@ -1,135 +0,0 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
""" rwlock.py
|
||||
|
||||
A class to implement read-write locks on top of the standard threading
|
||||
library.
|
||||
|
||||
This is implemented with two mutexes (threading.Lock instances) as per this
|
||||
wikipedia pseudocode:
|
||||
|
||||
https://en.wikipedia.org/wiki/Readers%E2%80%93writer_lock#Using_two_mutexes
|
||||
|
||||
Code written by Tyler Neylon at Unbox Research.
|
||||
|
||||
This file is public domain.
|
||||
|
||||
Modified to add a w_demote function to convert a writer lock to a reader lock
|
||||
"""
|
||||
|
||||
|
||||
# _______________________________________________________________________
|
||||
# Imports
|
||||
|
||||
from contextlib import contextmanager
|
||||
from threading import Lock
|
||||
|
||||
|
||||
# _______________________________________________________________________
|
||||
# Class
|
||||
|
||||
|
||||
class RWLock(object):
|
||||
""" RWLock class; this is meant to allow an object to be read from by
|
||||
multiple threads, but only written to by a single thread at a time. See:
|
||||
https://en.wikipedia.org/wiki/Readers%E2%80%93writer_lock
|
||||
|
||||
Usage:
|
||||
|
||||
from rwlock import RWLock
|
||||
|
||||
my_obj_rwlock = RWLock()
|
||||
|
||||
# When reading from my_obj:
|
||||
with my_obj_rwlock.r_locked():
|
||||
do_read_only_things_with(my_obj)
|
||||
|
||||
# When writing to my_obj:
|
||||
with my_obj_rwlock.w_locked():
|
||||
mutate(my_obj)
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
|
||||
self.w_lock = Lock()
|
||||
self.num_r_lock = Lock()
|
||||
self.num_r = 0
|
||||
|
||||
# The d_lock is needed to handle the demotion case,
|
||||
# so that the writer can become a reader without releasing the w_lock.
|
||||
# the d_lock is held by the writer, and prevents any other thread from taking the
|
||||
# num_r_lock during that time, which means the writer thread is able to take the
|
||||
# num_r_lock to update the num_r.
|
||||
self.d_lock = Lock()
|
||||
|
||||
# ___________________________________________________________________
|
||||
# Reading methods.
|
||||
|
||||
def r_acquire(self):
|
||||
self.d_lock.acquire()
|
||||
self.num_r_lock.acquire()
|
||||
self.num_r += 1
|
||||
|
||||
if self.num_r == 1:
|
||||
self.w_lock.acquire()
|
||||
|
||||
self.num_r_lock.release()
|
||||
self.d_lock.release()
|
||||
|
||||
def r_release(self):
|
||||
assert self.num_r > 0
|
||||
self.num_r_lock.acquire()
|
||||
self.num_r -= 1
|
||||
if self.num_r == 0:
|
||||
self.w_lock.release()
|
||||
|
||||
self.num_r_lock.release()
|
||||
|
||||
@contextmanager
|
||||
def r_locked(self):
|
||||
""" This method is designed to be used via the `with` statement. """
|
||||
try:
|
||||
self.r_acquire()
|
||||
yield
|
||||
finally:
|
||||
self.r_release()
|
||||
|
||||
# ___________________________________________________________________
|
||||
# Writing methods.
|
||||
|
||||
def w_acquire(self):
|
||||
self.d_lock.acquire()
|
||||
self.w_lock.acquire()
|
||||
|
||||
def w_acquire_non_blocking(self):
|
||||
# if d_lock and w_lock can be acquired without blocking, acquire and return True,
|
||||
# else immediately return False.
|
||||
if self.d_lock.acquire(blocking=False):
|
||||
if self.w_lock.acquire(blocking=False):
|
||||
return True
|
||||
else:
|
||||
self.d_lock.release()
|
||||
return False
|
||||
|
||||
def w_release(self):
|
||||
self.w_lock.release()
|
||||
self.d_lock.release()
|
||||
|
||||
def w_demote(self):
|
||||
"""demote a writer lock to a reader lock"""
|
||||
|
||||
# the d_lock is already held from w_acquire.
|
||||
# releasing the d_lock at the end of this function allows multiple readers.
|
||||
# incrementing num_r makes this thread one of those readers.
|
||||
self.num_r_lock.acquire()
|
||||
self.num_r += 1
|
||||
self.num_r_lock.release()
|
||||
self.d_lock.release()
|
||||
|
||||
@contextmanager
|
||||
def w_locked(self):
|
||||
""" This method is designed to be used via the `with` statement. """
|
||||
try:
|
||||
self.w_acquire()
|
||||
yield
|
||||
finally:
|
||||
self.w_release()
|
||||
Reference in New Issue
Block a user