From 02515c3876fdca8f6615dc83d61938bf496145d7 Mon Sep 17 00:00:00 2001
From: Aditya Hegde <adityahegderocks@gmail.com>
Date: Thu, 9 May 2024 20:03:10 +0530
Subject: [PATCH] fix: Make the preview table column width reactive (#4864)

* Make the preview table column width reactive

* Add test
---
 .../virtualized-table/VirtualTable.svelte     | 27 ++++-----
 .../virtualized-table/columnSizes.spec.ts     | 50 +++++++++++++++++
 .../virtualized-table/columnSizes.ts          | 56 +++++++++++++++++++
 3 files changed, 116 insertions(+), 17 deletions(-)
 create mode 100644 web-common/src/components/virtualized-table/columnSizes.spec.ts
 create mode 100644 web-common/src/components/virtualized-table/columnSizes.ts

diff --git a/web-common/src/components/virtualized-table/VirtualTable.svelte b/web-common/src/components/virtualized-table/VirtualTable.svelte
index e0c012e5694..970b078fac6 100644
--- a/web-common/src/components/virtualized-table/VirtualTable.svelte
+++ b/web-common/src/components/virtualized-table/VirtualTable.svelte
@@ -1,19 +1,7 @@
 <script lang="ts" context="module">
-  const columnSizes = (() => {
-    const sizes = new Map<string, number[]>();
-
-    return {
-      get: (key: string, calculator: () => number[]): number[] => {
-        let array = sizes.get(key);
-        if (!array) {
-          array = calculator();
-          sizes.set(key, array);
-        }
-        return array;
-      },
-      set: (key: string, value: number[]) => sizes.set(key, value),
-    };
-  })();
+  import { VirtualizedTableColumnSizes } from "@rilldata/web-common/components/virtualized-table/columnSizes";
+
+  const columnSizes = new VirtualizedTableColumnSizes();
 
   export const ROW_HEIGHT = 24;
   export const MIN_COL_WIDTH = 108;
@@ -107,7 +95,7 @@
   let scrollLeft = 0;
   let nextPinnedColumnPosition = ROW_HEADER_WIDTH;
 
-  $: columnWidths = columnSizes.get(name, () =>
+  $: columnWidths = columnSizes.get(name, columns, columnAccessor, () =>
     initColumnWidths({
       columns,
       rows,
@@ -173,6 +161,11 @@
         resizing.initialPixelWidth + delta,
         maxColWidth,
       );
+      columnSizes.set(
+        name,
+        columns[resizing.columnIndex].name as string,
+        columnWidths[resizing.columnIndex],
+      );
     });
   }
 
@@ -198,7 +191,7 @@
     const value =
       description ?? isHeader
         ? column
-        : (rows[index][column] as string | number | null);
+        : (rows[index]?.[column] as string | number | null);
     const type = columns.find((c) => c.name === column)?.type ?? "string";
 
     hovering = {
diff --git a/web-common/src/components/virtualized-table/columnSizes.spec.ts b/web-common/src/components/virtualized-table/columnSizes.spec.ts
new file mode 100644
index 00000000000..018c686a576
--- /dev/null
+++ b/web-common/src/components/virtualized-table/columnSizes.spec.ts
@@ -0,0 +1,50 @@
+import { VirtualizedTableColumnSizes } from "@rilldata/web-common/components/virtualized-table/columnSizes";
+import { describe, it, expect } from "vitest";
+
+describe("VirtualizedTableColumnSizes", () => {
+  it("Sanity check", () => {
+    const sizes = new VirtualizedTableColumnSizes();
+    expect(
+      sizes.get(
+        "model",
+        [{ name: "col1" }, { name: "col2" }, { name: "col3" }],
+        "name",
+        () => [50, 50, 50],
+      ),
+    ).toEqual([50, 50, 50]);
+
+    // set the width for col2 to something specific
+    sizes.set("model", "col2", 100);
+    // the width is saved
+    expect(
+      sizes.get(
+        "model",
+        [{ name: "col1" }, { name: "col2" }, { name: "col3" }],
+        "name",
+        () => [50, 50, 50],
+      ),
+    ).toEqual([50, 100, 50]);
+
+    // the width is saved when one of the columns is removed
+    expect(
+      sizes.get("model", [{ name: "col1" }, { name: "col2" }], "name", () => [
+        50, 50,
+      ]),
+    ).toEqual([50, 100]);
+
+    expect(
+      sizes.get("model", [{ name: "col1" }, { name: "col4" }], "name", () => [
+        50, 50,
+      ]),
+    ).toEqual([50, 50]);
+    // removing a column will remove any saved sizes. this is for cleaning up leaks
+    expect(
+      sizes.get(
+        "model",
+        [{ name: "col1" }, { name: "col2" }, { name: "col4" }],
+        "name",
+        () => [50, 50, 50],
+      ),
+    ).toEqual([50, 50, 50]);
+  });
+});
diff --git a/web-common/src/components/virtualized-table/columnSizes.ts b/web-common/src/components/virtualized-table/columnSizes.ts
new file mode 100644
index 00000000000..3adf3ab009f
--- /dev/null
+++ b/web-common/src/components/virtualized-table/columnSizes.ts
@@ -0,0 +1,56 @@
+import type { VirtualizedTableColumns } from "@rilldata/web-common/components/virtualized-table/types";
+import { V1MetricsViewColumn } from "@rilldata/web-common/runtime-client";
+
+export class VirtualizedTableColumnSizes {
+  private readonly sizesCache = new Map<string, Map<string, number>>();
+
+  public get(
+    key: string,
+    columns: (VirtualizedTableColumns | V1MetricsViewColumn)[],
+    columnAccessor: keyof VirtualizedTableColumns,
+    calculator: () => number[],
+  ): number[] {
+    const cache = this.sizesCache.get(key);
+    if (!cache) {
+      const sizes = calculator();
+      this.sizesCache.set(
+        key,
+        new Map<string, number>(
+          sizes.map((s, i) => [columns[i][columnAccessor], s]),
+        ),
+      );
+      return sizes;
+    }
+
+    let missingSize = false;
+    const sizes = columns.map((column) => {
+      if (!cache.has(column[columnAccessor])) {
+        missingSize = true;
+        return 0;
+      }
+      return cache.get(column[columnAccessor]) as number;
+    });
+    if (!missingSize) return sizes;
+
+    const newSizes = calculator();
+    // retain sizes from cache
+    newSizes.forEach((_, i) => {
+      if (cache.has(columns[i][columnAccessor])) {
+        newSizes[i] = cache.get(columns[i][columnAccessor]) as number;
+      }
+    });
+    // reset the cache
+    this.sizesCache.set(
+      key,
+      new Map<string, number>(
+        newSizes.map((s, i) => [columns[i][columnAccessor], s]),
+      ),
+    );
+
+    return newSizes;
+  }
+
+  public set(name: string, colName: string, value: number) {
+    this.sizesCache.get(name)?.set(colName, value);
+  }
+}