From 9aedf13c47089a8b60474f45a38212b9aef3f583 Mon Sep 17 00:00:00 2001
From: Benjamin Schmidt <bmschmidt@gmail.com>
Date: Fri, 25 Oct 2024 20:55:17 -0400
Subject: [PATCH] better mimic js array in selection getters (#168)

---
 src/selection.ts      | 25 ++++++++++++++++++++-----
 tests/dataset.spec.js | 15 +++++++++++----
 2 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/src/selection.ts b/src/selection.ts
index 8259683f4..4a2d3f813 100644
--- a/src/selection.ts
+++ b/src/selection.ts
@@ -805,14 +805,17 @@ export class DataSelection {
    * that contains the nth match, then iterate through the matches in that
    * tile until we find the nth match.
    *
-   * @param i the index of the row to get
+   * @param i the index of the row to get. If less than zero, will return
    */
   get(i: number | undefined = undefined): StructRowProxy | undefined {
     if (i === undefined) {
       i = this.cursor;
     }
-    if (i >= this.selectionSize) {
-      throw new Error('Selection out of bounds');
+    if (i >= this.selectionSize || i < -this.selectionSize) {
+      return undefined;
+    }
+    if (i < 0) {
+      i = this.selectionSize + i;
     }
     let currentOffset = 0;
     let relevantTile: Tile | undefined = undefined;
@@ -1103,6 +1106,15 @@ export class SortedDataSelection extends DataSelection {
    * This implementation uses Quickselect with a pivot selected from actual data.
    */
   get(k: number): StructRowProxy | undefined {
+    if (k >= this.selectionSize || k < -this.selectionSize) {
+      return undefined;
+    }
+
+    if (k < 0) {
+      k = this.selectionSize + k;
+    }
+
+    // Implement Quickselect over the combined data
     const actualK = this.order === 'ascending' ? k : this.selectionSize - k - 1;
     // Implement Quickselect over the combined data
     const result = quickSelect(actualK, this.tiles, this.key, true);
@@ -1236,7 +1248,10 @@ export class TileSorter
     // Finally add to heap
     for (const sortInfo of sortInfos) {
       // Only add if there are indices and we haven't reached the end of the values
-      if (sortInfo.indices.length > 0 && sortInfo.pointer < sortInfo.values.length) {
+      if (
+        sortInfo.indices.length > 0 &&
+        sortInfo.pointer < sortInfo.values.length
+      ) {
         this.valueHeap.insert(sortInfo);
       }
     }
@@ -1257,7 +1272,7 @@ export class TileSorter
         const rawSortInfo = tile.sorts[this.sortKey];
         // Skip tiles with empty selections
         if (rawSortInfo.indices.length <= 0) {
-          continue
+          continue;
         }
         const sortInfo: SortInfoWithPointer = {
           ...rawSortInfo,
diff --git a/tests/dataset.spec.js b/tests/dataset.spec.js
index f0fffaba3..f6979a1a9 100644
--- a/tests/dataset.spec.js
+++ b/tests/dataset.spec.js
@@ -95,6 +95,7 @@ test('Test composition of selections', async () => {
 
   await selectNothing.applyToAllLoadedTiles();
   assert.is(selectNothing.selectionSize, 0);
+  assert.is(selectNothing.get(), undefined);
 });
 
 test('Test sorting of selections', async () => {
@@ -121,6 +122,12 @@ test('Test sorting of selections', async () => {
   const mid = sorted.get(Math.floor(sorted.selectionSize / 2));
   assert.ok(mid.random > 0.45);
   assert.ok(mid.random < 0.55);
+
+  // Check negative offsets
+  const end = sorted.get(-1);
+  const end2 = sorted.get(8191);
+  assert.ok(end.random > 0.99);
+  assert.ok(end.ix === end2.ix);
 });
 
 test('Test iterated sorting of selections', async () => {
@@ -169,14 +176,14 @@ test('Test iterated sorting of selections', async () => {
   const second = sorted.iterator(5);
 
   const numbers = [0, 1, 2, 3, 5, 6, 7, 8, 9, 10];
-  const firstVals = numbers.map(d => first.next().value[sortKey])
-  
+  const firstVals = numbers.map((d) => first.next().value[sortKey]);
+
   for (let i = 0; i < 5; i++) {
     assert.ok(firstVals[5 + i] === second.next().value[sortKey]);
   }
 });
 
-test ('Iterated sorting of empty selection', async() => {
+test('Iterated sorting of empty selection', async () => {
   const dataset = createIntegerDataset();
   await dataset.root_tile.preprocessRootTileInfo();
   const emptySelection = new DataSelection(dataset, {
@@ -207,7 +214,7 @@ test ('Iterated sorting of empty selection', async() => {
     thrown = true;
   }
   assert.ok(thrown);
-})
+});
 
 test('Edge cases for iterated sorting of selections', async () => {
   const dataset = createIntegerDataset();