diff --git a/src/chart_types/heatmap/layout/types/viewmodel_types.ts b/src/chart_types/heatmap/layout/types/viewmodel_types.ts
index cc60ee38a6..6c29339bec 100644
--- a/src/chart_types/heatmap/layout/types/viewmodel_types.ts
+++ b/src/chart_types/heatmap/layout/types/viewmodel_types.ts
@@ -66,7 +66,7 @@ export interface HeatmapViewModel {
 export type PickFunction = (x: Pixels, y: Pixels) => Cell[] | TextBox;
 
 /** @internal */
-export type PickDragFunction = (points: [Point, Point]) => Cell[];
+export type PickDragFunction = (points: [Point, Point]) => { cells: Cell[]; x: any[]; y: any[] };
 
 /** @internal */
 export type PickDragShapeFunction = (
@@ -107,6 +107,6 @@ export const nullShapeViewModel = (specifiedConfig?: Config): ShapeViewModel =>
   config: specifiedConfig || config,
   heatmapViewModel: nullHeatmapViewModel,
   pickQuads: () => [],
-  pickDragArea: () => [],
+  pickDragArea: () => ({ cells: [], x: [], y: [] }),
   pickDragShape: () => {},
 });
diff --git a/src/chart_types/heatmap/layout/viewmodel/viewmodel.ts b/src/chart_types/heatmap/layout/viewmodel/viewmodel.ts
index d3c8300e1b..4b3146abdd 100644
--- a/src/chart_types/heatmap/layout/viewmodel/viewmodel.ts
+++ b/src/chart_types/heatmap/layout/viewmodel/viewmodel.ts
@@ -232,40 +232,6 @@ export function shapeViewModel(
     }
     return [];
   };
-
-  /**
-   * Return selected cells based on drag selection.
-   * @param start
-   * @param end
-   */
-  const pickDragArea: PickDragFunction = ([start, end]) => {
-    const result: Cell[] = [];
-
-    const startX = Math.min(start.x, end.x);
-    const startY = Math.min(start.y, end.y);
-
-    const endX = Math.max(start.x, end.x);
-    const endY = Math.max(start.y, end.y);
-
-    const [startPoint] = pickQuads(startX, startY);
-    result.push(startPoint);
-
-    let { x, y } = startPoint;
-    x += cellWidth + maxTextWidth;
-
-    while (y <= endY) {
-      while (x <= endX) {
-        result.push(pickQuads(x, y)[0]);
-        x += cellWidth;
-      }
-      // move to the next line
-      x = startPoint.x + maxTextWidth;
-      y += cellHeight;
-    }
-
-    return result;
-  };
-
   /**
    * Resolves coordinates and metrics of the selected rect area.
    * @param start
@@ -292,6 +258,44 @@ export function shapeViewModel(
     };
   };
 
+  /**
+   * Return selected cells based on drag selection.
+   */
+  const pickDragArea: PickDragFunction = (bound) => {
+    const result = {
+      cells: [] as Cell[],
+      x: new Set(),
+      y: new Set(),
+    };
+
+    const shape = pickDragShape(bound);
+
+    let { x, y } = shape;
+
+    while (y < shape.height + shape.y) {
+      result.y.add(yInvertedScale(y));
+      while (x < shape.width + shape.x) {
+        const [cell] = pickQuads(x, y);
+        if (cell) {
+          result.cells.push(cell);
+        }
+        result.x.add(xInvertedScale(x));
+        x += cellWidth;
+      }
+      // move to the next line
+      x = shape.x;
+      y += cellHeight;
+    }
+
+    const x = [...result.x];
+
+    return {
+      ...result,
+      x: isXAxisTimeScale ? [x[0], x[x.length - 1]] : x,
+      y: [...result.y],
+    };
+  };
+
   // vertical lines
   const xLines = [];
   for (let i = 0; i < xValues.length + 1; i++) {
diff --git a/src/chart_types/heatmap/state/selectors/on_brush_end_caller.ts b/src/chart_types/heatmap/state/selectors/on_brush_end_caller.ts
index afffff3d8b..15d2f4af5d 100644
--- a/src/chart_types/heatmap/state/selectors/on_brush_end_caller.ts
+++ b/src/chart_types/heatmap/state/selectors/on_brush_end_caller.ts
@@ -43,8 +43,8 @@ export function createOnBrushEndCaller(): (state: GlobalChartState) => void {
     if (selector === null && state.chartType === ChartTypes.Heatmap) {
       selector = createCachedSelector(
         [getSpecOrNull, getSettingsSpecSelector, getPickedCells],
-        (spec, { onBrushEnd, rotation, brushAxis, minBrushDelta }, pickedCells): void => {
-          if (!spec || !onBrushEnd || pickedCells.length === 0) {
+        (spec, { onBrushEnd }, pickedCells): void => {
+          if (!spec || !onBrushEnd || pickedCells === null) {
             return;
           }
 
diff --git a/src/chart_types/heatmap/state/selectors/picked_shapes.ts b/src/chart_types/heatmap/state/selectors/picked_shapes.ts
index 12eee58634..39c234fc8c 100644
--- a/src/chart_types/heatmap/state/selectors/picked_shapes.ts
+++ b/src/chart_types/heatmap/state/selectors/picked_shapes.ts
@@ -22,7 +22,7 @@ import createCachedSelector from 're-reselect';
 import { GlobalChartState } from '../../../../state/chart_state';
 import { getChartIdSelector } from '../../../../state/selectors/get_chart_id';
 import { getLastDragSelector } from '../../../../state/selectors/get_last_drag';
-import { Cell } from '../../layout/types/viewmodel_types';
+import { Cell, PickDragFunction } from '../../layout/types/viewmodel_types';
 import { TextBox } from '../../layout/viewmodel/viewmodel';
 import { geometries } from './geometries';
 
@@ -40,9 +40,11 @@ export const getPickedShapes = createCachedSelector([geometries, getCurrentPoint
 })(getChartIdSelector);
 
 /** @internal */
-export const getPickedCells = createCachedSelector([geometries, getLastDragSelector], (geoms, dragState): Cell[] => {
+export const getPickedCells = createCachedSelector([geometries, getLastDragSelector], (geoms, dragState): ReturnType<
+  PickDragFunction
+> | null => {
   if (!dragState) {
-    return [];
+    return null;
   }
 
   const {