chore: added caching layer for table widget data

app/client/src/widgets/TableWidgetV2/component/VirtualTable/InifiniteScrollBody/index.tsx

@@ -1,11 +1,10 @@
-import React, { type Ref } from "react";
+import React, { useMemo, type Ref } from "react";
 import { type ReactElementType } from "react-window";
 import InfiniteLoader from "react-window-infinite-loader";
 import type SimpleBar from "simplebar-react";
 import { FixedInfiniteVirtualList } from "../../TableBodyCoreComponents/VirtualList";
-import { useAppsmithTable } from "../../TableContext";
 import { LoadingIndicator } from "../../LoadingIndicator";
-import { useInfiniteVirtualization } from "./useInfiniteVirtualization";
+import { useAppsmithTable } from "../../TableContext";
 
 interface InfiniteScrollBodyProps {
   innerElementType: ReactElementType;
@@ -15,28 +14,25 @@ const InfiniteScrollBodyComponent = React.forwardRef(
   (props: InfiniteScrollBodyProps, ref: Ref<SimpleBar>) => {
     const {
       height,
+      isItemLoaded,
       isLoading,
       nextPageClick,
       pageSize,
       subPage: rows,
       tableSizes,
       totalRecordsCount,
     } = useAppsmithTable();
-    const { cachedRows, isItemLoaded, itemCount, loadMoreItems } =
-      useInfiniteVirtualization({
-        rows,
-        totalRecordsCount,
-        isLoading,
-        loadMore: nextPageClick,
-        pageSize,
-      });
+
+    const itemCount = useMemo(() => {
+      return totalRecordsCount ?? rows.length;
+    }, [totalRecordsCount, rows]);
 
     return (
       <div className="simplebar-content-wrapper">
         <InfiniteLoader
           isItemLoaded={isItemLoaded}
           itemCount={itemCount}
-          loadMoreItems={loadMoreItems}
+          loadMoreItems={nextPageClick}
           minimumBatchSize={pageSize}
         >
           {({ onItemsRendered, ref: infiniteLoaderRef }) => (
@@ -48,7 +44,7 @@ const InfiniteScrollBodyComponent = React.forwardRef(
               onItemsRendered={onItemsRendered}
               outerRef={ref}
               pageSize={pageSize}
-              rows={cachedRows}
+              rows={rows}
               tableSizes={tableSizes}
             />
           )}

app/client/src/widgets/TableWidgetV2/component/index.tsx

@@ -15,6 +15,7 @@ import equal from "fast-deep-equal/es6";
 import { useCallback } from "react";
 import type { EditableCell, TableVariant } from "../constants";
 import { ColumnTypes } from "../constants";
+import { useCachingVirtualization } from "./useCachingVirtualization";
 
 export interface ColumnMenuOptionProps {
   content: string | JSX.Element;
@@ -171,6 +172,18 @@ function ReactTableComponent(props: ReactTableComponentProps) {
     width,
   } = props;
 
+  const { cachedRows, isItemLoaded, isLoadingData, loadMoreNextPage } =
+    useCachingVirtualization({
+      data: tableData,
+      editableCell,
+      isLoading,
+      nextPageClick,
+      isInfiniteScrollEnabled,
+      pageSize,
+      totalRecordsCount,
+      isAddRowInProgress,
+    });
+
   const sortTableColumn = useCallback(
     (columnIndex: number, asc: boolean) => {
       if (allowSorting) {
@@ -236,7 +249,7 @@ function ReactTableComponent(props: ReactTableComponentProps) {
       columnWidthMap={columnWidthMap}
       columns={columns}
       compactMode={compactMode || CompactModeTypes.DEFAULT}
-      data={tableData}
+      data={cachedRows}
       delimiter={delimiter}
       disableDrag={memoziedDisableDrag}
       disabledAddNewRowSave={disabledAddNewRowSave}
@@ -250,14 +263,15 @@ function ReactTableComponent(props: ReactTableComponentProps) {
       height={height}
       isAddRowInProgress={isAddRowInProgress}
       isInfiniteScrollEnabled={isInfiniteScrollEnabled}
-      isLoading={isLoading}
+      isItemLoaded={isItemLoaded}
+      isLoading={isLoadingData}
       isSortable={isSortable}
       isVisibleDownload={isVisibleDownload}
       isVisibleFilters={isVisibleFilters}
       isVisiblePagination={isVisiblePagination}
       isVisibleSearch={isVisibleSearch}
       multiRowSelection={multiRowSelection}
-      nextPageClick={nextPageClick}
+      nextPageClick={loadMoreNextPage}
       onAddNewRow={onAddNewRow}
       onAddNewRowAction={onAddNewRowAction}
       onBulkEditDiscard={onBulkEditDiscard}

app/client/src/widgets/TableWidgetV2/component/types.ts

@@ -80,6 +80,7 @@ export interface TableProps {
   showConnectDataOverlay: boolean;
   onConnectData: () => void;
   isInfiniteScrollEnabled: boolean;
+  isItemLoaded: (index: number) => boolean;
 }
 
 export interface TableProviderProps extends TableProps {

app/client/src/widgets/TableWidgetV2/component/useCachingVirtualization.tsx

@@ -0,0 +1,184 @@
+import { concat } from "lodash";
+import { useCallback, useEffect, useMemo, useRef, useState } from "react";
+import type { EditableCell } from "../constants";
+
+export interface UseCachingVirtualizationProps {
+  editableCell: EditableCell;
+  data: Array<Record<string, unknown>>;
+  isLoading: boolean;
+  isInfiniteScrollEnabled: boolean;
+  nextPageClick: () => void;
+  pageSize: number;
+  totalRecordsCount?: number;
+  isAddRowInProgress: boolean;
+}
+
+export interface UseCachingTableDataReturn {
+  isLoadingData: boolean;
+  cachedRows: Array<Record<string, unknown>>;
+  loadMoreNextPage: (
+    startIndex?: number,
+    stopIndex?: number,
+  ) => Promise<void> | void;
+  isItemLoaded: (index: number) => boolean;
+}
+
+export const useCachingVirtualization = ({
+  data,
+  editableCell,
+  isAddRowInProgress,
+  isInfiniteScrollEnabled,
+  isLoading,
+  nextPageClick,
+  pageSize,
+  totalRecordsCount,
+}: UseCachingVirtualizationProps): UseCachingTableDataReturn => {
+  const [cachedRows, setCachedRows] = useState<Array<Record<string, unknown>>>(
+    [],
+  );
+  const [isLoadingData, setIsLoadingData] = useState(false);
+  const lastLoadedPageRef = useRef<number>(0);
+  const hasMoreDataRef = useRef<boolean>(true);
+  const isAddRowInProgressRef = useRef<{
+    isAddRowInProgress: boolean;
+    addedNewRow: boolean;
+  }>({ isAddRowInProgress: false, addedNewRow: false });
+
+  const editableCellRef = useRef<{
+    isEditing: boolean;
+    editingCell?: EditableCell;
+  }>({ isEditing: false, editingCell: undefined });
+
+  const maxPages = useMemo(() => {
+    if (!totalRecordsCount) return Infinity;
+
+    return Math.ceil(totalRecordsCount / pageSize);
+  }, [totalRecordsCount, pageSize]);
+
+  const isItemLoaded = useCallback(
+    (index: number) => {
+      const pageIndex = Math.floor(index / pageSize);
+
+      return (
+        pageIndex >= maxPages ||
+        pageIndex < lastLoadedPageRef.current ||
+        !hasMoreDataRef.current
+      );
+    },
+    [pageSize, maxPages],
+  );
+
+  useEffect(() => {
+    if (isInfiniteScrollEnabled) {
+      setCachedRows([]);
+    }
+  }, [isInfiniteScrollEnabled]);
+
+  useEffect(() => {
+    isAddRowInProgressRef.current.isAddRowInProgress = isAddRowInProgress;
+  }, [isAddRowInProgress]);
+
+  useEffect(() => {
+    if (editableCell.index !== -1) {
+      editableCellRef.current.editingCell = editableCell;
+      editableCellRef.current.isEditing = true;
+      setCachedRows((e) => {
+        const newRows = [...e];
+
+        newRows[editableCell.index][editableCell.column] =
+          editableCell.inputValue;
+
+        return newRows;
+      });
+    } else {
+      editableCellRef.current.editingCell = undefined;
+    }
+  }, [editableCell]);
+
+  useEffect(() => {
+    if (isInfiniteScrollEnabled) {
+      if (isAddRowInProgressRef.current.isAddRowInProgress) {
+        if (isAddRowInProgressRef.current.addedNewRow) {
+          setCachedRows((e) => concat([data[0]], e.slice(1)));
+        } else {
+          setCachedRows((e) => concat([data[0]], e));
+          isAddRowInProgressRef.current.addedNewRow = true;
+        }
+
+        // TODO: If new rows are added, we would have to update the cached rows
+        return;
+      } else if (data.length > 0) {
+        if (isAddRowInProgressRef.current.addedNewRow) {
+          isAddRowInProgressRef.current.addedNewRow = false;
+          setCachedRows((e) => e.slice(1));
+        } else if (editableCellRef.current.isEditing) {
+          editableCellRef.current.isEditing = false;
+        } else {
+          const currentPageIndex = lastLoadedPageRef.current;
+
+          // Only increment if we got a full page or some data
+          if (data.length === pageSize) {
+            lastLoadedPageRef.current = currentPageIndex + 1;
+          } else if (data.length < pageSize && data.length > 0) {
+            // If we got less than a full page, assume this is the last page
+            hasMoreDataRef.current = false;
+          }
+
+          setCachedRows((e) => concat(e, data));
+        }
+      } else if (data.length === 0 && lastLoadedPageRef.current > 0) {
+        if (isAddRowInProgressRef.current.addedNewRow) {
+          isAddRowInProgressRef.current.addedNewRow = false;
+          setCachedRows((e) => concat(e.slice(1), data));
+        } else {
+          // If no rows are returned and we've loaded at least one page, assume end of data
+          hasMoreDataRef.current = false;
+        }
+      }
+    } else {
+      if (data.length > 0) {
+        setCachedRows(data);
+      }
+    }
+  }, [data, isInfiniteScrollEnabled]);
+
+  useEffect(() => {
+    setIsLoadingData(isLoading);
+  }, [isLoading]);
+
+  const loadMoreNextPage = useCallback(
+    async (startIndex?: number, stopIndex?: number) => {
+      if (isInfiniteScrollEnabled) {
+        if (!isLoadingData && hasMoreDataRef.current && !isAddRowInProgress) {
+          const targetPage = Math.floor(stopIndex! / pageSize);
+
+          if (
+            targetPage >= lastLoadedPageRef.current &&
+            targetPage < maxPages
+          ) {
+            nextPageClick();
+          }
+        }
+
+        return Promise.resolve();
+      } else {
+        nextPageClick();
+      }
+    },
+    [
+      isLoadingData,
+      nextPageClick,
+      isInfiniteScrollEnabled,
+      pageSize,
+      maxPages,
+      isAddRowInProgress,
+    ],
+  );
+
+  return {
+    cachedRows,
+    isItemLoaded,
+    isLoadingData,
+    loadMoreNextPage,
+  };
+};

app/client/src/widgets/TableWidgetV2/widget/index.tsx

@@ -1957,7 +1957,8 @@ class TableWidgetV2 extends BaseWidget<TableWidgetProps, WidgetState> {
      */
     if (this.props.isAddRowInProgress) {
       row = filteredTableData[rowIndex - 1];
-      originalIndex = rowIndex === 0 ? -1 : row[ORIGINAL_INDEX_KEY] ?? rowIndex;
+      originalIndex =
+        rowIndex === 0 ? -1 : row?.[ORIGINAL_INDEX_KEY] ?? rowIndex;
     } else {
       row = filteredTableData[rowIndex];
       originalIndex = row ? row[ORIGINAL_INDEX_KEY] ?? rowIndex : rowIndex;