[DataGrid] Keep focused cell in the DOM

packages/grid/x-data-grid-pro/src/hooks/features/rowReorder/useGridRowReorder.tsx

@@ -81,6 +81,7 @@ export const useGridRowReorder = (
       });
 
       originRowIndex.current = apiRef.current.getRowIndexRelativeToVisibleRows(params.id);
+      apiRef.current.setCellFocus(params.id, '__reorder__');
     },
     [isRowReorderDisabled, classes.rowDragging, logger, apiRef],
   );

packages/grid/x-data-grid/src/components/GridRow.tsx

@@ -450,11 +450,14 @@ const GridRow = React.forwardRef<
   const rowType = apiRef.current.getRowNode(rowId)!.type;
   const cells: JSX.Element[] = [];
 
+  const isAdditionalColumnAdded = renderedColumns.length < lastColumnToRender - firstColumnToRender;
   for (let i = 0; i < renderedColumns.length; i += 1) {
     const column = renderedColumns[i];
     const indexRelativeToAllColumns = firstColumnToRender + i;
 
-    const isLastColumn = indexRelativeToAllColumns === visibleColumns.length - 1;
+    const isLastColumn = isAdditionalColumnAdded
+      ? indexRelativeToAllColumns === visibleColumns.length
+      : indexRelativeToAllColumns === visibleColumns.length - 1;
     const removeLastBorderRight = isLastColumn && hasScrollX && !hasScrollY;
     const showRightBorder = !isLastColumn
       ? rootProps.showCellVerticalBorder

packages/grid/x-data-grid/src/hooks/features/virtualization/useGridVirtualScroller.tsx

@@ -418,6 +418,9 @@ export const useGridVirtualScroller = (props: UseGridVirtualScrollerProps) => {
     const rowBuffer = !disableVirtualization ? rootProps.rowBuffer : 0;
     const columnBuffer = !disableVirtualization ? rootProps.columnBuffer : 0;
 
+    let isFocusedCellRendered = false;
+    const { cell } = apiRef.current.state.focus;
+
     const [firstRowToRender, lastRowToRender] = getRenderableIndexes({
       firstIndex: nextRenderContext.firstRowIndex,
       lastIndex: nextRenderContext.lastRowIndex,
@@ -445,6 +448,10 @@ export const useGridVirtualScroller = (props: UseGridVirtualScrollerProps) => {
 
       for (let i = firstRowToRender; i < lastRowToRender; i += 1) {
         const row = currentPage.rows[i];
+        if (cell && row.id === cell?.id) {
+          isFocusedCellRendered = true;
+        }
+
         renderedRows.push(row);
         apiRef.current.calculateColSpan({
           rowId: row.id,
@@ -453,6 +460,29 @@ export const useGridVirtualScroller = (props: UseGridVirtualScrollerProps) => {
           columns: visibleColumns,
         });
       }
+
+      // If the selected row is not within the current range of rows being displayed,
+      // we need to render it at either the top or bottom of the rows,
+      // depending on whether it is above or below the range.
+      if (!isFocusedCellRendered && cell && !disableVirtualization) {
+        const rows = currentPage.rows.filter((row) => row.id === cell.id);
+        const focusedRow = rows.length > 0 && rows[0];
+
+        if (focusedRow) {
+          const index = currentPage.rows.indexOf(focusedRow);
+          if (index > firstRowToRender) {
+            renderedRows.push(focusedRow);
+          } else {
+            renderedRows.unshift(focusedRow);
+          }
+          apiRef.current.calculateColSpan({
+            rowId: focusedRow.id,
+            minFirstColumn,
+            maxLastColumn,
+            columns: visibleColumns,
+          });
+        }
+      }
     }
 
     const [initialFirstColumnToRender, lastColumnToRender] = getRenderableIndexes({
@@ -473,11 +503,34 @@ export const useGridVirtualScroller = (props: UseGridVirtualScrollerProps) => {
 
     const renderedColumns = visibleColumns.slice(firstColumnToRender, lastColumnToRender);
 
+    // If the selected column is not within the current range of columns being displayed,
+    // we need to render it at either the left or right of the columns,
+    // depending on whether it is above or below the range.
+    const isAllColumnRendered =
+      maxLastColumn - minFirstColumn === lastColumnToRender - firstColumnToRender;
+    if (cell && !disableVirtualization && !isAllColumnRendered) {
+      const selectedColumns = visibleColumns.filter((column) => column.field === cell.field);
+      if (selectedColumns.length > 0) {
+        const focusedColumn = selectedColumns[0];
+        const focusedColumnIndex = visibleColumns.indexOf(focusedColumn);
+
+        if (firstColumnToRender > focusedColumnIndex) {
+          renderedColumns.unshift(focusedColumn);
+        } else if (lastColumnToRender < focusedColumnIndex) {
+          renderedColumns.push(focusedColumn);
+        }
+      }
+    }
+
     const rows: JSX.Element[] = [];
 
     for (let i = 0; i < renderedRows.length; i += 1) {
       const { id, model } = renderedRows[i];
-      const lastVisibleRowIndex = firstRowToRender + i === currentPage.rows.length - 1;
+
+      const lastVisibleRowIndex =
+        !isFocusedCellRendered && cell
+          ? firstRowToRender + i === currentPage.rows.length
+          : firstRowToRender + i === currentPage.rows.length - 1;
       const baseRowHeight = !apiRef.current.rowHasAutoHeight(id)
         ? apiRef.current.unstable_getRowHeight(id)
         : 'auto';