Read AnnData's nullable-encoded columns and v3 categoricals

.changeset/nullable-anndata-columns.md

@@ -0,0 +1,30 @@
+---
+'@spatialdata/core': minor
+---
+
+Read AnnData's nullable-encoded columns, and report their kind.
+
+A nullable column is a **group** of `values` + `mask`, not an array, so opening its
+path as an array fails outright. The visible symptom is usually a missing *index*
+rather than a missing value, because `obs/_index` and `var/_index` are ordinary
+columns — a table would load with `varN` in place of gene names, or with no row ids.
+This is not a legacy shape to tolerate: AnnData 0.13 defaults to zarr v3 and writes
+string columns this way by default, `_index` included, so it is what a freshly
+written `spatialdata` store looks like.
+
+All three nullable encodings are read (`nullable-string-array`, `nullable-integer`,
+`nullable-boolean`), with the mask honoured — a masked entry decodes as `null`, which
+the missing-value handling already treats as absent, so it stays distinguishable from
+a real `0` or `''` rather than being silently rendered as one.
+
+`getObsColumnKinds` recognises the same three. Without this the kind lookup fell
+through to array metadata that a group does not have and returned `undefined`,
+sending `'auto'` mode back to sniffing decoded values for exactly the columns AnnData
+now writes by default — the case that lookup exists to avoid.
+
+Also fixes zarr v3 categoricals decoding to their raw integer codes. The categories
+array is written as `string` on v3, and the text check tested for one v2 spelling of
+that dtype (`v2:object`), so the column resolved to codes — plausible-looking numbers
+rather than an error. The check now asks zarrita whether the dtype is text, which
+covers v3 `string` and v2 `U`/`S` alike, and pandas' `-1` missing code maps to `null`
+instead of indexing off the end of the categories.

packages/core/src/models/VAnnDataSource.ts

@@ -1,7 +1,9 @@
+import * as zarr from 'zarrita';
 import { get as zarrGet, open as zarrOpen } from 'zarrita';
 import type { TableColumnData } from '../types';
 import type { DataSourceParams } from '../Vutils';
 import { dirname } from '../Vutils';
+import { isNullableEncoding, readNullableArray } from './nullableArrays';
 import ZarrDataSource from './VZarrDataSource';
 
 function prependSlash(path: string) {
@@ -93,20 +95,20 @@ export default class AnnDataSource extends ZarrDataSource {
     let categoriesValues: string[] | undefined;
     let codesPath: string | undefined;
     if (categories) {
-      const { dtype } = await zarrOpen(storeRoot.resolve(`${prefix}/${categories}`), {
-        kind: 'array',
-      });
-      if (dtype === 'v2:object') {
-        categoriesValues = await this.getFlatArrDecompressed(`${prefix}/${categories}`);
+      const categoriesPath = `${prefix}/${categories}`;
+      if (await this.hasTextValues(categoriesPath)) {
+        categoriesValues = await this.getFlatArrDecompressed(categoriesPath);
       }
     } else if (encodingType === 'categorical') {
-      const { dtype } = await zarrOpen(storeRoot.resolve(`${path}/categories`), { kind: 'array' });
-      if (dtype === 'v2:object') {
-        categoriesValues = await this.getFlatArrDecompressed(`${path}/categories`);
+      const categoriesPath = `${path}/categories`;
+      if (await this.hasTextValues(categoriesPath)) {
+        categoriesValues = await this.getFlatArrDecompressed(categoriesPath);
       }
       codesPath = `${path}/codes`;
     } else if (encodingType === 'string-array') {
       return this.getFlatArrDecompressed(path);
+    } else if (isNullableEncoding(encodingType)) {
+      return readNullableArray(storeRoot.resolve(path));
     } else {
       const { dtype } = await zarrOpen(storeRoot.resolve(path), { kind: 'array' });
       if (dtype === 'v2:object') {
@@ -119,7 +121,24 @@ export default class AnnDataSource extends ZarrDataSource {
     if (!categoriesValues) {
       return data as TableColumnData;
     }
-    return Array.from(data, (i) => categoriesValues[i as number]);
+    // Pandas encodes a missing categorical as code -1, which indexes nothing.
+    return Array.from(data, (code) => {
+      const index = Number(code);
+      return index < 0 ? null : categoriesValues[index];
+    });
+  }
+
+  /**
+   * Whether the array at *path* holds text, and so needs decoding to strings.
+   *
+   * Covers zarr v3 `string` arrays as well as v2's `object` and fixed-width
+   * unicode. Testing for `v2:object` alone silently leaves categorical columns
+   * resolving to their raw integer codes, which look like plausible data rather
+   * than an error.
+   */
+  private async hasTextValues(path: string): Promise<boolean> {
+    const arr = await zarrOpen(this.storeRoot.resolve(path), { kind: 'array' });
+    return arr.is('string') || arr.is('object');
   }
 
   /**
@@ -168,9 +187,21 @@ export default class AnnDataSource extends ZarrDataSource {
    */
   async getFlatArrDecompressed(path: string) {
     const { storeRoot } = this;
-    const arr = await zarrOpen(storeRoot.resolve(path), { kind: 'array' });
+    const location = storeRoot.resolve(path);
+    // A nullable column is a group, so opening it as an array throws. Resolve the
+    // node first rather than assuming, since index paths reach here directly.
+    const node = await zarrOpen(location);
+    if (node instanceof zarr.Group) {
+      if (!isNullableEncoding(node.attrs['encoding-type'])) {
+        throw new Error(
+          `Expected an array at ${path}, but found a group encoded as ` +
+            `${String(node.attrs['encoding-type'] ?? 'unknown')}.`
+        );
+      }
+      return (await readNullableArray(location)) as string[];
+    }
     // Zarrita supports decoding vlen-utf8-encoded string arrays.
-    const data = await zarrGet(arr);
+    const data = await zarrGet(node);
     if (data.data?.[Symbol.iterator]) {
       return Array.from(data.data) as string[];
     }

packages/core/src/models/index.ts

@@ -29,6 +29,7 @@ import type {
   ZarrTree,
 } from '../types';
 import { ATTRS_KEY, Err, Ok, ZARRAY_KEY } from '../types';
+import { NULLABLE_ENCODING_KINDS } from './nullableArrays';
 import SpatialDataPointsSource from './VPointsSource';
 import SpatialDataShapesSource from './VShapesSource';
 import SpatialDataTableSource from './VTableSource';
@@ -211,10 +212,20 @@ export type TableKeys = {
   instanceKey: string;
 };
 
-/** AnnData `encoding-type` values that settle the kind on their own. */
+/**
+ * AnnData `encoding-type` values that settle the kind on their own.
+ *
+ * The nullable encodings have to be listed here rather than left to the dtype
+ * branch below: they are groups of `values` + `mask`, so there is no array
+ * metadata on the node itself to fall back to. Leaving them out is not a
+ * missing edge case — AnnData 0.13 writes string columns this way by default on
+ * zarr v3, `obs/_index` included, so the columns that most need a declared kind
+ * are exactly the ones that would arrive without one.
+ */
 const OBS_KIND_BY_ENCODING: Record<string, TableColumnKind> = {
   categorical: 'categorical',
   'string-array': 'string',
+  ...NULLABLE_ENCODING_KINDS,
 };
 
 /**
@@ -415,9 +426,8 @@ export class TableElement extends AbstractElement<'tables'> {
    * `undefined` for a column that is absent or whose node we do not recognise.
    */
   getObsColumnKinds(columnNames: string[]): Array<TableColumnKind | undefined> {
-    const node = this.parsed as ZarrTree;
-    const obsNode = node.obs as ZarrTree | undefined;
-    if (!obsNode || typeof obsNode !== 'object') {
+    const obsNode = this.getObsGroup();
+    if (!obsNode) {
       return columnNames.map(() => undefined);
     }
     return columnNames.map((columnName) => classifyObsColumnNode(obsNode[columnName]));

packages/core/src/models/nullableArrays.ts

@@ -0,0 +1,71 @@
+import { type Location, type Readable, get as zarrGet, open as zarrOpen } from 'zarrita';
+import type { TableColumnKind, TableValue } from '../types';
+
+/**
+ * AnnData's nullable encodings, which store a column as a *group* rather than an
+ * array: a `values` array plus a boolean `mask` marking the null positions.
+ *
+ * Readers that open a column path as an array fail outright on these — the
+ * symptom is a missing index rather than a missing value, because
+ * `obs/_index` and `var/_index` are themselves ordinary columns. AnnData writes
+ * this layout by default from 0.13 onwards (and `spatialdata` inherits it), so
+ * it turns up in freshly written stores, not only in ones that have been
+ * rewritten.
+ *
+ * See the AnnData on-disk specification, "Nullable integers, booleans and
+ * strings" — all three share this group layout at encoding-version 0.1.0.
+ *
+ * The mapped value is the kind of the column's `values`, which is what the
+ * encoding name already tells us — the mask changes which entries are present,
+ * never what type they are. Declared here rather than at the classifier so the
+ * three encoding names are listed once: reading them and classifying them have
+ * to stay in step, and they live in different modules.
+ */
+export const NULLABLE_ENCODING_KINDS: Record<string, TableColumnKind> = {
+  'nullable-string-array': 'string',
+  'nullable-integer': 'numeric',
+  'nullable-boolean': 'boolean',
+};
+
+export const NULLABLE_ENCODING_TYPES = new Set(Object.keys(NULLABLE_ENCODING_KINDS));
+
+export function isNullableEncoding(encodingType: unknown): boolean {
+  return typeof encodingType === 'string' && NULLABLE_ENCODING_TYPES.has(encodingType);
+}
+
+/**
+ * Read a nullable-encoded column into a flat array, with `null` at masked
+ * positions.
+ *
+ * The mask is read alongside the values rather than ignored: a masked entry is
+ * absent, not empty, and callers that render it (tooltips, legends) need to be
+ * able to tell those apart.
+ */
+export async function readNullableArray(location: Location<Readable>): Promise<TableValue[]> {
+  const valuesArray = await zarrOpen(location.resolve('values'), { kind: 'array' });
+  const values = await zarrGet(valuesArray);
+
+  // The specification requires a mask, but a values-only group is still
+  // unambiguous, so read what is there rather than failing on a missing mask.
+  let mask: boolean[] | undefined;
+  try {
+    const maskArray = await zarrOpen(location.resolve('mask'), { kind: 'array' });
+    mask = toArray((await zarrGet(maskArray)).data).map(Boolean);
+  } catch {
+    mask = undefined;
+  }
+
+  return toArray(values.data).map((value, index) => (mask?.[index] ? null : value));
+}
+
+/**
+ * Materialise a zarrita chunk as a plain array.
+ *
+ * Zarrita returns several backing types — native TypedArrays plus its own
+ * `BoolArray`/`ByteStringArray`/`UnicodeStringArray`. They are all iterable but
+ * do not share an indexable interface, so the union has no common supertype to
+ * narrow to; iterating is the one operation defined across all of them.
+ */
+function toArray(data: unknown): TableValue[] {
+  return Array.from(data as Iterable<TableValue>);
+}