From 83d572cde25d1ce8468329b3b4078ca130b9c554 Mon Sep 17 00:00:00 2001
From: Bruce Martin <bruce@chanzuckerberg.com>
Date: Wed, 22 Jul 2020 17:02:19 -0700
Subject: [PATCH] force categorical treatment of user annotations (#1663)

---
 client/src/actions/annotation.js              | 29 ++++++++++++-------
 client/src/annoMatrix/loader.js               | 14 ++++-----
 client/src/annoMatrix/schema.js               |  8 ++++-
 .../src/components/continuous/continuous.js   |  1 +
 .../src/util/stateManager/controlsHelpers.js  |  7 ++++-
 5 files changed, 39 insertions(+), 20 deletions(-)

diff --git a/client/src/actions/annotation.js b/client/src/actions/annotation.js
index de89f1d1..8a7a154b 100644
--- a/client/src/actions/annotation.js
+++ b/client/src/actions/annotation.js
@@ -34,7 +34,8 @@ export const annotationCreateCategoryAction = (
     throw new Error("name collision on annotation category create");
 
   let initialValue;
-  let categories;
+  let newSchema;
+  let ctor;
   if (categoryToDuplicate) {
     /* if we are duplicating a category, retrieve it */
     const catDupSchema = schema.annotations.obsByName[categoryToDuplicate];
@@ -47,25 +48,33 @@ export const annotationCreateCategoryAction = (
       .fetch("obs", categoryToDuplicate);
     const col = catToDupDf.col(categoryToDuplicate);
     initialValue = col.asArray();
-    ({ categories } = col.summarize());
+    const { categories } = col.summarizeCategorical();
     // all user-created annotations must have the unassigned category
     if (!categories.includes(globals.unassignedCategoryLabel)) {
       categories.push(globals.unassignedCategoryLabel);
     }
+    ctor = initialValue.constructor;
+    newSchema = {
+      ...catDupSchema,
+      name: newCategoryName,
+      categories,
+      writable: true,
+    };
   } else {
     /* else assign to the standard default value */
     initialValue = globals.unassignedCategoryLabel;
-    categories = [globals.unassignedCategoryLabel];
+    ctor = Array;
+    newSchema = {
+      name: newCategoryName,
+      type: "categorical",
+      categories: [globals.unassignedCategoryLabel],
+      writable: true,
+    };
   }
 
   const obsCrossfilter = prevObsCrossfilter.addObsColumn(
-    {
-      name: newCategoryName,
-      type: "categorical",
-      categories,
-      writable: true,
-    },
-    Array,
+    newSchema,
+    ctor,
     initialValue
   );
 
diff --git a/client/src/annoMatrix/loader.js b/client/src/annoMatrix/loader.js
index 7ee27553..e10e3ee1 100644
--- a/client/src/annoMatrix/loader.js
+++ b/client/src/annoMatrix/loader.js
@@ -90,10 +90,10 @@ export default class AnnoMatrixLoader extends AnnoMatrix {
     If an array, it must be of same size as nObs and same type as Ctor
 		*/
     colSchema.writable = true;
-    const col = colSchema.name;
+    const colName = colSchema.name;
     if (
-      _getColumnSchema(this.schema, "obs", col) ||
-      this._cache.obs.hasCol(col)
+      _getColumnSchema(this.schema, "obs", colName) ||
+      this._cache.obs.hasCol(colName)
     ) {
       throw new Error("column already exists");
     }
@@ -109,11 +109,9 @@ export default class AnnoMatrixLoader extends AnnoMatrix {
     } else {
       data = new Ctor(this.nObs).fill(value);
     }
-    o._cache.obs = this._cache.obs.withCol(col, data);
-    o.schema = addObsAnnoColumn(this.schema, col, {
-      ...colSchema,
-      writable: true,
-    });
+    o._cache.obs = this._cache.obs.withCol(colName, data);
+    _normalizeCategoricalSchema(colSchema, o._cache.obs.col(colName));
+    o.schema = addObsAnnoColumn(this.schema, colName, colSchema);
     return o;
   }
 
diff --git a/client/src/annoMatrix/schema.js b/client/src/annoMatrix/schema.js
index 1c66febe..29c7e959 100644
--- a/client/src/annoMatrix/schema.js
+++ b/client/src/annoMatrix/schema.js
@@ -66,7 +66,12 @@ export function _isContinuousType(schema) {
 
 export function _normalizeCategoricalSchema(colSchema, col) {
   const { type, writable } = colSchema;
-  if (type === "string" || type === "boolean" || type === "categorical") {
+  if (
+    type === "string" ||
+    type === "boolean" ||
+    type === "categorical" ||
+    writable
+  ) {
     const categorySet = new Set(
       col.summarizeCategorical().categories.concat(colSchema.categories ?? [])
     );
@@ -79,4 +84,5 @@ export function _normalizeCategoricalSchema(colSchema, col) {
   if (colSchema.categories) {
     colSchema.categories = catLabelSort(writable, colSchema.categories);
   }
+  return colSchema;
 }
diff --git a/client/src/components/continuous/continuous.js b/client/src/components/continuous/continuous.js
index 3752aee6..ad5c908d 100644
--- a/client/src/components/continuous/continuous.js
+++ b/client/src/components/continuous/continuous.js
@@ -16,6 +16,7 @@ class Continuous extends React.PureComponent {
     const allContinuousNames = schema.annotations.obs.columns
       .filter((col) => col.type === "int32" || col.type === "float32")
       .filter((col) => col.name !== obsIndex)
+      .filter((col) => !col.writable) // skip user annotations - they will be treated as categorical
       .map((col) => col.name);
 
     return (
diff --git a/client/src/util/stateManager/controlsHelpers.js b/client/src/util/stateManager/controlsHelpers.js
index a0091138..e2aa8379 100644
--- a/client/src/util/stateManager/controlsHelpers.js
+++ b/client/src/util/stateManager/controlsHelpers.js
@@ -64,7 +64,12 @@ function topNCategories(colSchema, summary, N) {
 
 export function isSelectableCategoryName(schema, name) {
   const { index } = schema.annotations.obs;
-  return name && name !== index && isCategoricalAnnotation(schema, name);
+  const colSchema = schema.annotations.obsByName[name];
+  return (
+    name &&
+    name !== index &&
+    (isCategoricalAnnotation(schema, name) || colSchema.writable)
+  );
 }
 
 export function selectableCategoryNames(schema, names) {