feat(frontend): skip Arrow-to-object materialization via columnar engine

Avoid converting Arrow Tables into DataRecord[] arrays (~84 ms at 200K
rows) by reading column values directly through a new ColumnarDataSource
interface. Controlled by FEATURE_FLAGS.arrowColumnar (on by default).

Author: mayagbarnes

streamlit_pivot/frontend/src/PivotRoot.tsx

@@ -42,11 +42,13 @@ import {
 } from "./engine/types";
 import { PivotData, type PivotDataOptions } from "./engine/PivotData";
 import {
+  createArrowDataSource,
   parseArrowToRecords,
   getArrowColumnNames,
   getNumericColumns,
 } from "./engine/parseArrow";
 import {
+  FEATURE_FLAGS,
   measureSync,
   logMetrics,
   checkBudgets,
@@ -131,9 +133,20 @@ const PivotRoot: FC<PivotRootProps> = ({
     [currentConfig],
   );
 
-  const { records, parseMs } = useMemo(() => {
+  const { pivotInput, parseMs } = useMemo(() => {
+    if (FEATURE_FLAGS.arrowColumnar) {
+      const measured = measureSync(() => createArrowDataSource(dataframe));
+      const ds = measured.result;
+      return {
+        pivotInput: ds && ds.numRows > 0 ? ds : null,
+        parseMs: measured.elapsedMs,
+      };
+    }
     const measured = measureSync(() => parseArrowToRecords(dataframe));
-    return { records: measured.result, parseMs: measured.elapsedMs };
+    return {
+      pivotInput: measured.result.length > 0 ? measured.result : null,
+      parseMs: measured.elapsedMs,
+    };
   }, [dataframe]);
 
   const rawAllColumns = useMemo(
@@ -185,12 +198,12 @@ const PivotRoot: FC<PivotRootProps> = ({
   );
 
   const { pivotData, computeMs } = useMemo(() => {
-    if (records.length === 0) return { pivotData: null, computeMs: 0 };
+    if (!pivotInput) return { pivotData: null, computeMs: 0 };
     const measured = measureSync(
-      () => new PivotData(records, currentConfig, pivotOptions),
+      () => new PivotData(pivotInput, currentConfig, pivotOptions),
     );
     return { pivotData: measured.result, computeMs: measured.elapsedMs };
-  }, [records, currentConfig, pivotOptions]);
+  }, [pivotInput, currentConfig, pivotOptions]);
 
   const budget = useMemo(() => {
     if (!pivotData) return null;
@@ -268,7 +281,7 @@ const PivotRoot: FC<PivotRootProps> = ({
       firstMountMs:
         debugMetrics?.firstMountMs ??
         Math.round((parseMs + computeMs + renderMs) * 100) / 100,
-      sourceRows: records.length,
+      sourceRows: pivotData.recordCount,
       sourceCols: rawAllColumns.length,
       totalRows: pivotData.uniqueRowKeyCount,
       totalCols: pivotData.uniqueColKeyCount,
@@ -304,7 +317,7 @@ const PivotRoot: FC<PivotRootProps> = ({
     parseMs,
     pivotData,
     rawAllColumns.length,
-    records.length,
+    pivotInput,
   ]);
 
   // perf_metrics are exposed via the data-perf-metrics DOM attribute (set

streamlit_pivot/frontend/src/engine/PivotData.test.ts

@@ -27,8 +27,10 @@ import {
   normalizeAggregationConfig,
   type AggregationConfig,
   type AggregationType,
+  type ColumnarDataSource,
   type PivotConfigV1,
 } from "./types";
+import { DataRecordSource } from "./parseArrow";
 import { measureSync, DEFAULT_BUDGETS } from "./perf";
 
 type TestConfigOverrides = Partial<Omit<PivotConfigV1, "aggregation">> & {
@@ -1648,3 +1650,79 @@ describe("PivotData - getColumnNames", () => {
     expect(cols).toHaveLength(4);
   });
 });
+
+describe("PivotData - DataRecordSource wrapper", () => {
+  it("matches array-backed aggregates and keys", () => {
+    const cfg = makeConfig();
+    const fromArray = new PivotData(SAMPLE_DATA, cfg);
+    const source = new DataRecordSource(
+      SAMPLE_DATA,
+      Object.keys(SAMPLE_DATA[0]!),
+    );
+    const fromSource = new PivotData(source, cfg);
+    expect(fromSource.getRowKeys()).toEqual(fromArray.getRowKeys());
+    expect(fromSource.getColKeys()).toEqual(fromArray.getColKeys());
+    expect(fromSource.getAggregator(["US"], ["2023"]).value()).toEqual(
+      fromArray.getAggregator(["US"], ["2023"]).value(),
+    );
+  });
+});
+
+class TestColumnarSource implements ColumnarDataSource {
+  constructor(private readonly rows: DataRecord[]) {}
+
+  get numRows(): number {
+    return this.rows.length;
+  }
+
+  getValue(rowIndex: number, fieldName: string): unknown {
+    return this.rows[rowIndex]![fieldName];
+  }
+
+  getColumnNames(): string[] {
+    return Object.keys(this.rows[0] ?? {});
+  }
+}
+
+describe("PivotData - non-materialized columnar source", () => {
+  it("getMatchingRecords materializes rows on demand", () => {
+    const cfg = makeConfig();
+    const src = new TestColumnarSource([...SAMPLE_DATA]);
+    const pd = new PivotData(src, cfg);
+    const r = pd.getMatchingRecords({ region: "US", year: "2023" });
+    expect(r.totalCount).toBe(2);
+    expect(r.records).toHaveLength(2);
+    expect(
+      r.records.every((row) => row.region === "US" && row.year === "2023"),
+    ).toBe(true);
+  });
+
+  it("getUniqueValues scans the columnar source", () => {
+    const src = new TestColumnarSource(SAMPLE_DATA);
+    const pd = new PivotData(src, makeConfig());
+    expect(pd.getUniqueValues("region")).toEqual(["EU", "US"]);
+  });
+
+  it("aggregates match array-backed results", () => {
+    const cfg = makeConfig();
+    const fromArray = new PivotData(SAMPLE_DATA, cfg);
+    const fromColumnar = new PivotData(
+      new TestColumnarSource(SAMPLE_DATA),
+      cfg,
+    );
+    expect(fromColumnar.getRowKeys()).toEqual(fromArray.getRowKeys());
+    expect(fromColumnar.getColKeys()).toEqual(fromArray.getColKeys());
+    expect(fromColumnar.getGrandTotal().value()).toEqual(
+      fromArray.getGrandTotal().value(),
+    );
+    expect(fromColumnar.recordCount).toBe(fromArray.recordCount);
+  });
+
+  it("getColumnNames works without materialized records", () => {
+    const src = new TestColumnarSource(SAMPLE_DATA);
+    const pd = new PivotData(src, makeConfig());
+    expect(pd.getColumnNames()).toEqual(
+      expect.arrayContaining(["region", "year", "revenue", "profit"]),
+    );
+  });
+});