load embeddings in parallel (#1352)

* load embeddings in parallel

* correctly capture unclipped

* test

* another test

* add convenient copy assets target

* cleanup

Makefile

@@ -42,6 +42,10 @@ build-cli: build-client
 build-for-server-dev: clean-server build-client
 	$(call copy_client_assets,client/build,server)
 
+.PHONY: copy-client-assets
+copy-client-assets:
+	$(call copy_client_assets,client/build,server)
+
 # TESTING
 .PHONY: test
 test: unit-test smoke-test

client/src/actions/index.js

@@ -1,6 +1,7 @@
 import _ from "lodash";
 import * as globals from "../globals";
 import { Universe, MatrixFBS } from "../util/stateManager";
+import * as Dataframe from "../util/dataframe";
 import {
   catchErrorsWrap,
   doJsonRequest,
@@ -74,24 +75,29 @@ function varAnnotationFetchAndLoad(dispatch, schema) {
 /*
 return promise fetching layout we need
 */
-function layoutFetchAndLoad(dispatch) {
+function layoutFetchAndLoad(dispatch, schema) {
+  const embeddings = schema?.schema?.layout?.obs ?? [];
+  const embNames = embeddings.map(e => e.name);
+  const baseURL = `${globals.API.prefix}${globals.API.version}layout/obs`;
+
+  const plimit = new PromiseLimit(4);
   return Promise.all(
-    ["layout/obs"]
-      .map(path => {
-        const url = `${globals.API.prefix}${globals.API.version}${path}`;
-        return doBinaryRequest(url);
+    embNames.map(e =>
+      plimit.add(() => {
+        const url = `${baseURL}?layout-name=${encodeURIComponent(e)}`;
+        return doBinaryRequest(url).then(buffer =>
+          Universe.matrixFBSToDataframe(buffer)
+        );
       })
-      .map(rqst => rqst.then(buffer => Universe.matrixFBSToDataframe(buffer)))
-      .map(resp =>
-        resp.then(df =>
-          dispatch({
-            type: "universe: column load success",
-            dim: "obsLayout",
-            dataframe: df
-          })
-        )
-      )
-  );
+    )
+  ).then(dfs => {
+    const df = Dataframe.Dataframe.empty().withColsFromAll(dfs);
+    dispatch({
+      type: "universe: column load success",
+      dim: "obsLayout",
+      dataframe: df
+    });
+  });
 }
 
 /*
@@ -130,7 +136,7 @@ const doInitialDataLoad = () =>
       Step 2 - load the minimum stuff required to display.
       */
       await Promise.all([
-        layoutFetchAndLoad(dispatch),
+        layoutFetchAndLoad(dispatch, schema),
         varAnnotationFetchAndLoad(dispatch, schema)
       ]);
 

client/src/reducers/world.js

@@ -39,14 +39,20 @@ const WorldReducer = (
       /* incremental initial data load - always assumes world == universe */
       const { universe } = nextSharedState;
       const { dim } = action;
+
+      // we don't clip anything except for varData and obsAnnotations
+      let unclipped = state.unclipped;
+      if (dim == "varData" || dim == "obsAnnotations") {
+        unclipped = {
+          ...unclipped,
+          [dim]: universe[dim].clone()
+        };
+      }
       return {
         ...state,
         schema: universe.schema,
         [dim]: universe[dim].clone(),
-        unclipped: {
-          ...state.unclipped,
-          [dim]: universe[dim].clone()
-        }
+        unclipped
       };
     }
 
@@ -106,10 +112,15 @@ const WorldReducer = (
       const { userDefinedGenes, diffexpGenes } = prevSharedState;
       const allTheGenesWeNeed = [
         ...new Set(
-          [userDefinedGenes, diffexpGenes, Object.keys(action.expressionData)].filter(ele => ele).flat()
+          [userDefinedGenes, diffexpGenes, Object.keys(action.expressionData)]
+            .filter(ele => ele)
+            .flat()
         )
       ];
-      unclippedVarData = ControlsHelpers.pruneVarDataCache(unclippedVarData, allTheGenesWeNeed);
+      unclippedVarData = ControlsHelpers.pruneVarDataCache(
+        unclippedVarData,
+        allTheGenesWeNeed
+      );
 
       // at this point, we have the unclipped data in unclippedVarData.
       // Now create clipped.

server/common/rest.py

@@ -10,6 +10,7 @@ from server.common.errors import (
     PrepareError,
     DisabledFeatureError,
     ExceedsLimitError,
+    DatasetAccessError,
 )
 
 import json
@@ -179,14 +180,21 @@ def diffexp_obs_post(request, data_adaptor):
 
 
 def layout_obs_get(request, data_adaptor):
+    fields = request.args.getlist("layout-name", None)
+    num_columns_requested = len(data_adaptor.get_embedding_names()) if len(fields) == 0 else len(fields)
+    if data_adaptor.config.exceeds_limit("column_request_max", num_columns_requested):
+        return abort(HTTPStatus.BAD_REQUEST)
+
     preferred_mimetype = request.accept_mimetypes.best_match(["application/octet-stream"])
     if preferred_mimetype != "application/octet-stream":
         return abort(HTTPStatus.NOT_ACCEPTABLE)
 
     try:
         return make_response(
-            data_adaptor.layout_to_fbs_matrix(), HTTPStatus.OK, {"Content-Type": "application/octet-stream"}
+            data_adaptor.layout_to_fbs_matrix(fields), HTTPStatus.OK, {"Content-Type": "application/octet-stream"}
         )
+    except (KeyError, DatasetAccessError) as e:
+        return abort_and_log(HTTPStatus.BAD_REQUEST, str(e), include_exc_info=True)
     except PrepareError:
         return abort_and_log(
             HTTPStatus.NOT_IMPLEMENTED,

server/data_common/data_adaptor.py

@@ -331,10 +331,9 @@ class DataAdaptor(metaclass=ABCMeta):
         normalized_layout = normalized_layout.astype(dtype=np.float32)
         return normalized_layout
 
-    def layout_to_fbs_matrix(self):
-        """ same as layout, except returns a flatbuffer """
+    def layout_to_fbs_matrix(self, fields):
         """
-        return all embeddings as a flatbuffer, using the cellxgene matrix fbs encoding.
+        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.
@@ -343,8 +342,7 @@ class DataAdaptor(metaclass=ABCMeta):
         * does not support filtering
 
         """
-
-        embeddings = self.get_embedding_names()
+        embeddings = self.get_embedding_names() if fields is None or len(fields) == 0 else fields
         layout_data = []
         with ServerTiming.time(f"layout.query"):
             for ename in embeddings:

server/test/schema.json

@@ -67,6 +67,16 @@
         "name": "umap",
         "type": "float32",
         "dims": ["umap_0", "umap_1"]
+      },
+      {
+        "name": "tsne",
+        "type": "float32",
+        "dims": ["tsne_0", "tsne_1"]
+      },
+      {
+        "name": "pca",
+        "type": "float32",
+        "dims": ["pca_0", "pca_1"]
       }
     ]
   }

server/test/test_anndata_adaptor.py

@@ -34,7 +34,7 @@ Test the anndata adaptor using the pbmc3k data set.
 class AdaptorTest(unittest.TestCase):
     def setUp(self):
         args = {
-            "embeddings__names": ["umap"],
+            "embeddings__names": ["umap", "tsne", "pca"],
             "presentation__max_categories": 100,
             "single_dataset__obs_names": None,
             "single_dataset__var_names": None,
@@ -122,9 +122,9 @@ class AdaptorTest(unittest.TestCase):
         check_feature("PUT", "/annotations/obs", False)
 
     def test_layout(self):
-        fbs = self.data.layout_to_fbs_matrix()
+        fbs = self.data.layout_to_fbs_matrix(fields=None)
         layout = decode_fbs.decode_matrix_FBS(fbs)
-        self.assertEqual(layout["n_cols"], 2)
+        self.assertEqual(layout["n_cols"], 6)
         self.assertEqual(layout["n_rows"], 2638)
 
         X = layout["columns"][0]
@@ -132,6 +132,20 @@ class AdaptorTest(unittest.TestCase):
         Y = layout["columns"][1]
         self.assertTrue((Y >= 0).all() and (Y <= 1).all())
 
+    def test_layout_fields(self):
+        """ X_pca, X_tsne, X_umap are available """
+        fbs = self.data.layout_to_fbs_matrix(["pca"])
+        layout = decode_fbs.decode_matrix_FBS(fbs)
+        self.assertEqual(layout["n_cols"], 2)
+        self.assertEqual(layout["n_rows"], 2638)
+        self.assertCountEqual(layout["col_idx"], ["pca_0", "pca_1"])
+
+        fbs = self.data.layout_to_fbs_matrix(["tsne", "pca"])
+        layout = decode_fbs.decode_matrix_FBS(fbs)
+        self.assertEqual(layout["n_cols"], 4)
+        self.assertEqual(layout["n_rows"], 2638)
+        self.assertCountEqual(layout["col_idx"], ["tsne_0", "tsne_1", "pca_0", "pca_1"])
+
     def test_annotations(self):
         fbs = self.data.annotation_to_fbs_matrix("obs")
         annotations = decode_fbs.decode_matrix_FBS(fbs)