[Lens] Random sampling feature

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

@@ -16,6 +16,15 @@ import type { DataView } from '@kbn/data-views-plugin/common';
 import { stubIndexPattern } from '../../stubs';
 import { IEsSearchResponse } from '..';
 
+// Mute moment.tz warnings about not finding a mock timezone
+jest.mock('../utils', () => {
+  const original = jest.requireActual('../utils');
+  return {
+    ...original,
+    getUserTimeZone: jest.fn(() => 'US/Pacific'),
+  };
+});
+
 describe('AggConfigs', () => {
   const indexPattern: DataView = stubIndexPattern;
   let typesRegistry: AggTypesRegistryStart;
@@ -563,6 +572,82 @@ describe('AggConfigs', () => {
         '1-bucket>_count'
       );
     });
+
+    it('prepends a sampling agg whenever sampling is enabled', () => {
+      const configStates = [
+        {
+          enabled: true,
+          id: '1',
+          type: 'avg_bucket',
+          schema: 'metric',
+          params: {
+            customBucket: {
+              id: '1-bucket',
+              type: 'date_histogram',
+              schema: 'bucketAgg',
+              params: {
+                field: '@timestamp',
+                interval: '10s',
+              },
+            },
+            customMetric: {
+              id: '1-metric',
+              type: 'count',
+              schema: 'metricAgg',
+              params: {},
+            },
+          },
+        },
+        {
+          enabled: true,
+          id: '2',
+          type: 'terms',
+          schema: 'bucket',
+          params: {
+            field: 'clientip',
+          },
+        },
+        {
+          enabled: true,
+          id: '3',
+          type: 'terms',
+          schema: 'bucket',
+          params: {
+            field: 'machine.os.raw',
+          },
+        },
+      ];
+
+      const ac = new AggConfigs(
+        indexPattern,
+        configStates,
+        { typesRegistry, hierarchical: true, probability: 0.5 },
+        jest.fn()
+      );
+      const topLevelDsl = ac.toDsl();
+
+      expect(Object.keys(topLevelDsl)).toContain('sampling');
+      expect(Object.keys(topLevelDsl.sampling)).toEqual(['random_sampler', 'aggs']);
+      expect(Object.keys(topLevelDsl.sampling.aggs)).toContain('2');
+      expect(Object.keys(topLevelDsl.sampling.aggs['2'].aggs)).toEqual(['1', '3', '1-bucket']);
+    });
+
+    it('should not prepend a sampling agg when no nested agg is avaialble', () => {
+      const ac = new AggConfigs(
+        indexPattern,
+        [
+          {
+            enabled: true,
+            type: 'count',
+            schema: 'metric',
+          },
+        ],
+        { typesRegistry, probability: 0.5 },
+        jest.fn()
+      );
+      const topLevelDsl = ac.toDsl();
+      expect(Object.keys(topLevelDsl)).not.toContain('sampling');
+    });
   });
 
   describe('#postFlightTransform', () => {
@@ -854,4 +939,74 @@ describe('AggConfigs', () => {
       `);
     });
   });
+
+  describe('isSamplingEnabled', () => {
+    it('should return false if probability is 1', () => {
+      const ac = new AggConfigs(
+        indexPattern,
+        [{ enabled: true, type: 'avg', schema: 'metric', params: { field: 'bytes' } }],
+        { typesRegistry, probability: 1 },
+        jest.fn()
+      );
+
+      expect(ac.isSamplingEnabled()).toBeFalsy();
+    });
+
+    it('should return true if probability is less than 1', () => {
+      const ac = new AggConfigs(
+        indexPattern,
+        [{ enabled: true, type: 'avg', schema: 'metric', params: { field: 'bytes' } }],
+        { typesRegistry, probability: 0.1 },
+        jest.fn()
+      );
+
+      expect(ac.isSamplingEnabled()).toBeTruthy();
+    });
+
+    it('should return false when all aggs have hasNoDsl flag enabled', () => {
+      const ac = new AggConfigs(
+        indexPattern,
+        [
+          {
+            enabled: true,
+            type: 'count',
+            schema: 'metric',
+          },
+        ],
+        { typesRegistry, probability: 1 },
+        jest.fn()
+      );
+
+      expect(ac.isSamplingEnabled()).toBeFalsy();
+    });
+
+    it('should return false when no nested aggs are avaialble', () => {
+      const ac = new AggConfigs(
+        indexPattern,
+        [{ enabled: false, type: 'avg', schema: 'metric', params: { field: 'bytes' } }],
+        { typesRegistry, probability: 1 },
+        jest.fn()
+      );
+
+      expect(ac.isSamplingEnabled()).toBeFalsy();
+    });
+
+    it('should return true if at least one nested agg is available and probability < 1', () => {
+      const ac = new AggConfigs(
+        indexPattern,
+        [
+          {
+            enabled: true,
+            type: 'count',
+            schema: 'metric',
+          },
+          { enabled: true, type: 'avg', schema: 'metric', params: { field: 'bytes' } },
+        ],
+        { typesRegistry, probability: 0.1 },
+        jest.fn()
+      );
+
+      expect(ac.isSamplingEnabled()).toBeTruthy();
+    });
+  });
 });

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

@@ -26,6 +26,7 @@ import { AggTypesDependencies, GetConfigFn, getUserTimeZone } from '../..';
 import { getTime, calculateBounds } from '../..';
 import type { IBucketAggConfig } from './buckets';
 import { insertTimeShiftSplit, mergeTimeShifts } from './utils/time_splits';
+import { createSamplerAgg, isSamplingEnabled } from './utils/sampler';
 
 function removeParentAggs(obj: any) {
   for (const prop in obj) {
@@ -55,6 +56,8 @@ export interface AggConfigsOptions {
   hierarchical?: boolean;
   aggExecutionContext?: AggTypesDependencies['aggExecutionContext'];
   partialRows?: boolean;
+  probability?: number;
+  samplerSeed?: number;
 }
 
 export type CreateAggConfigParams = Assign<AggConfigSerialized, { type: string | IAggType }>;
@@ -107,6 +110,17 @@ export class AggConfigs {
     return this.opts.partialRows ?? false;
   }
 
+  public get samplerConfig() {
+    return { probability: this.opts.probability ?? 1, seed: this.opts.samplerSeed };
+  }
+
+  isSamplingEnabled() {
+    return (
+      isSamplingEnabled(this.opts.probability) &&
+      this.getRequestAggs().filter((agg) => !agg.type.hasNoDsl).length > 0
+    );
+  }
+
   setTimeFields(timeFields: string[] | undefined) {
     this.timeFields = timeFields;
   }
@@ -225,7 +239,7 @@ export class AggConfigs {
   }
 
   toDsl(): Record<string, any> {
-    const dslTopLvl = {};
+    const dslTopLvl: Record<string, any> = {};
     let dslLvlCursor: Record<string, any>;
     let nestedMetrics: Array<{ config: AggConfig; dsl: Record<string, any> }> | [];
 
@@ -254,10 +268,21 @@ export class AggConfigs {
       (config) => 'splitForTimeShift' in config.type && config.type.splitForTimeShift(config, this)
     );
 
+    if (this.isSamplingEnabled()) {
+      dslTopLvl.sampling = createSamplerAgg({
+        probability: this.opts.probability ?? 1,
+        seed: this.opts.samplerSeed,
+      });
+    }
+
     requestAggs.forEach((config: AggConfig, i: number, list) => {
       if (!dslLvlCursor) {
         // start at the top level
         dslLvlCursor = dslTopLvl;
+        // when sampling jump directly to the aggs
+        if (this.isSamplingEnabled()) {
+          dslLvlCursor = dslLvlCursor.sampling.aggs;
+        }
       } else {
         const prevConfig: AggConfig = list[i - 1];
         const prevDsl = dslLvlCursor[prevConfig.id];
@@ -452,7 +477,12 @@ export class AggConfigs {
         doc_count: response.rawResponse.hits?.total as estypes.AggregationsAggregate,
       };
     }
-    const aggCursor = transformedRawResponse.aggregations!;
+    const aggCursor = this.isSamplingEnabled()
+      ? (transformedRawResponse.aggregations!.sampling! as Record<
+          string,
+          estypes.AggregationsAggregate
+        >)
+      : transformedRawResponse.aggregations!;
 
     mergeTimeShifts(this, aggCursor);
     return {
@@ -531,6 +561,8 @@ export class AggConfigs {
         metricsAtAllLevels: this.hierarchical,
         partialRows: this.partialRows,
         aggs: this.aggs.map((agg) => buildExpression(agg.toExpressionAst())),
+        probability: this.opts.probability,
+        samplerSeed: this.opts.samplerSeed,
       }),
     ]).toAst();
   }