fix(heatmap): pick correct brush end value

packages/charts/src/chart_types/heatmap/layout/viewmodel/viewmodel.ts

@@ -275,25 +275,10 @@ export function shapeViewModel(
     const startY = yInvertedScale(clamp(topLeft[1], 0, currentGridHeight - 1));
     const endY = yInvertedScale(clamp(bottomRight[1], 0, currentGridHeight - 1));
 
-    let allXValuesInRange: Array<NonNullable<PrimitiveValue>> = [];
-    const invertedXValues: Array<NonNullable<PrimitiveValue>> = [];
-
-    if (timeScale && typeof endX === 'number') {
-      invertedXValues.push(startX);
-      invertedXValues.push(endX + xDomain.minInterval);
-      let [startXValue] = invertedXValues;
-      if (typeof startXValue === 'number') {
-        while (startXValue < invertedXValues[1]) {
-          allXValuesInRange.push(startXValue);
-          startXValue += xDomain.minInterval;
-        }
-      }
-    } else {
-      allXValuesInRange = getValuesInRange(xValues, startX, endX);
-      invertedXValues.push(...allXValuesInRange);
-    }
-
+    const allXValuesInRange: Array<NonNullable<PrimitiveValue>> = getValuesInRange(xValues, startX, endX);
     const allYValuesInRange: Array<NonNullable<PrimitiveValue>> = getValuesInRange(yValues, startY, endY);
+    const invertedXValues: Array<NonNullable<PrimitiveValue>> =
+      timeScale && typeof endX === 'number' ? [startX, endX + xDomain.minInterval] : [...allXValuesInRange];
 
     const cells: Cell[] = [];
 
@@ -325,7 +310,7 @@ export function shapeViewModel(
 
     // find X coordinated based on the time range
     const leftIndex = typeof startValue === 'number' ? bisectLeft(xValues, startValue) : xValues.indexOf(startValue);
-    const rightIndex = typeof endValue === 'number' ? bisectLeft(xValues, endValue) : xValues.indexOf(endValue);
+    const rightIndex = typeof endValue === 'number' ? bisectLeft(xValues, endValue) : xValues.indexOf(endValue) + 1;
 
     const isRightOutOfRange = rightIndex > xValues.length - 1 || rightIndex < 0;
     const isLeftOutOfRange = leftIndex > xValues.length - 1 || leftIndex < 0;
@@ -340,7 +325,7 @@ export function shapeViewModel(
     const xStart = chartDimensions.left + startFromScale;
 
     // extend the range in case the right boundary has been selected
-    const width = endFromScale - startFromScale + cellWidth; // (isRightOutOfRange || isLeftOutOfRange ? cellWidth : 0);
+    const width = endFromScale - startFromScale + (isRightOutOfRange || isLeftOutOfRange ? cellWidth : 0);
 
     // resolve Y coordinated making sure the order is correct
     const { y: yStart, totalHeight } = y

packages/charts/src/chart_types/heatmap/state/selectors/get_brushed_highlighted_shapes.test.ts

@@ -17,6 +17,7 @@
  * under the License.
  */
 
+import { DateTime } from 'luxon';
 import { Store } from 'redux';
 
 import { MockGlobalSpec, MockSeriesSpec } from '../../../../mocks/specs/specs';
@@ -26,7 +27,7 @@ import { onMouseDown, onMouseUp, onPointerMove } from '../../../../state/actions
 import { GlobalChartState } from '../../../../state/chart_state';
 import { createOnBrushEndCaller } from './on_brush_end_caller';
 
-describe('Heatmap brush', () => {
+describe('Categorical heatmap brush', () => {
   let store: Store<GlobalChartState>;
   let onBrushEndMock = jest.fn();
 
@@ -87,3 +88,79 @@ describe('Heatmap brush', () => {
     expect(brushEvent.y).toEqual(['ya', 'yb', 'yc']);
   });
 });
+describe('Temporal heatmap brush', () => {
+  let store: Store<GlobalChartState>;
+  let onBrushEndMock = jest.fn();
+  const start = DateTime.fromISO('2021-07-01T00:00:00.000Z');
+  beforeEach(() => {
+    store = MockStore.default({ width: 300, height: 300, top: 0, left: 0 }, 'chartId');
+    onBrushEndMock = jest.fn();
+    MockStore.addSpecs(
+      [
+        MockGlobalSpec.settingsNoMargins(),
+        MockSeriesSpec.heatmap({
+          xScaleType: ScaleType.Time,
+          data: [
+            { x: start.toMillis(), y: 'ya', value: 1 },
+            { x: start.plus({ days: 1 }).toMillis(), y: 'ya', value: 2 },
+            { x: start.plus({ days: 2 }).toMillis(), y: 'ya', value: 3 },
+            { x: start.toMillis(), y: 'yb', value: 4 },
+            { x: start.plus({ days: 1 }).toMillis(), y: 'yb', value: 5 },
+            { x: start.plus({ days: 2 }).toMillis(), y: 'yb', value: 6 },
+            { x: start.toMillis(), y: 'yc', value: 7 },
+            { x: start.plus({ days: 1 }).toMillis(), y: 'yc', value: 8 },
+            { x: start.plus({ days: 2 }).toMillis(), y: 'yc', value: 9 },
+          ],
+          config: {
+            grid: {
+              cellHeight: {
+                max: 'fill',
+              },
+              cellWidth: {
+                max: 'fill',
+              },
+            },
+            xAxisLabel: {
+              visible: false,
+            },
+            yAxisLabel: {
+              visible: false,
+            },
+            margin: { top: 0, bottom: 0, left: 0, right: 0 },
+            onBrushEnd: onBrushEndMock,
+          },
+        }),
+      ],
+      store,
+    );
+  });
+
+  it('should brush on the x scale + minInterval', () => {
+    const caller = createOnBrushEndCaller();
+    store.dispatch(onPointerMove({ x: 50, y: 50 }, 0));
+    store.dispatch(onMouseDown({ x: 50, y: 50 }, 100));
+    store.dispatch(onPointerMove({ x: 250, y: 250 }, 200));
+    store.dispatch(onMouseUp({ x: 250, y: 250 }, 300));
+    caller(store.getState());
+    expect(onBrushEndMock).toBeCalledTimes(1);
+    const brushEvent = onBrushEndMock.mock.calls[0][0];
+    expect(brushEvent.cells).toHaveLength(6);
+    // it covers from the beginning of the cell to the end of the next cell
+    expect(brushEvent.x).toEqual([start.toMillis(), start.plus({ days: 2 }).toMillis()]);
+    expect(brushEvent.y).toEqual(['ya', 'yb', 'yc']);
+  });
+  it('should brush on the x scale + minInterval on a single cell', () => {
+    const caller = createOnBrushEndCaller();
+    store.dispatch(onPointerMove({ x: 50, y: 50 }, 0));
+    store.dispatch(onMouseDown({ x: 50, y: 50 }, 100));
+    store.dispatch(onPointerMove({ x: 60, y: 60 }, 200));
+    store.dispatch(onMouseUp({ x: 60, y: 60 }, 300));
+    caller(store.getState());
+    expect(onBrushEndMock).toBeCalledTimes(1);
+    const brushEvent = onBrushEndMock.mock.calls[0][0];
+    expect(brushEvent.cells).toHaveLength(1);
+    // it covers from the beginning of the cell to the end of the next cell
+    expect(brushEvent.x).toEqual([start.toMillis(), start.plus({ days: 1 }).toMillis()]);
+    expect(brushEvent.y).toEqual(['ya']);
+  });
+});

packages/charts/src/chart_types/heatmap/state/selectors/get_picked_cells.ts

@@ -30,18 +30,6 @@ export const getPickedCells = createCustomCachedSelector(
       return null;
     }
 
-    const {
-      start: {
-        position: { x: startX, y: startY },
-      },
-      end: {
-        position: { x: endX, y: endY },
-      },
-    } = dragState;
-
-    return geoms.pickDragArea([
-      { x: startX, y: startY },
-      { x: endX, y: endY },
-    ]);
+    return geoms.pickDragArea([dragState.start.position, dragState.end.position]);
   },
 );