[Tabify] Handle doc_count with care on time shift scenarios (elastic#…

…178394)

## Summary

Fixes elastic#178073

This PR restructure a bit the `tabify` code to explain and handle the
count with time shift with extra care.

The fix itself is one line, but I thought it was worth expanding a
little bit more this obscure part of the code AND scope it down the fix
to a specific scenario. The `hasMultipleDocCountAtRootWithFilters` can
be perhaps removed once many more tests can be performed.

### Checklist
- [x]
[Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html)
was added for features that require explanation or tutorials
- [x] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios

src/plugins/data/common/search/aggs/agg_configs.ts

@@ -117,7 +117,7 @@ export class AggConfigs {
   isSamplingEnabled() {
     return (
       isSamplingEnabled(this.opts.probability) &&
-      this.getRequestAggs().filter((agg) => !agg.type.hasNoDsl).length > 0
+      this.getRequestAggs().some((agg) => !agg.type.hasNoDsl)
     );
   }
 

src/plugins/data/common/search/tabify/fixtures/fake_timeoffset_data.ts

@@ -0,0 +1,19 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+export const timeOffsetFiltersWithZeroDocCountResponse = {
+  hits: {
+    total: 474,
+    max_score: 0,
+    hits: [],
+  },
+  aggregations: {
+    doc_count: 0,
+    doc_count_86400000: 234,
+  },
+};

src/plugins/data/common/search/tabify/tabify.test.ts

@@ -12,15 +12,19 @@ import { AggConfigs, BucketAggParam, IAggConfig, IAggConfigs } from '../aggs';
 import { mockAggTypesRegistry } from '../aggs/test_helpers';
 import { metricOnly, threeTermBuckets } from './fixtures/fake_hierarchical_data';
 import { isSamplingEnabled } from '../aggs/utils/sampler';
+import { timeOffsetFiltersWithZeroDocCountResponse } from './fixtures/fake_timeoffset_data';
 
 describe('tabifyAggResponse Integration', () => {
   const typesRegistry = mockAggTypesRegistry();
 
   for (const probability of [1, 0.5, undefined]) {
     function getTitlePostfix() {
-      if (!isSamplingEnabled(probability)) {
+      if (probability == null) {
         return '';
       }
+      if (probability === 1) {
+        return ` - with no sampling (probability = 1)`;
+      }
       return ` - with sampling (probability = ${probability})`;
     }
 
@@ -243,5 +247,20 @@ describe('tabifyAggResponse Integration', () => {
         });
       });
     });
+
+    describe(`edge cases${getTitlePostfix()}`, () => {
+      test('it should correctly report zero doc count for unshifted bucket', () => {
+        const aggConfigs = createAggConfigs([
+          mockAggConfig({ type: 'count', schema: 'metric' }),
+          mockAggConfig({ type: 'count', schema: 'metric', params: { timeShift: '1d' } }),
+        ]);
+
+        // no need to wrap with sampling as count is not affected by it
+        const tabbed = tabifyAggResponse(aggConfigs, timeOffsetFiltersWithZeroDocCountResponse, {
+          metricsAtAllLevels: false,
+        });
+        expect(tabbed.rows[0]).toEqual({ 'col-0-1': 0, 'col-1-2': 234 });
+      });
+    });
   }
 });

src/plugins/data/common/search/tabify/tabify.ts

@@ -146,13 +146,27 @@ export function tabifyAggResponse(
   }
 
   const write = new TabbedAggResponseWriter(aggConfigs, respOpts || {});
+  // Check whether there's a time shift for a count operation at root level
+  const hasMultipleDocCountAtRootWithFilters = Object.keys(esResponse.aggregations ?? {}).some(
+    (key) => /doc_count_/.test(key)
+  );
+
   const topLevelBucket: AggResponseBucket = {
     ...(aggConfigs.isSamplingEnabled()
       ? esResponse.aggregations?.sampling
       : esResponse.aggregations),
-    doc_count: esResponse.aggregations?.doc_count || esResponse.hits?.total,
+    doc_count: esResponse.aggregations?.doc_count,
   };
 
+  // The fix itself is one line, but it's a bit hard to clear assess the full impact of it
+  // therefore here's a check to scope down the impact to the known scenario with the bug https://github.com/elastic/kibana/issues/178073
+  // this can be lifted off once the full impact is assessed
+  if (!topLevelBucket.doc_count) {
+    if (!hasMultipleDocCountAtRootWithFilters) {
+      topLevelBucket.doc_count = esResponse.hits?.total;
+    }
+  }
+
   collectBucket(aggConfigs, write, topLevelBucket, '', 1);
 
   return {