[Discover][New table] Fix selected documents no longer in the sample …

…of 500 documents (elastic#99114) (elastic#99800)

* Don't count and make use of selected docs no longer in the given 500 sample of the result set

* When the underlying dataset changes, set selected docs filter to inactive if there are no filtered docs

* Add test

src/plugins/discover/public/application/components/discover_grid/discover_grid.test.tsx

@@ -124,6 +124,33 @@ describe('DiscoverGrid', () => {
       expect(getDisplayedDocNr(component)).toBe(5);
     });
 
+    test('showing selected documents, underlying data changes, all documents are displayed, selection is gone', async () => {
+      await toggleDocSelection(component, esHits[0]);
+      await toggleDocSelection(component, esHits[1]);
+      expect(getSelectedDocNr(component)).toBe(2);
+      findTestSubject(component, 'dscGridSelectionBtn').simulate('click');
+      findTestSubject(component, 'dscGridShowSelectedDocuments').simulate('click');
+      expect(getDisplayedDocNr(component)).toBe(2);
+      component.setProps({
+        rows: [
+          {
+            _index: 'i',
+            _id: '6',
+            _score: 1,
+            _type: '_doc',
+            _source: {
+              date: '2020-20-02T12:12:12.128',
+              name: 'test6',
+              extension: 'doc',
+              bytes: 50,
+            },
+          },
+        ],
+      });
+      expect(getDisplayedDocNr(component)).toBe(1);
+      expect(getSelectedDocNr(component)).toBe(0);
+    });
+
     test('showing only selected documents and remove filter deselecting each doc manually', async () => {
       await toggleDocSelection(component, esHits[0]);
       findTestSubject(component, 'dscGridSelectionBtn').simulate('click');

src/plugins/discover/public/application/components/discover_grid/discover_grid.tsx

@@ -164,17 +164,33 @@ export const DiscoverGrid = ({
   const [isFilterActive, setIsFilterActive] = useState(false);
   const displayedColumns = getDisplayedColumns(columns, indexPattern);
   const defaultColumns = displayedColumns.includes('_source');
+  const usedSelectedDocs = useMemo(() => {
+    if (!selectedDocs.length || !rows?.length) {
+      return [];
+    }
+    const idMap = rows.reduce((map, row) => map.set(getDocId(row), true), new Map());
+    // filter out selected docs that are no longer part of the current data
+    const result = selectedDocs.filter((docId) => idMap.get(docId));
+    if (result.length === 0 && isFilterActive) {
+      setIsFilterActive(false);
+    }
+    return result;
+  }, [selectedDocs, rows, isFilterActive]);
+
   const displayedRows = useMemo(() => {
     if (!rows) {
       return [];
     }
-    if (!isFilterActive || selectedDocs.length === 0) {
+    if (!isFilterActive || usedSelectedDocs.length === 0) {
+      return rows;
+    }
+    const rowsFiltered = rows.filter((row) => usedSelectedDocs.includes(getDocId(row)));
+    if (!rowsFiltered.length) {
+      // in case the selected docs are no longer part of the sample of 500, show all docs
       return rows;
     }
-    return rows.filter((row) => {
-      return selectedDocs.includes(getDocId(row));
-    });
-  }, [rows, selectedDocs, isFilterActive]);
+    return rowsFiltered;
+  }, [rows, usedSelectedDocs, isFilterActive]);
 
   /**
    * Pagination
@@ -258,16 +274,16 @@ export const DiscoverGrid = ({
 
   const additionalControls = useMemo(
     () =>
-      selectedDocs.length ? (
+      usedSelectedDocs.length ? (
         <DiscoverGridDocumentToolbarBtn
           isFilterActive={isFilterActive}
           rows={rows!}
-          selectedDocs={selectedDocs}
+          selectedDocs={usedSelectedDocs}
           setSelectedDocs={setSelectedDocs}
           setIsFilterActive={setIsFilterActive}
         />
       ) : null,
-    [selectedDocs, isFilterActive, rows, setIsFilterActive]
+    [usedSelectedDocs, isFilterActive, rows, setIsFilterActive]
   );
 
   if (!rowCount) {
@@ -291,7 +307,7 @@ export const DiscoverGrid = ({
         onFilter,
         indexPattern,
         isDarkMode: services.uiSettings.get('theme:darkMode'),
-        selectedDocs,
+        selectedDocs: usedSelectedDocs,
         setSelectedDocs: (newSelectedDocs) => {
           setSelectedDocs(newSelectedDocs);
           if (isFilterActive && newSelectedDocs.length === 0) {