@parca/profile 0.19.12 → 0.19.14

package/src/ProfileView/index.tsx

@@ -19,7 +19,12 @@ import ColorStackLegend from './components/ColorStackLegend';
 import {getDashboardItem} from './components/DashboardItems';
 import {DashboardLayout} from './components/DashboardLayout';
 import {ProfileHeader} from './components/ProfileHeader';
-import {IcicleGraphToolbar, TableToolbar, VisualisationToolbar} from './components/Toolbars';
+import {
+  IcicleGraphToolbar,
+  SandwichIcicleGraphToolbar,
+  TableToolbar,
+  VisualisationToolbar,
+} from './components/Toolbars';
 import {DashboardProvider} from './context/DashboardContext';
 import {ProfileViewContextProvider} from './context/ProfileViewContext';
 import {useProfileMetadata} from './hooks/useProfileMetadata';
@@ -63,7 +68,6 @@ export const ProfileView = ({
     clearSelection,
     setGroupByLabels,
     sandwichFunctionName,
-    setSandwichFunctionName,
     resetSandwichFunctionName,
   } = useVisualizationState();
 
@@ -119,6 +123,12 @@ export const ProfileView = ({
         currentSearchString={currentSearchString}
       />
     ),
+    sandwich: (
+      <SandwichIcicleGraphToolbar
+        resetSandwichFunctionName={resetSandwichFunctionName}
+        sandwichFunctionName={sandwichFunctionName}
+      />
+    ),
   };
 
   const hasProfileSource = profileSource !== undefined && profileSource.toString(timezone) !== '';
@@ -154,8 +164,6 @@ export const ProfileView = ({
             setGroupByLabels={setGroupByLabels}
             showVisualizationSelector={showVisualizationSelector}
             sandwichFunctionName={sandwichFunctionName}
-            setSandwichFunctionName={setSandwichFunctionName}
-            resetSandwichFunctionName={resetSandwichFunctionName}
           />
 
           {isColorStackLegendEnabled && (

package/src/Sandwich/components/CalleesSection.tsx

@@ -21,7 +21,6 @@ import {type ProfileSource} from '../../ProfileSource';
 
 interface CalleesSectionProps {
   calleesRef: React.RefObject<HTMLDivElement>;
-  isHalfScreen: boolean;
   calleesFlamegraphResponse?: {
     report: {
       oneofKind: string;
@@ -40,7 +39,6 @@ interface CalleesSectionProps {
 
 export function CalleesSection({
   calleesRef,
-  isHalfScreen,
   calleesFlamegraphResponse,
   calleesFlamegraphLoading,
   calleesFlamegraphError,
@@ -68,11 +66,7 @@ export function CalleesSection({
         error={calleesFlamegraphError}
         isHalfScreen={true}
         width={
-          calleesRef.current != null
-            ? isHalfScreen
-              ? (calleesRef.current.getBoundingClientRect().width - 54) / 2
-              : calleesRef.current.getBoundingClientRect().width - 16
-            : 0
+          calleesRef.current != null ? calleesRef.current.getBoundingClientRect().width - 25 : 0
         }
         metadataMappingFiles={metadataMappingFiles}
         metadataLoading={false}

package/src/Sandwich/components/CallersSection.tsx

@@ -11,17 +11,33 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import React from 'react';
+import React, {useMemo} from 'react';
+
+import {Vector, tableFromIPC} from 'apache-arrow';
+import {Tooltip} from 'react-tooltip';
 
 import {type FlamegraphArrow} from '@parca/client';
+import {Button} from '@parca/components';
 
 import ProfileIcicleGraph from '../../ProfileIcicleGraph';
 import {type CurrentPathFrame} from '../../ProfileIcicleGraph/IcicleGraphArrow/utils';
 import {type ProfileSource} from '../../ProfileSource';
 
+const FIELD_DEPTH = 'depth';
+
+function getMaxDepth(depthColumn: Vector<any> | null): number {
+  if (depthColumn === null) return 0;
+
+  let max = 0;
+  for (const val of depthColumn) {
+    const numVal = Number(val);
+    if (numVal > max) max = numVal;
+  }
+  return max;
+}
+
 interface CallersSectionProps {
   callersRef: React.RefObject<HTMLDivElement>;
-  isHalfScreen: boolean;
   callersFlamegraphResponse?: {
     report: {
       oneofKind: string;
@@ -36,11 +52,13 @@ interface CallersSectionProps {
   curPathArrow: CurrentPathFrame[];
   setCurPathArrow: (path: CurrentPathFrame[]) => void;
   metadataMappingFiles?: string[];
+  isExpanded: boolean;
+  setIsExpanded: (isExpanded: boolean) => void;
+  defaultMaxFrames: number;
 }
 
 export function CallersSection({
   callersRef,
-  isHalfScreen,
   callersFlamegraphResponse,
   callersFlamegraphLoading,
   callersFlamegraphError,
@@ -49,40 +67,79 @@ export function CallersSection({
   curPathArrow,
   setCurPathArrow,
   metadataMappingFiles,
+  isExpanded,
+  setIsExpanded,
+  defaultMaxFrames,
 }: CallersSectionProps): JSX.Element {
+  const maxDepth = useMemo(() => {
+    if (
+      callersFlamegraphResponse?.report.oneofKind === 'flamegraphArrow' &&
+      callersFlamegraphResponse?.report?.flamegraphArrow != null
+    ) {
+      const table = tableFromIPC(callersFlamegraphResponse.report.flamegraphArrow.record);
+      const depthColumn = table.getChild(FIELD_DEPTH);
+      return getMaxDepth(depthColumn);
+    }
+    return 0;
+  }, [callersFlamegraphResponse]);
+
+  const shouldShowButton = maxDepth > defaultMaxFrames;
+
   return (
-    <div className="flex relative flex-row" ref={callersRef}>
-      <div className="[writing-mode:vertical-lr] -rotate-180 px-1 uppercase text-[10px] text-left">
-        Callers {'->'}
+    <>
+      {shouldShowButton && (
+        <Button
+          variant="neutral"
+          onClick={() => setIsExpanded(!isExpanded)}
+          className="absolute right-8 top-[-46px] z-10"
+          type="button"
+        >
+          <span
+            data-tooltip-content={
+              !isExpanded
+                ? `This profile has ${maxDepth} frames, showing only the top ${defaultMaxFrames} frames. Click to show more frames.`
+                : `This profile has ${maxDepth} frames, showing all frames. Click to hide frames.`
+            }
+            data-tooltip-id="show-more-frames"
+          >
+            {isExpanded ? 'Hide frames' : 'Show more frames'}
+          </span>
+          <Tooltip id="show-more-frames" />
+        </Button>
+      )}
+      <div className="flex relative flex-row overflow-hidden" ref={callersRef}>
+        <div className="[writing-mode:vertical-lr] -rotate-180 px-1 uppercase text-[10px] text-left flex-shrink-0">
+          Callers {'->'}
+        </div>
+        <div className="flex-1 overflow-hidden relative">
+          <ProfileIcicleGraph
+            arrow={
+              callersFlamegraphResponse?.report.oneofKind === 'flamegraphArrow'
+                ? callersFlamegraphResponse?.report?.flamegraphArrow
+                : undefined
+            }
+            total={BigInt(callersFlamegraphResponse?.total ?? '0')}
+            filtered={filtered}
+            profileType={profileSource?.ProfileType()}
+            loading={callersFlamegraphLoading}
+            error={callersFlamegraphError}
+            isHalfScreen={true}
+            width={
+              callersRef.current != null ? callersRef.current.getBoundingClientRect().width - 25 : 0
+            }
+            metadataMappingFiles={metadataMappingFiles}
+            metadataLoading={false}
+            isSandwichIcicleGraph={true}
+            curPathArrow={curPathArrow}
+            setNewCurPathArrow={setCurPathArrow}
+            isFlamegraph={true}
+            profileSource={profileSource}
+            tooltipId="callers"
+            maxFrameCount={defaultMaxFrames}
+            isExpanded={isExpanded}
+          />
+        </div>
       </div>
-      <ProfileIcicleGraph
-        arrow={
-          callersFlamegraphResponse?.report.oneofKind === 'flamegraphArrow'
-            ? callersFlamegraphResponse?.report?.flamegraphArrow
-            : undefined
-        }
-        total={BigInt(callersFlamegraphResponse?.total ?? '0')}
-        filtered={filtered}
-        profileType={profileSource?.ProfileType()}
-        loading={callersFlamegraphLoading}
-        error={callersFlamegraphError}
-        isHalfScreen={true}
-        width={
-          callersRef.current != null
-            ? isHalfScreen
-              ? (callersRef.current.getBoundingClientRect().width - 54) / 2
-              : callersRef.current.getBoundingClientRect().width - 16
-            : 0
-        }
-        metadataMappingFiles={metadataMappingFiles}
-        metadataLoading={false}
-        isSandwichIcicleGraph={true}
-        curPathArrow={curPathArrow}
-        setNewCurPathArrow={setCurPathArrow}
-        isFlamegraph={true}
-        profileSource={profileSource}
-        tooltipId="callers"
-      />
-    </div>
+    </>
   );
 }

package/src/Sandwich/components/TableSection.tsx

@@ -46,8 +46,8 @@ export function TableSection({
   return (
     <div
       style={{height: height !== undefined ? `${height}px` : '80vh'}}
-      className={`font-robotoMono w-full cursor-pointer ${
-        selectedRow != null && sandwichFunctionName !== undefined ? 'w-[50%]' : ''
+      className={`font-robotoMono cursor-pointer ${
+        selectedRow != null && sandwichFunctionName !== undefined ? 'w-[50%]' : 'w-full'
       }`}
     >
       <TableComponent

package/src/Sandwich/index.tsx

@@ -11,32 +11,28 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import React, {useCallback, useEffect, useMemo, useRef, useState} from 'react';
+import React, {useEffect, useMemo, useRef, useState} from 'react';
 
 import {type Row as TableRow} from '@tanstack/table-core';
 import {tableFromIPC} from 'apache-arrow';
 import {AnimatePresence, motion} from 'framer-motion';
 
 import {QueryRequest_ReportType, QueryServiceClient} from '@parca/client';
-import {TableSkeleton, useParcaContext, useURLState} from '@parca/components';
+import {useParcaContext, useURLState} from '@parca/components';
 import {useCurrentColorProfile} from '@parca/hooks';
 import {ProfileType} from '@parca/parser';
-import {isSearchMatch} from '@parca/utilities';
 
 import useMappingList, {
   useFilenamesList,
 } from '../ProfileIcicleGraph/IcicleGraphArrow/useMappingList';
 import {ProfileSource} from '../ProfileSource';
-import {useProfileViewContext} from '../ProfileView/context/ProfileViewContext';
+import {useDashboard} from '../ProfileView/context/DashboardContext';
 import {useVisualizationState} from '../ProfileView/hooks/useVisualizationState';
 import {FIELD_FUNCTION_NAME, Row} from '../Table';
 import {useColorManagement} from '../Table/hooks/useColorManagement';
-import {useTableConfiguration} from '../Table/hooks/useTableConfiguration';
-import {type DataRow} from '../Table/utils/functions';
 import {useQuery} from '../useQuery';
 import {CalleesSection} from './components/CalleesSection';
 import {CallersSection} from './components/CallersSection';
-import {TableSection} from './components/TableSection';
 import {processRowData} from './utils/processRowData';
 
 interface Props {
@@ -45,7 +41,6 @@ interface Props {
   filtered: bigint;
   profileType?: ProfileType;
   loading: boolean;
-  isHalfScreen: boolean;
   unit?: string;
   metadataMappingFiles?: string[];
   queryClient?: QueryServiceClient;
@@ -54,30 +49,26 @@ interface Props {
 
 const Sandwich = React.memo(function Sandwich({
   data,
-  total,
   filtered,
   profileType,
   loading,
-  isHalfScreen,
   unit,
   metadataMappingFiles,
   queryClient,
   profileSource,
 }: Props): React.JSX.Element {
   const currentColorProfile = useCurrentColorProfile();
+  const {dashboardItems} = useDashboard();
+  const [sandwichFunctionName] = useURLState<string | undefined>('sandwich_function_name');
 
-  const [sandwichFunctionName, setSandwichFunctionName] = useURLState<string | undefined>(
-    'sandwich_function_name'
-  );
   const {isDarkMode} = useParcaContext();
   const [selectedRow, setSelectedRow] = useState<TableRow<Row> | null>(null);
   const callersRef = React.useRef<HTMLDivElement | null>(null);
   const calleesRef = React.useRef<HTMLDivElement | null>(null);
+  const [isExpanded, setIsExpanded] = useState(false);
+  const defaultMaxFrames = 10;
 
   const callersCalleesContainerRef = useRef<HTMLDivElement | null>(null);
-  const [tableHeight, setTableHeight] = useState<number | undefined>(undefined);
-
-  const {compareMode} = useProfileViewContext();
 
   const {colorBy, setColorBy, curPathArrow, setCurPathArrow} = useVisualizationState();
 
@@ -160,15 +151,6 @@ const Sandwich = React.memo(function Sandwich({
 
   unit = useMemo(() => unit ?? profileType?.sampleUnit ?? '', [unit, profileType?.sampleUnit]);
 
-  const tableConfig = useTableConfiguration({
-    unit,
-    total,
-    filtered,
-    compareMode,
-  });
-
-  const {columns, initialSorting, columnVisibility} = tableConfig;
-
   const rows = useMemo(() => {
     if (table == null || table.numRows === 0) {
       return [];
@@ -194,71 +176,6 @@ const Sandwich = React.memo(function Sandwich({
     }
   }, [sandwichFunctionName, rows, selectedRow]);
 
-  // Update table height based on callers/callees container height
-  useEffect(() => {
-    const updateTableHeight = (): void => {
-      if (callersCalleesContainerRef.current != null) {
-        const containerHeight = callersCalleesContainerRef.current.getBoundingClientRect().height;
-        setTableHeight(containerHeight);
-      }
-    };
-
-    // Initial measurement
-    updateTableHeight();
-
-    // Update on window resize
-    window.addEventListener('resize', updateTableHeight);
-
-    // Use ResizeObserver if available for more accurate updates
-    let resizeObserver: ResizeObserver | null = null;
-    if (callersCalleesContainerRef.current != null && 'ResizeObserver' in window) {
-      resizeObserver = new ResizeObserver(updateTableHeight);
-      resizeObserver.observe(callersCalleesContainerRef.current);
-    }
-
-    return () => {
-      window.removeEventListener('resize', updateTableHeight);
-      if (resizeObserver != null) {
-        resizeObserver.disconnect();
-      }
-    };
-  }, [sandwichFunctionName, callersFlamegraphResponse, calleesFlamegraphResponse]);
-
-  const onRowClick = useCallback(
-    (row: DataRow) => {
-      setSelectedRow(row as unknown as TableRow<Row>);
-      setSandwichFunctionName(row.name.trim());
-    },
-    [setSandwichFunctionName]
-  );
-
-  const enableHighlighting = useMemo(() => {
-    return sandwichFunctionName != null && sandwichFunctionName?.length > 0;
-  }, [sandwichFunctionName]);
-
-  const shouldHighlightRow = useCallback(
-    (row: Row) => {
-      if (!('name' in row)) {
-        return false;
-      }
-      const name = row.name;
-      return isSearchMatch(sandwichFunctionName as string, name);
-    },
-    [sandwichFunctionName]
-  );
-
-  if (loading) {
-    return (
-      <div className="overflow-clip h-[700px] min-h-[700px]">
-        <TableSkeleton isHalfScreen={isHalfScreen} isDarkMode={isDarkMode} />
-      </div>
-    );
-  }
-
-  if (rows.length === 0) {
-    return <div className="mx-auto text-center">Profile has no samples</div>;
-  }
-
   return (
     <section className="flex flex-row h-full w-full">
       <AnimatePresence>
@@ -270,24 +187,10 @@ const Sandwich = React.memo(function Sandwich({
           transition={{duration: 0.5}}
         >
           <div className="relative flex flex-row">
-            <TableSection
-              rows={rows}
-              columns={columns}
-              initialSorting={initialSorting}
-              columnVisibility={columnVisibility}
-              selectedRow={selectedRow}
-              onRowClick={onRowClick}
-              shouldHighlightRow={shouldHighlightRow}
-              enableHighlighting={enableHighlighting}
-              height={tableHeight}
-              sandwichFunctionName={sandwichFunctionName}
-            />
-
-            {sandwichFunctionName !== undefined && (
-              <div className="w-[50%] flex flex-col" ref={callersCalleesContainerRef}>
+            {sandwichFunctionName !== undefined ? (
+              <div className="w-full flex flex-col" ref={callersCalleesContainerRef}>
                 <CallersSection
                   callersRef={callersRef}
-                  isHalfScreen={isHalfScreen}
                   callersFlamegraphResponse={
                     callersFlamegraphResponse?.report.oneofKind === 'flamegraphArrow'
                       ? {
@@ -306,11 +209,13 @@ const Sandwich = React.memo(function Sandwich({
                   curPathArrow={curPathArrow}
                   setCurPathArrow={setCurPathArrow}
                   metadataMappingFiles={metadataMappingFiles}
+                  isExpanded={isExpanded}
+                  setIsExpanded={setIsExpanded}
+                  defaultMaxFrames={defaultMaxFrames}
                 />
                 <div className="h-4" />
                 <CalleesSection
                   calleesRef={calleesRef}
-                  isHalfScreen={isHalfScreen}
                   calleesFlamegraphResponse={
                     calleesFlamegraphResponse?.report.oneofKind === 'flamegraphArrow'
                       ? {
@@ -331,6 +236,14 @@ const Sandwich = React.memo(function Sandwich({
                   metadataMappingFiles={metadataMappingFiles}
                 />
               </div>
+            ) : (
+              <div className="items-center justify-center flex h-full w-full">
+                <p className="text-sm">
+                  {dashboardItems.includes('table')
+                    ? 'Please select a function to view its callers and callees.'
+                    : 'Use the right-click menu on the flame graph to choose a function to view its callers and callees.'}
+                </p>
+              </div>
             )}
           </div>
         </motion.div>

package/src/SimpleMatchers/index.tsx

@@ -20,7 +20,7 @@ import cx from 'classnames';
 import {QueryServiceClient} from '@parca/client';
 import {useGrpcMetadata} from '@parca/components';
 import {Query} from '@parca/parser';
-import {sanitizeLabelValue} from '@parca/utilities';
+import {millisToProtoTimestamp, sanitizeLabelValue} from '@parca/utilities';
 
 import {LabelProvider, useLabels} from '../contexts/SimpleMatchersLabelContext';
 import {useUtilizationLabels} from '../contexts/UtilizationLabelsContext';
@@ -33,6 +33,8 @@ interface Props {
   currentQuery: Query;
   profileType: string;
   queryBrowserRef: React.RefObject<HTMLDivElement>;
+  start?: number;
+  end?: number;
 }
 
 interface QueryRow {
@@ -116,6 +118,8 @@ const SimpleMatchers = ({
   currentQuery,
   profileType,
   queryBrowserRef,
+  start,
+  end,
 }: Props): JSX.Element => {
   const utilizationLabels = useUtilizationLabels();
   const [queryRows, setQueryRows] = useState<QueryRow[]>([
@@ -140,10 +144,20 @@ const SimpleMatchers = ({
       }
       try {
         const values = await reactQueryClient.fetchQuery(
-          [labelName, profileType],
+          [labelName, profileType, start, end],
           async () => {
             const response = await queryClient.values(
-              {labelName, match: [], profileType},
+              {
+                labelName,
+                match: [],
+                profileType,
+                ...(start !== undefined && end !== undefined
+                  ? {
+                      start: millisToProtoTimestamp(start),
+                      end: millisToProtoTimestamp(end),
+                    }
+                  : {}),
+              },
               {meta: metadata}
             ).response;
             const sanitizedValues = sanitizeLabelValue(response.labelValues);
@@ -159,7 +173,7 @@ const SimpleMatchers = ({
         return [];
       }
     },
-    [queryClient, metadata, profileType, reactQueryClient]
+    [queryClient, metadata, profileType, reactQueryClient, start, end]
   );
 
   const fetchLabelValuesUtilization = useCallback(
@@ -419,6 +433,8 @@ export default function SimpleMathersWithProvider(props: Props): JSX.Element {
       queryClient={props.queryClient}
       profileType={props.profileType}
       labelNameFromMatchers={labelNameFromMatchers}
+      start={props.start}
+      end={props.end}
     >
       <SimpleMatchers {...props} />
     </LabelProvider>

package/src/Table/MoreDropdown.tsx

@@ -18,7 +18,6 @@ import {useParcaContext, useURLState} from '@parca/components';
 
 const MoreDropdown = ({functionName}: {functionName: string}): React.JSX.Element | null => {
   const [_, setSandwichFunctionName] = useURLState<string | undefined>('sandwich_function_name');
-  // eslint-disable-next-line @typescript-eslint/no-unused-vars
   const [dashboardItems, setDashboardItems] = useURLState<string[]>('dashboard_items', {
     alwaysReturnArray: true,
   });
@@ -26,7 +25,7 @@ const MoreDropdown = ({functionName}: {functionName: string}): React.JSX.Element
 
   const onSandwichViewSelect = (): void => {
     setSandwichFunctionName(functionName.trim());
-    setDashboardItems(['sandwich']);
+    setDashboardItems([...dashboardItems, 'sandwich']);
   };
 
   const menuItems: Array<{label: string; action: () => void}> = [];

package/src/Table/TableContextMenu.tsx

@@ -0,0 +1,70 @@
+// Copyright 2022 The Parca Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import {Icon} from '@iconify/react';
+import cx from 'classnames';
+import {Item, Menu} from 'react-contexify';
+
+import 'react-contexify/dist/ReactContexify.css';
+
+import {useParcaContext, useURLState} from '@parca/components';
+
+import {type Row} from '.';
+
+interface TableContextMenuProps {
+  menuId: string;
+  row: Row | null;
+}
+
+const TableContextMenu = ({menuId, row}: TableContextMenuProps): React.JSX.Element => {
+  const [_, setSandwichFunctionName] = useURLState<string | undefined>('sandwich_function_name');
+  const [dashboardItems, setDashboardItems] = useURLState<string[]>('dashboard_items', {
+    alwaysReturnArray: true,
+  });
+  const {enableSandwichView, isDarkMode} = useParcaContext();
+
+  const onSandwichViewSelect = (): void => {
+    if (row?.name != null && row.name.length > 0) {
+      setSandwichFunctionName(row.name.trim());
+      if (!dashboardItems.includes('sandwich')) {
+        setDashboardItems([...dashboardItems, 'sandwich']);
+      }
+    }
+  };
+
+  const isMenuDisabled = row === null || enableSandwichView !== true;
+
+  return (
+    <Menu
+      id={menuId}
+      theme={isDarkMode ? 'dark' : ''}
+      className={cx(
+        dashboardItems.includes('sandwich') ? 'min-w-[350px] w-[350px]' : 'min-w-[260px] w-[260px]'
+      )}
+    >
+      <Item id="sandwich-view" onClick={onSandwichViewSelect} disabled={isMenuDisabled}>
+        <div className="flex w-full items-center gap-2">
+          <Icon icon="tdesign:sandwich-filled" />
+          <div className="relative">
+            {dashboardItems.includes('sandwich')
+              ? 'Focus sandwich on this frame.'
+              : 'Show in sandwich'}
+            <span className="absolute top-[-2px] text-xs lowercase text-red-500">&nbsp;alpha</span>
+          </div>
+        </div>
+      </Item>
+    </Menu>
+  );
+};
+
+export default TableContextMenu;

package/src/Table/TableContextMenuWrapper.tsx

@@ -0,0 +1,48 @@
+// Copyright 2022 The Parca Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import {forwardRef, useImperativeHandle, useState} from 'react';
+
+import {type Row} from '.';
+import TableContextMenu from './TableContextMenu';
+
+interface TableContextMenuWrapperProps {
+  menuId: string;
+}
+
+export interface TableContextMenuWrapperRef {
+  setRow: (row: Row | null, callback?: () => void) => void;
+}
+
+const TableContextMenuWrapper = forwardRef<
+  TableContextMenuWrapperRef,
+  TableContextMenuWrapperProps
+>(({menuId}, ref) => {
+  const [row, setRow] = useState<Row | null>(null);
+
+  useImperativeHandle(ref, () => ({
+    setRow: (newRow: Row | null, callback?: () => void) => {
+      setRow(newRow);
+      // Execute callback after state update using requestAnimationFrame
+      if (callback != null) {
+        requestAnimationFrame(callback);
+      }
+    },
+  }));
+
+  return <TableContextMenu menuId={menuId} row={row} />;
+});
+
+TableContextMenuWrapper.displayName = 'TableContextMenuWrapper';
+
+export default TableContextMenuWrapper;