[Cascade] Virtualizer Improvements

src/platform/packages/shared/shared-ux/document_data_cascade/impl/index.ts

@@ -7,15 +7,22 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-export { DataCascade, DataCascadeRow, DataCascadeRowCell } from './src/components';
+export {
+  DataCascade,
+  DataCascadeRow,
+  DataCascadeRowCell,
+  toRestorableState,
+} from './src/components';
 export type {
   GroupNode,
   LeafNode,
   DataCascadeProps,
   DataCascadeImplRef,
   DataCascadeRowProps,
   DataCascadeRowCellProps,
-  CascadeRowCellNestedVirtualizationAnchorProps,
+  CascadeGroupNodeUIInteraction,
+  DataCascadeUISnapshot,
+  DataCascadeRestorableState,
 } from './src/components';
 export * from './src/lib';
 export { NumberBadge } from './src/components/helpers';

src/platform/packages/shared/shared-ux/document_data_cascade/impl/src/components/data_cascade_impl/data_cascade_impl.styles.tsx

@@ -31,6 +31,16 @@ export const dataCascadeImplStyles = (euiTheme: UseEuiTheme['euiTheme']) => ({
     isolation: 'isolate',
   }),
   containerInner: css([relativePosition, { height: '100%' }]),
+  cascadeLoadingContainer: css({
+    height: '100%',
+    width: '100%',
+    position: 'absolute',
+    top: 0,
+    left: 0,
+    display: 'flex',
+    justifyContent: 'center',
+    alignItems: 'center',
+  }),
   cascadeTreeGridBlock: css([
     overflowYAuto,
     relativePosition,

src/platform/packages/shared/shared-ux/document_data_cascade/impl/src/components/data_cascade_impl/data_cascade_impl.tsx

@@ -12,10 +12,13 @@ import {
   EuiAutoSizer,
   EuiFlexGroup,
   EuiFlexItem,
+  EuiLoadingSpinner,
   useEuiTheme,
   useGeneratedHtmlId,
   useIsWithinMaxBreakpoint,
+  type EuiAutoSize,
 } from '@elastic/eui';
+import { css } from '@emotion/react';
 import { CascadeHeaderPrimitive } from './data_cascade_header';
 import { CascadeRowPrimitive, CascadeRowHeaderSlotsScrollSyncProvider } from './data_cascade_row';
 import { CascadeRowCellPrimitive } from './data_cascade_row_cell';
@@ -63,8 +66,7 @@ export function DataCascadeImpl<G extends GroupNode, L extends LeafNode>({
   enableRowSelection = false,
   enableStickyGroupHeader = true,
   allowMultipleRowToggle = false,
-  initialScrollOffset,
-  initialRect,
+  initialState,
   cascadeRef,
 }: DataCascadeImplProps<G, L>) {
   const rowElement = Children.only(children);
@@ -138,43 +140,42 @@ export function DataCascadeImpl<G extends GroupNode, L extends LeafNode>({
   const { collectVirtualizerStateChanges } = useExposePublicApi<G, L>(cascadeRef, {
     rows,
     enableStickyGroupHeader,
+    childController: virtualizerInstance.current?.childController,
   });
 
-  // persist the virtualizer instance to ref, so that invocations of getVirtualizer will always return the latest instance
+  const initialPersistedAnchors = useMemo(() => {
+    if (!initialState?.connectedChildren) return undefined;
+    const anchors: Record<string, number | null> = {};
+    for (const [cellId, child] of Object.entries(initialState.connectedChildren)) {
+      anchors[cellId] = child.scrollAnchorItemIndex;
+    }
+    return anchors;
+  }, [initialState?.connectedChildren]);
+
   virtualizerInstance.current = useCascadeVirtualizer<G>({
     rows,
     overscan,
     getScrollElement,
     enableStickyGroupHeader,
     estimatedRowHeight: size === 's' ? 32 : size === 'm' ? 40 : 48,
     onStateChange: collectVirtualizerStateChanges,
-    initialOffset: initialScrollOffset,
-    initialRect,
+    initialRect: initialState?.scrollRect,
+    initialAnchorItemIndex: initialState?.scrollAnchorItemIndex ?? undefined,
+    initialPersistedAnchors,
   });
 
-  const {
-    getVirtualItems,
-    getTotalSize,
-    range,
-    measureElement,
-    virtualizedRowComputedTranslateValue,
-    scrollToVirtualizedIndex,
-    scrollOffset: virtualizerScrollOffset,
-    isScrolling,
-  } = virtualizerInstance.current;
-
   // Calculate activeStickyIndex directly from the virtualizer's current range.
   // This ensures the value is always current and never stale from intermediate memoization.
   const activeStickyIndex = calculateActiveStickyIndex(
     rows,
-    range?.startIndex ?? 0,
+    virtualizerInstance.current?.range?.startIndex ?? 0,
     enableStickyGroupHeader
   );
 
   useRegisterCascadeAccessibilityHelpers<G>({
     tableRows: rows,
     tableWrapperElement: cascadeWrapperRef.current!,
-    scrollToRowIndex: scrollToVirtualizedIndex,
+    scrollToRowIndex: virtualizerInstance.current!.scrollToVirtualizedIndex,
   });
 
   const virtualCascadeRowRenderer = useCallback<VirtualizedCascadeListProps<G>['listItemRenderer']>(
@@ -188,13 +189,15 @@ export function DataCascadeImpl<G extends GroupNode, L extends LeafNode>({
           virtualRow: virtualItem,
           virtualRowStyle,
           isMobile,
-          innerRef: measureElement,
+          innerRef: virtualizerInstance.current!.measureElement,
           activeStickyRenderSlotRef,
+          // getVirtualizer will not return undefined here as it is set immediately after the first render
+          getVirtualizer: getVirtualizer as () => ReturnType<typeof useCascadeVirtualizer>,
           ...rowElement.props,
         }}
       />
     ),
-    [size, enableRowSelection, isMobile, measureElement, rowElement.props]
+    [size, enableRowSelection, isMobile, getVirtualizer, rowElement.props]
   );
 
   const treeGridContainerARIAAttributes = useTreeGridContainerARIAAttributes(headerId);
@@ -203,59 +206,98 @@ export function DataCascadeImpl<G extends GroupNode, L extends LeafNode>({
     return (
       enableStickyGroupHeader &&
       activeStickyIndex !== null &&
-      (virtualizerScrollOffset ?? 0) > (virtualizedRowComputedTranslateValue.get(0) ?? 0)
+      (virtualizerInstance.current?.scrollOffset ?? 0) >
+        (virtualizerInstance.current?.virtualizedRowComputedTranslateValue.get(0) ?? 0)
     );
-  }, [
-    activeStickyIndex,
-    enableStickyGroupHeader,
-    virtualizerScrollOffset,
-    virtualizedRowComputedTranslateValue,
-  ]);
+  }, [activeStickyIndex, enableStickyGroupHeader]);
+
+  const cascadeTreeGridRenderer = useCallback(
+    (containerSize: EuiAutoSize) => {
+      return (
+        <React.Fragment>
+          <div
+            css={css([
+              styles.cascadeLoadingContainer,
+              {
+                visibility:
+                  containerSize.height === 0 || containerSize.width === 0 ? 'visible' : 'hidden',
+              },
+            ])}
+          >
+            <EuiLoadingSpinner size="l" />
+          </div>
+          <div
+            ref={scrollElementRef}
+            css={css([
+              styles.cascadeTreeGridBlock,
+              {
+                visibility:
+                  containerSize.height === 0 || containerSize.width === 0 ? 'hidden' : 'visible',
+              },
+            ])}
+            style={containerSize}
+            data-test-subj="data-cascade-scroll-container"
+          >
+            {/* Always render the slot so the ref is available immediately.
+            Use hidden style when not visible to avoid layout impact. */}
+            <div
+              css={
+                shouldRenderStickyHeader
+                  ? styles.cascadeTreeGridHeaderStickyRenderSlot
+                  : styles.cascadeTreeGridHeaderStickyRenderSlotHidden
+              }
+            >
+              <div ref={activeStickyRenderSlotRef} />
+            </div>
+            <div
+              css={styles.cascadeTreeGridWrapper}
+              style={{ height: virtualizerInstance.current?.getTotalSize() }}
+            >
+              <div {...treeGridContainerARIAAttributes} css={relativePosition}>
+                <CascadeRowHeaderSlotsScrollSyncProvider
+                  disableScrollSync={virtualizerInstance.current?.isScrolling}
+                >
+                  <VirtualizedCascadeRowList<G>
+                    {...{
+                      activeStickyIndex,
+                      getVirtualItems: virtualizerInstance.current!.getVirtualItems,
+                      virtualizedRowComputedTranslateValue:
+                        virtualizerInstance.current!.virtualizedRowComputedTranslateValue,
+                      rows,
+                      listItemRenderer: virtualCascadeRowRenderer,
+                    }}
+                  />
+                </CascadeRowHeaderSlotsScrollSyncProvider>
+              </div>
+            </div>
+          </div>
+        </React.Fragment>
+      );
+    },
+    [
+      styles,
+      shouldRenderStickyHeader,
+      activeStickyRenderSlotRef,
+      treeGridContainerARIAAttributes,
+      rows,
+      virtualCascadeRowRenderer,
+      activeStickyIndex,
+    ]
+  );
 
   return (
     <div ref={cascadeWrapperRef} data-test-subj="data-cascade" css={styles.container}>
       <EuiFlexGroup direction="column" gutterSize="none" css={styles.containerInner}>
         <EuiFlexItem grow={false}>
           <TableHeader headerColumns={headerColumns} />
         </EuiFlexItem>
-        <EuiFlexItem grow={true}>
-          <EuiAutoSizer doNotBailOutOnEmptyChildren>
-            {(scrollContainerSize) => (
-              <div
-                ref={scrollElementRef}
-                css={styles.cascadeTreeGridBlock}
-                style={{
-                  ...scrollContainerSize,
-                }}
-              >
-                {/* Always render the slot so the ref is available immediately.
-                    Use hidden style when not visible to avoid layout impact. */}
-                <div
-                  css={
-                    shouldRenderStickyHeader
-                      ? styles.cascadeTreeGridHeaderStickyRenderSlot
-                      : styles.cascadeTreeGridHeaderStickyRenderSlotHidden
-                  }
-                >
-                  <div ref={activeStickyRenderSlotRef} />
-                </div>
-                <div css={styles.cascadeTreeGridWrapper} style={{ height: getTotalSize() }}>
-                  <div {...treeGridContainerARIAAttributes} css={relativePosition}>
-                    <CascadeRowHeaderSlotsScrollSyncProvider disableScrollSync={isScrolling}>
-                      <VirtualizedCascadeRowList<G>
-                        {...{
-                          activeStickyIndex,
-                          getVirtualItems,
-                          virtualizedRowComputedTranslateValue,
-                          rows,
-                          listItemRenderer: virtualCascadeRowRenderer,
-                        }}
-                      />
-                    </CascadeRowHeaderSlotsScrollSyncProvider>
-                  </div>
-                </div>
-              </div>
-            )}
+        <EuiFlexItem grow={true} style={{ position: 'relative' }}>
+          <EuiAutoSizer
+            defaultHeight={initialState?.scrollRect?.height}
+            defaultWidth={initialState?.scrollRect?.width}
+            doNotBailOutOnEmptyChildren
+          >
+            {cascadeTreeGridRenderer}
           </EuiAutoSizer>
         </EuiFlexItem>
       </EuiFlexGroup>

src/platform/packages/shared/shared-ux/document_data_cascade/impl/src/components/data_cascade_impl/data_cascade_row/cascade_row.tsx

@@ -7,14 +7,19 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import React, { useMemo } from 'react';
+import React, { useEffect, useMemo } from 'react';
 import { createPortal } from 'react-dom';
 import { EuiFlexGroup, EuiFlexItem, useEuiTheme } from '@elastic/eui';
 import type { CascadeRowPrimitiveProps } from '../types';
 import { type LeafNode, type GroupNode, useDataCascadeState } from '../../../store_provider';
 import { TableCellRender, useAdaptedTableRows } from '../../../lib/core/table';
 import { useTreeGridRowARIAAttributes } from '../../../lib/core/accessibility';
-import { isCascadeGroupRowNode } from '../../../lib/utils';
+import {
+  isCascadeGroupRowNode,
+  getCascadeRowNodePath,
+  getCascadeRowNodePathValueRecord,
+  getCascadeRowLeafDataCacheKey,
+} from '../../../lib/utils';
 import {
   styles as cascadeRowStyles,
   rootRowAttribute,
@@ -43,6 +48,7 @@ export function CascadeRowPrimitive<G extends GroupNode, L extends LeafNode>({
   virtualRow,
   virtualRowStyle,
   enableRowSelection,
+  getVirtualizer,
   enableSecondaryExpansionAction,
 }: CascadeRowPrimitiveProps<G, L>) {
   const { euiTheme } = useEuiTheme();
@@ -56,6 +62,19 @@ export function CascadeRowPrimitive<G extends GroupNode, L extends LeafNode>({
 
   const isGroupNode = isCascadeGroupRowNode(currentGroupByColumns, rowInstance);
 
+  const cellId = useMemo(() => {
+    if (isGroupNode) return null;
+    const nodePath = getCascadeRowNodePath(currentGroupByColumns, rowInstance);
+    const nodePathMap = getCascadeRowNodePathValueRecord(currentGroupByColumns, rowInstance);
+    return getCascadeRowLeafDataCacheKey(nodePath, nodePathMap, rowInstance.id);
+  }, [currentGroupByColumns, rowInstance, isGroupNode]);
+
+  useEffect(() => {
+    if (!rowIsExpanded && cellId) {
+      getVirtualizer().childController?.clearPersistedAnchor(cellId);
+    }
+  }, [rowIsExpanded, cellId, getVirtualizer]);
+
   const styles = useMemo(() => {
     return cascadeRowStyles(
       euiTheme,