Fancier caching of selection get (#195)

We've seen some performance issue when calling selection.get very often; this should speed it up a bit.

We already cache the tile bitmasks, but currently for each Selection.get we have to identify where the nth TRUE element is for one tile. This is currently handled by getting in each individual position. This introduces two improvements:

it calls bitmask.which(), which returns an int16 array with locations for all the set bits. Ordinarily we try to avoid materializing this because for a fully set mask it's 16x larger than the bitmask itself. But for a nearly empty bitmask it iterates much faster, because it works a byte at a time instead of a bit at a time.
We cache the output of bitmask.which(), so we don't have to run it repeatedly at all. The 1000 fifo cache is global for deepscatter: assuming 50,000 items per bitmask and 50% filled bitmasks, that could be 100MB of CPU data in this cache. That ain't nothing! I might also make the cache expire after 1-10 seconds to avoid gunking things up too much.






<!-- ELLIPSIS_HIDDEN -->


----

> [!IMPORTANT]
> Improves `Selection.get` performance by caching bitmask indices using a new `FifoTupleMap` class.
> 
>   - **Caching**:
>     - Introduces `FifoTupleMap` in `utilityFunctions.ts` for FIFO caching.
>     - Caches `bitmask.which()` results in `cachedWhiches` in `selection.ts` to improve `Selection.get` performance.
>   - **Behavior**:
>     - `Selection.get` in `selection.ts` now uses cached bitmask indices, reducing repeated calculations.
>   - **Misc**:
>     - Adds `check` script to `package.json` for building and testing.
> 
> <sup>This description was created by </sup>[<img alt="Ellipsis" src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=nomic-ai%2Fdeepscatter&utm_source=github&utm_medium=referral)<sup> for 3b055be. You can [customize](https://app.ellipsis.dev/nomic-ai/settings/summaries) this summary. It will automatically update as commits are pushed.</sup>


<!-- ELLIPSIS_HIDDEN -->

Author: bmschmidt

package.json

@@ -18,6 +18,7 @@
   },
   "scripts": {
     "dev": "vite --mode dev --port 3344 --host",
+    "check": "npm run build && npm run test",
     "format": "prettier --write src",
     "prepare": "npm run build",
     "build": "vite build && tsc",

src/selection.ts

@@ -6,7 +6,7 @@ import { getTileFromRow } from './tixrixqid';
 import type * as DS from './types';
 import { Bool, StructRowProxy, Utf8, Vector, makeData } from 'apache-arrow';
 import { bisectLeft, bisectRight, range } from 'd3-array';
-import { MinHeap } from './utilityFunctions';
+import { FifoTupleMap, MinHeap } from './utilityFunctions';
 interface SelectParams {
   name: string;
   useNameCache?: boolean; // If true and a selection with that name already exists, use it and ignore all passed parameters. Otherwise, throw an error.
@@ -382,6 +382,7 @@ class SelectionTile {
  * more than once a second or so.
  */
 
+const cachedWhiches = new FifoTupleMap<SelectionTile, Uint16Array>(250);
 export class DataSelection {
   deeptable: Deeptable;
   plot: Scatterplot;
@@ -818,31 +819,33 @@ export class DataSelection {
       i = this.selectionSize + i;
     }
     let currentOffset = 0;
-    let relevantTile: Tile | undefined = undefined;
-    for (const { tile, matchCount } of this.tiles) {
-      if (i < currentOffset + matchCount) {
-        relevantTile = tile;
+    let relevantTile: SelectionTile | undefined = undefined;
+    for (const selectionTile of this.tiles) {
+      if (i < currentOffset + selectionTile.matchCount) {
+        relevantTile = selectionTile;
         break;
       }
-      currentOffset += matchCount;
+      currentOffset += selectionTile.matchCount;
     }
     if (relevantTile === undefined) {
       return undefined;
     }
-    const column = relevantTile.record_batch.getChild(
-      this.name,
-    ) as Vector<Bool>;
+
     const offset = i - currentOffset;
-    let ix_in_match = 0;
-    for (let j = 0; j < column.length; j++) {
-      if (column.get(j)) {
-        if (ix_in_match === offset) {
-          return relevantTile.record_batch.get(j) || undefined;
-        }
-        ix_in_match++;
-      }
+    let cached = cachedWhiches.get([relevantTile]);
+    if (cached === undefined) {
+      const mask = Bitmask.from_arrow(relevantTile.bitmask);
+      cached = mask.which();
+      cachedWhiches.set([relevantTile], cached)
     }
-    throw new Error(`unable to locate point ${i}`);
+
+    if (offset >= cached.length) {
+      throw new Error(`unable to locate point ${i}`);
+    }
+    const tix = cached[offset];
+    
+    return relevantTile.tile.record_batch.get(tix) || undefined;
+
   }
 
   // Iterate over the points in raw order.

src/utilityFunctions.ts

@@ -107,7 +107,12 @@ export class TupleMap<K = object, V = object> {
     }
   }
 
+  get size(): number {
+    return this.map.size;
+  }
+
   set(keys: Some<K>, value: V): void {
+    // eslint-disable-next-line @typescript-eslint/no-this-alias
     let currentMap: TupleMap<K, V> = this;
     for (const key of keys) {
       if (!currentMap.map.has(key)) {
@@ -329,3 +334,25 @@ export class MinHeap<T> {
     }
   }
 }
+
+
+export class FifoTupleMap<K, V> extends TupleMap<K, V> {
+  // private f: (key: K, arg: ArgT) => V;
+  max: number;
+  constructor(max: number) {
+    super();
+    this.max = max;
+  }
+
+  set(key: Some<K>, value: V) {
+    super.set(key, value);
+    if (this.size > this.max) {
+      const next = this.keys().next();
+      if (next.value) {        
+        // using 'as' here because we lose the type in the iteration.
+        this.delete(next.value as Some<K>);
+      }
+    }
+    return this;
+  }
+}