Speed up JSON mappings (#7706)

* Use Map instead of Object for JSON mappings. Also, no longer maintain extra mappingKeys array.
* update changelog
* Keep frontend api backwards compatible
* Merge branch 'master' into use-map-for-mappings

Author: daniel-wer

CHANGELOG.unreleased.md

@@ -20,7 +20,7 @@ For upgrade instructions, please check the [migration guide](MIGRATIONS.released
 - If storage scan is enabled, the measured used storage is now displayed in the dashboard’s dataset detail view. [#7677](https://github.com/scalableminds/webknossos/pull/7677)
 - Prepared support to download full stl meshes via the HTTP api. [#7587](https://github.com/scalableminds/webknossos/pull/7587)
 - You can now place segment index files with your on-disk segmentation layers, which makes segment stats available when viewing these segmentations, and also when working on volume annotations based on these segmentation layers. [#7437](https://github.com/scalableminds/webknossos/pull/7437)
-- Added an action to delete erronous, unimported datasets directly from the dashboard. [#7448](https://github.com/scalableminds/webknossos/pull/7448)
+- Added an action to delete erroneous, unimported datasets directly from the dashboard. [#7448](https://github.com/scalableminds/webknossos/pull/7448)
 - Added support for `window`, `active`, `inverted` keys from the `omero` info in the NGFF metadata. [7685](https://github.com/scalableminds/webknossos/pull/7685)
 - Added getSegment function to JavaScript API. Also, createSegmentGroup returns the id of the new group now. [#7694](https://github.com/scalableminds/webknossos/pull/7694)
 - Added support for importing N5 datasets without multiscales metadata. [#7664](https://github.com/scalableminds/webknossos/pull/7664)
@@ -37,6 +37,7 @@ For upgrade instructions, please check the [migration guide](MIGRATIONS.released
 - If the current dataset folder in the dashboard cannot be found (e.g., because somebody else deleted it), the page navigates to the root folder automatically. [#7669](https://github.com/scalableminds/webknossos/pull/7669)
 - Voxelytics logs are now stored by organization name, rather than id, in Loki. This is in preparation of the unification of these two concepts. [#7687](https://github.com/scalableminds/webknossos/pull/7687)
 - Using a segment index file with a different data type than uint16 will now result in an error. [#7698](https://github.com/scalableminds/webknossos/pull/7698)
+- Improved performance of JSON mappings in preparation of frontend HDF5 mappings. [#7706](https://github.com/scalableminds/webknossos/pull/7706)
 
 ### Fixed
 - Fixed rare SIGBUS crashes of the datastore module that were caused by memory mapping on unstable file systems. [#7528](https://github.com/scalableminds/webknossos/pull/7528)

frontend/javascripts/oxalis/api/api_latest.ts

@@ -1497,7 +1497,7 @@ class DataApi {
    */
   setMapping(
     layerName: string,
-    mapping: Mapping,
+    mapping: Mapping | Record<number, number>,
     options: {
       colors?: Array<number>;
       hideUnmappedIds?: boolean;
@@ -1517,10 +1517,10 @@ class DataApi {
       sendAnalyticsEvent("setMapping called with custom colors");
     }
     const mappingProperties = {
-      mapping: _.clone(mapping),
-      // Object.keys is sorted for numerical keys according to the spec:
-      // http://www.ecma-international.org/ecma-262/6.0/#sec-ordinary-object-internal-methods-and-internal-slots-ownpropertykeys
-      mappingKeys: Object.keys(mapping).map((x) => parseInt(x, 10)),
+      mapping:
+        mapping instanceof Map
+          ? new Map(mapping)
+          : new Map(Object.entries(mapping).map(([key, value]) => [parseInt(key, 10), value])),
       mappingColors,
       hideUnmappedIds,
       showLoadingIndicator,

frontend/javascripts/oxalis/api/api_v2.ts

@@ -555,7 +555,7 @@ class DataApi {
    *
    * api.setMapping("segmentation", mapping);
    */
-  setMapping(layerName: string, mapping: Mapping) {
+  setMapping(layerName: string, mapping: Mapping | Record<number, number>) {
     const segmentationLayer = this.model.getLayerByName(layerName);
     const segmentationLayerName = segmentationLayer != null ? segmentationLayer.name : null;
 
@@ -564,8 +564,10 @@ class DataApi {
     }
 
     const mappingProperties = {
-      mapping: _.clone(mapping),
-      mappingKeys: Object.keys(mapping).map((x) => parseInt(x, 10)),
+      mapping:
+        mapping instanceof Map
+          ? new Map(mapping)
+          : new Map(Object.entries(mapping).map(([key, value]) => [parseInt(key, 10), value])),
     };
     Store.dispatch(setMappingAction(layerName, "<custom mapping>", "JSON", mappingProperties));
   }

frontend/javascripts/oxalis/merger_mode.ts

@@ -26,7 +26,7 @@ import { type AdditionalCoordinate } from "types/api_flow_types";
 
 type MergerModeState = {
   treeIdToRepresentativeSegmentId: Record<number, number | null | undefined>;
-  idMapping: Record<number, number>;
+  idMapping: Map<number, number>;
   nodesPerSegment: Record<number, number>;
   nodes: Array<NodeWithTreeId>;
   // A properly initialized merger mode should always
@@ -48,7 +48,7 @@ function mapSegmentToRepresentative(
   mergerModeState: MergerModeState,
 ) {
   const representative = getRepresentativeForTree(treeId, segId, mergerModeState);
-  mergerModeState.idMapping[segId] = representative;
+  mergerModeState.idMapping.set(segId, representative);
 }
 
 function getRepresentativeForTree(treeId: number, segId: number, mergerModeState: MergerModeState) {
@@ -66,7 +66,7 @@ function getRepresentativeForTree(treeId: number, segId: number, mergerModeState
 
 function deleteIdMappingOfSegment(segId: number, treeId: number, mergerModeState: MergerModeState) {
   // Remove segment from color mapping
-  delete mergerModeState.idMapping[segId];
+  mergerModeState.idMapping.delete(segId);
   delete mergerModeState.treeIdToRepresentativeSegmentId[treeId];
 }
 
@@ -420,9 +420,9 @@ function resetState(mergerModeState: Partial<MergerModeState> = {}) {
   const state = Store.getState();
   const visibleLayer = getVisibleSegmentationLayer(state);
   const segmentationLayerName = visibleLayer != null ? visibleLayer.name : null;
-  const defaults = {
+  const defaults: MergerModeState = {
     treeIdToRepresentativeSegmentId: {},
-    idMapping: {},
+    idMapping: new Map(),
     nodesPerSegment: {},
     nodes: getAllNodesWithTreeId(),
     segmentationLayerName,

frontend/javascripts/oxalis/model/accessors/dataset_accessor.ts

@@ -655,7 +655,6 @@ export function is2dDataset(dataset: APIDataset): boolean {
 const dummyMapping = {
   mappingName: null,
   mapping: null,
-  mappingKeys: null,
   mappingColors: null,
   hideUnmappedIds: false,
   mappingStatus: MappingStatusEnum.DISABLED,

frontend/javascripts/oxalis/model/actions/settings_actions.ts

@@ -174,7 +174,6 @@ export const setMappingEnabledAction = (layerName: string, isMappingEnabled: boo
 
 export type OptionalMappingProperties = {
   mapping?: Mapping;
-  mappingKeys?: Array<number>;
   mappingColors?: Array<number>;
   hideUnmappedIds?: boolean;
   showLoadingIndicator?: boolean;
@@ -183,21 +182,14 @@ export const setMappingAction = (
   layerName: string,
   mappingName: string | null | undefined,
   mappingType: MappingType = "JSON",
-  {
-    mapping,
-    mappingKeys,
-    mappingColors,
-    hideUnmappedIds,
-    showLoadingIndicator,
-  }: OptionalMappingProperties = {},
+  { mapping, mappingColors, hideUnmappedIds, showLoadingIndicator }: OptionalMappingProperties = {},
 ) =>
   ({
     type: "SET_MAPPING",
     layerName,
     mappingName,
     mappingType,
     mapping,
-    mappingKeys,
     mappingColors,
     hideUnmappedIds,
     showLoadingIndicator,

frontend/javascripts/oxalis/model/bucket_data_handling/data_cube.ts

@@ -175,7 +175,7 @@ class DataCube {
     const mapping = this.getMapping();
 
     if (mapping != null && this.isMappingEnabled()) {
-      mappedId = mapping[idToMap];
+      mappedId = mapping.get(idToMap);
     }
 
     if (this.shouldHideUnmappedIds() && mappedId == null) {
@@ -841,7 +841,7 @@ class DataCube {
       const dataValue = Number(data[voxelIndex]);
 
       if (mapping) {
-        const mappedValue = mapping[dataValue];
+        const mappedValue = mapping.get(dataValue);
 
         if (mappedValue != null) {
           return mappedValue;

frontend/javascripts/oxalis/model/bucket_data_handling/mappings.ts

@@ -50,30 +50,26 @@ class Mappings {
       (state) =>
         getMappingInfo(state.temporaryConfiguration.activeMappingByLayer, this.layerName).mapping,
       (mapping) => {
-        const { mappingKeys } = getMappingInfo(
-          Store.getState().temporaryConfiguration.activeMappingByLayer,
-          this.layerName,
-        );
-        this.updateMappingTextures(mapping, mappingKeys);
+        this.updateMappingTextures(mapping);
       },
       true,
     );
   }
 
-  async updateMappingTextures(
-    mapping: Mapping | null | undefined,
-    mappingKeys: Array<number> | null | undefined,
-  ): Promise<void> {
-    if (mapping == null || mappingKeys == null) return;
+  async updateMappingTextures(mapping: Mapping | null | undefined): Promise<void> {
+    if (mapping == null) return;
     console.time("Time to create mapping texture");
-    const mappingSize = mappingKeys.length;
+    const mappingSize = mapping.size;
     // The typed arrays need to be padded with 0s so that their length is a multiple of MAPPING_TEXTURE_WIDTH
     const paddedLength =
       mappingSize + MAPPING_TEXTURE_WIDTH - (mappingSize % MAPPING_TEXTURE_WIDTH);
     const keys = new Uint32Array(paddedLength);
     const values = new Uint32Array(paddedLength);
+    const mappingKeys = Array.from(mapping.keys());
+    mappingKeys.sort((a, b) => a - b);
     keys.set(mappingKeys);
-    values.set(mappingKeys.map((key) => mapping[key]));
+    // @ts-ignore mappingKeys are guaranteed to exist in mapping as they are mapping.keys()
+    values.set(mappingKeys.map((key) => mapping.get(key)));
     // Instantiate the Uint8Arrays with the array buffer from the Uint32Arrays, so that each 32-bit value is converted
     // to four 8-bit values correctly
     const uint8Keys = new Uint8Array(keys.buffer);

frontend/javascripts/oxalis/model/reducers/dataset_reducer.ts

@@ -43,7 +43,6 @@ function DatasetReducer(state: OxalisState, action: Action): OxalisState {
           activeMappingByLayer: createDictWithKeysAndValue(segmentationLayerNames, () => ({
             mappingName: null,
             mapping: null,
-            mappingKeys: null,
             mappingColors: null,
             hideUnmappedIds: false,
             mappingStatus: MappingStatusEnum.DISABLED,

frontend/javascripts/oxalis/model/reducers/settings_reducer.ts

@@ -250,7 +250,7 @@ function SettingsReducer(state: OxalisState, action: Action): OxalisState {
     }
 
     case "SET_MAPPING": {
-      const { mappingName, mapping, mappingKeys, mappingColors, mappingType, layerName } = action;
+      const { mappingName, mapping, mappingColors, mappingType, layerName } = action;
 
       // Editable mappings cannot be disabled or switched for now
       if (!isMappingActivationAllowed(state, mappingName, layerName)) return state;
@@ -265,10 +265,9 @@ function SettingsReducer(state: OxalisState, action: Action): OxalisState {
         {
           mappingName,
           mapping,
-          mappingKeys,
           mappingColors,
           mappingType,
-          mappingSize: mappingKeys != null ? mappingKeys.length : 0,
+          mappingSize: mapping != null ? mapping.size : 0,
           hideUnmappedIds,
           mappingStatus:
             mappingName != null ? MappingStatusEnum.ACTIVATING : MappingStatusEnum.DISABLED,

frontend/javascripts/oxalis/model/sagas/mapping_saga.ts

@@ -229,14 +229,9 @@ function* handleSetMapping(
   const fetchedMapping = fetchedMappings[mappingName];
   const { hideUnmappedIds, colors: mappingColors } = fetchedMapping;
 
-  const [mappingObject, mappingKeys] = yield* call(
-    buildMappingObject,
-    mappingName,
-    fetchedMappings,
-  );
+  const mapping = yield* call(buildMappingObject, mappingName, fetchedMappings);
   const mappingProperties = {
-    mapping: mappingObject,
-    mappingKeys,
+    mapping,
     mappingColors,
     hideUnmappedIds,
   };
@@ -259,9 +254,9 @@ function* handleSetMapping(
 
 function convertMappingObjectToClasses(existingMapping: Mapping) {
   const classesByRepresentative: Record<number, number[]> = {};
-  for (const unmappedStr of Object.keys(existingMapping)) {
-    const unmapped = Number(unmappedStr);
-    const mapped = existingMapping[unmapped];
+  for (const unmapped of existingMapping.keys()) {
+    // @ts-ignore unmapped is guaranteed to exist in existingMapping as it was obtained using existingMapping.keys()
+    const mapped: number = existingMapping.get(unmapped);
     classesByRepresentative[mapped] = classesByRepresentative[mapped] || [];
     classesByRepresentative[mapped].push(unmapped);
   }
@@ -280,10 +275,10 @@ function* setCustomColors(
   let classIdx = 0;
   for (const aClass of classes) {
     const firstIdEntry = aClass[0];
-    if (firstIdEntry == null) {
-      continue;
-    }
-    const representativeId = mappingProperties.mapping[firstIdEntry];
+    if (firstIdEntry == null) continue;
+
+    const representativeId = mappingProperties.mapping.get(firstIdEntry);
+    if (representativeId == null) continue;
 
     const hueValue = mappingProperties.mappingColors[classIdx];
     const color = jsHsv2rgb(360 * hueValue, 1, 1);
@@ -319,14 +314,8 @@ function* fetchMappings(
   }
 }
 
-function buildMappingObject(
-  mappingName: string,
-  fetchedMappings: APIMappings,
-): [Mapping, Array<number>] {
-  const mappingObject: Mapping = {};
-  // Performance optimization: Object.keys(...) is slow for large objects
-  // keeping track of the keys in a separate array is ~5x faster
-  const mappingKeys = [];
+function buildMappingObject(mappingName: string, fetchedMappings: APIMappings): Mapping {
+  const mappingObject: Mapping = new Map();
 
   for (const currentMappingName of getMappingChain(mappingName, fetchedMappings)) {
     const mapping = fetchedMappings[currentMappingName];
@@ -341,17 +330,15 @@ function buildMappingObject(
         // The class is empty and can be ignored
         continue;
       }
-      const mappedId = mappingObject[minId] || minId;
+      const mappedId = mappingObject.get(minId) || minId;
 
       for (const id of mappingClass) {
-        mappingObject[id] = mappedId;
-        mappingKeys.push(id);
+        mappingObject.set(id, mappedId);
       }
     }
   }
 
-  mappingKeys.sort((a, b) => a - b);
-  return [mappingObject, mappingKeys];
+  return mappingObject;
 }
 
 function getMappingChain(mappingName: string, fetchedMappings: APIMappings): Array<string> {

frontend/javascripts/oxalis/store.ts

@@ -389,12 +389,11 @@ export type RecommendedConfiguration = Partial<
 // A histogram value of undefined indicates that the histogram hasn't been fetched yet
 // whereas a value of null indicates that the histogram couldn't be fetched
 export type HistogramDataForAllLayers = Record<string, APIHistogramData | null>;
-export type Mapping = Record<number, number>;
+export type Mapping = Map<number, number>;
 export type MappingType = "JSON" | "HDF5";
 export type ActiveMappingInfo = {
   readonly mappingName: string | null | undefined;
   readonly mapping: Mapping | null | undefined;
-  readonly mappingKeys: number[] | null | undefined;
   readonly mappingColors: number[] | null | undefined;
   readonly hideUnmappedIds: boolean;
   readonly mappingStatus: MappingStatus;

frontend/javascripts/test/model/binary/cube.spec.ts

@@ -200,8 +200,8 @@ test("getDataValue() should return the mapping value if available", async (t) =>
   const { cube } = t.context;
   await cube._labelVoxelInResolution_DEPRECATED([0, 0, 0], null, 42, 0, null);
   await cube._labelVoxelInResolution_DEPRECATED([1, 1, 1], null, 43, 0, null);
-  const mapping = [];
-  mapping[42] = 1;
+  const mapping = new Map();
+  mapping.set(42, 1);
   t.is(cube.getDataValue([0, 0, 0], null, mapping), 1);
   t.is(cube.getDataValue([1, 1, 1], null, mapping), 43);
 });

frontend/javascripts/test/sagas/volumetracing/volumetracing_saga.spec.ts

@@ -120,8 +120,7 @@ const initialState = update(defaultState, {
 
 const dummyActiveMapping: ActiveMappingInfo = {
   mappingName: "dummy-mapping-name",
-  mapping: {},
-  mappingKeys: [],
+  mapping: new Map(),
   mappingColors: [],
   hideUnmappedIds: false,
   mappingStatus: "ENABLED",