[Lens] Random sampling

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

@@ -69,6 +69,7 @@ export const getAggTypes = () => ({
     { name: BUCKET_TYPES.GEOHASH_GRID, fn: buckets.getGeoHashBucketAgg },
     { name: BUCKET_TYPES.GEOTILE_GRID, fn: buckets.getGeoTitleBucketAgg },
     { name: BUCKET_TYPES.SAMPLER, fn: buckets.getSamplerBucketAgg },
+    { name: BUCKET_TYPES.RANDOM_SAMPLER, fn: buckets.getRandomSamplerBucketAgg },
     { name: BUCKET_TYPES.DIVERSIFIED_SAMPLER, fn: buckets.getDiversifiedSamplerBucketAgg },
   ],
 });
@@ -90,6 +91,7 @@ export const getAggTypesFunctions = () => [
   buckets.aggMultiTerms,
   buckets.aggRareTerms,
   buckets.aggSampler,
+  buckets.aggRandomSampler,
   buckets.aggDiversifiedSampler,
   metrics.aggAvg,
   metrics.aggBucketAvg,

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

@@ -69,6 +69,7 @@ describe('Aggs service', () => {
           "geotile_grid",
           "sampler",
           "diversified_sampler",
+          "random_sampler",
           "foo",
         ]
       `);
@@ -123,6 +124,7 @@ describe('Aggs service', () => {
           "geotile_grid",
           "sampler",
           "diversified_sampler",
+          "random_sampler",
         ]
       `);
       expect(bStart.types.getAll().metrics.map((t) => t.name)).toMatchInlineSnapshot(`

src/plugins/data/common/search/aggs/buckets/_terms_other_bucket_helper.ts

@@ -67,6 +67,10 @@ const getAggResultBuckets = (
         const aggKey = keys(responseAgg)[aggId];
         const aggConfig = find(aggConfigs.aggs, (agg) => agg.id === aggKey);
         if (aggConfig) {
+          if (keyParts[i] === '_no_bucket_') {
+            responseAgg = aggById;
+            break;
+          }
           const aggResultBucket = find(aggById.buckets, (bucket, bucketObjKey) => {
             const bucketKey = aggConfig
               .getKey(bucket, isNumber(bucketObjKey) ? undefined : bucketObjKey)
@@ -172,16 +176,27 @@ export const buildOtherBucketAgg = (
     filters: any[],
     key: string
   ) => {
+    const agg = aggregations[aggId];
+    const newAggIndex = aggIndex + 1;
+    const newAgg = bucketAggs[newAggIndex];
+    const currentAgg = bucketAggs[aggIndex];
+    if (agg && !Array.isArray(agg?.buckets)) {
+      // there is no buckets array, assum an agg which is not reflected in the response and recurse
+      walkBucketTree(
+        newAggIndex,
+        agg,
+        newAgg.id,
+        filters,
+        `${key}${OTHER_BUCKET_SEPARATOR}_no_bucket_`
+      );
+      return;
+    }
     // make sure there are actually results for the buckets
-    if (aggregations[aggId]?.buckets.length < 1) {
+    if (agg?.buckets.length < 1) {
       noAggBucketResults = true;
       return;
     }
 
-    const agg = aggregations[aggId];
-    const newAggIndex = aggIndex + 1;
-    const newAgg = bucketAggs[newAggIndex];
-    const currentAgg = bucketAggs[aggIndex];
     if (aggIndex === index && agg && agg.sum_other_doc_count > 0) {
       exhaustiveBuckets = false;
     }

src/plugins/data/common/search/aggs/buckets/bucket_agg_types.ts

@@ -22,5 +22,6 @@ export enum BUCKET_TYPES {
   GEOTILE_GRID = 'geotile_grid',
   DATE_HISTOGRAM = 'date_histogram',
   SAMPLER = 'sampler',
+  RANDOM_SAMPLER = 'random_sampler',
   DIVERSIFIED_SAMPLER = 'diversified_sampler',
 }

src/plugins/data/common/search/aggs/buckets/index.ts

@@ -48,4 +48,6 @@ export * from './sampler_fn';
 export * from './sampler';
 export * from './diversified_sampler_fn';
 export * from './diversified_sampler';
+export * from './random_sampler_fn';
+export * from './random_sampler';
 export { SHARD_DELAY_AGG_NAME } from './shard_delay';

src/plugins/data/common/search/aggs/buckets/random_sampler.ts

@@ -0,0 +1,43 @@
+/*
+ * 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.
+ */
+
+import { i18n } from '@kbn/i18n';
+import { BucketAggType } from './bucket_agg_type';
+import { BaseAggParams } from '../types';
+import { aggRandomSamplerFnName } from './random_sampler_fn';
+
+export const RANDOM_SAMPLER_AGG_NAME = 'random_sampler';
+
+const title = i18n.translate('data.search.aggs.buckets.randomSamplerTitle', {
+  defaultMessage: 'Random Sampler',
+  description: 'Random sampler aggregation title',
+});
+
+export interface AggParamsRandomSampler extends BaseAggParams {
+  /**
+   * The sampling probability
+   */
+  probability: number;
+}
+
+/**
+ * A filtering aggregation used to limit any sub aggregations' processing to a sample of the top-scoring documents.
+ */
+export const getRandomSamplerBucketAgg = () =>
+  new BucketAggType({
+    name: RANDOM_SAMPLER_AGG_NAME,
+    title,
+    customLabels: false,
+    expressionName: aggRandomSamplerFnName,
+    params: [
+      {
+        name: 'probability',
+        type: 'number',
+      },
+    ],
+  });

src/plugins/data/common/search/aggs/buckets/random_sampler_fn.test.ts

@@ -0,0 +1,34 @@
+/*
+ * 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.
+ */
+
+import { functionWrapper } from '../test_helpers';
+import { aggRandomSampler } from './random_sampler_fn';
+
+describe('aggRandomSampler', () => {
+  const fn = functionWrapper(aggRandomSampler());
+
+  test('includes params when they are provided', () => {
+    const actual = fn({
+      id: 'random_sampler',
+      schema: 'bucket',
+      probability: 0.1,
+    });
+
+    expect(actual.value).toMatchInlineSnapshot(`
+      Object {
+        "enabled": true,
+        "id": "random_sampler",
+        "params": Object {
+          "probability": 0.1,
+        },
+        "schema": "bucket",
+        "type": "random_sampler",
+      }
+    `);
+  });
+});

src/plugins/data/common/search/aggs/buckets/random_sampler_fn.ts

@@ -0,0 +1,76 @@
+/*
+ * 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.
+ */
+
+import { i18n } from '@kbn/i18n';
+import { ExpressionFunctionDefinition } from '@kbn/expressions-plugin/common';
+import { AggExpressionFunctionArgs, AggExpressionType, BUCKET_TYPES } from '..';
+import { RANDOM_SAMPLER_AGG_NAME } from './random_sampler';
+
+export const aggRandomSamplerFnName = 'aggRandomSampler';
+
+type Input = any;
+type Arguments = AggExpressionFunctionArgs<typeof BUCKET_TYPES.RANDOM_SAMPLER>;
+
+type Output = AggExpressionType;
+type FunctionDefinition = ExpressionFunctionDefinition<
+  typeof aggRandomSamplerFnName,
+  Input,
+  Arguments,
+  Output
+>;
+
+export const aggRandomSampler = (): FunctionDefinition => ({
+  name: aggRandomSamplerFnName,
+  help: i18n.translate('data.search.aggs.function.buckets.randomSampler.help', {
+    defaultMessage: 'Generates a serialized agg config for a Random sampler agg',
+  }),
+  type: 'agg_type',
+  args: {
+    id: {
+      types: ['string'],
+      help: i18n.translate('data.search.aggs.buckets.randomSampler.id.help', {
+        defaultMessage: 'ID for this aggregation',
+      }),
+    },
+    enabled: {
+      types: ['boolean'],
+      default: true,
+      help: i18n.translate('data.search.aggs.buckets.randomSampler.enabled.help', {
+        defaultMessage: 'Specifies whether this aggregation should be enabled',
+      }),
+    },
+    schema: {
+      types: ['string'],
+      help: i18n.translate('data.search.aggs.buckets.randomSampler.schema.help', {
+        defaultMessage: 'Schema to use for this aggregation',
+      }),
+    },
+    probability: {
+      types: ['number'],
+      help: i18n.translate('data.search.aggs.buckets.randomSampler.probability.help', {
+        defaultMessage: 'The sampling probability',
+      }),
+    },
+  },
+  fn: (input, args) => {
+    const { id, enabled, schema, ...rest } = args;
+
+    return {
+      type: 'agg_type',
+      value: {
+        id,
+        enabled,
+        schema,
+        type: RANDOM_SAMPLER_AGG_NAME,
+        params: {
+          ...rest,
+        },
+      },
+    };
+  },
+});

src/plugins/data/common/search/aggs/metrics/lib/sibling_pipeline_agg_helper.ts

@@ -31,7 +31,13 @@ const metricAggFilter: string[] = [
   '!filtered_metric',
   '!single_percentile',
 ];
-const bucketAggFilter: string[] = ['!filter', '!sampler', '!diversified_sampler', '!multi_terms'];
+const bucketAggFilter: string[] = [
+  '!filter',
+  '!sampler',
+  '!diversified_sampler',
+  '!random_sampler',
+  '!multi_terms',
+];
 
 export const siblingPipelineType = i18n.translate(
   'data.search.aggs.metrics.siblingPipelineAggregationsSubtypeTitle',

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

@@ -109,8 +109,10 @@ import {
   AggParamsMovingAvgSerialized,
   AggParamsSerialDiffSerialized,
   AggParamsTopHitSerialized,
+  aggRandomSampler,
 } from '.';
 import { AggParamsSampler } from './buckets/sampler';
+import { AggParamsRandomSampler } from './buckets/random_sampler';
 import { AggParamsDiversifiedSampler } from './buckets/diversified_sampler';
 import { AggParamsSignificantText } from './buckets/significant_text';
 import { aggTopMetrics } from './metrics/top_metrics_fn';
@@ -185,6 +187,7 @@ interface SerializedAggParamsMapping {
   [BUCKET_TYPES.MULTI_TERMS]: AggParamsMultiTermsSerialized;
   [BUCKET_TYPES.RARE_TERMS]: AggParamsRareTerms;
   [BUCKET_TYPES.SAMPLER]: AggParamsSampler;
+  [BUCKET_TYPES.RANDOM_SAMPLER]: AggParamsRandomSampler;
   [BUCKET_TYPES.DIVERSIFIED_SAMPLER]: AggParamsDiversifiedSampler;
   [METRIC_TYPES.AVG]: AggParamsAvg;
   [METRIC_TYPES.CARDINALITY]: AggParamsCardinality;
@@ -230,6 +233,7 @@ export interface AggParamsMapping {
   [BUCKET_TYPES.MULTI_TERMS]: AggParamsMultiTerms;
   [BUCKET_TYPES.RARE_TERMS]: AggParamsRareTerms;
   [BUCKET_TYPES.SAMPLER]: AggParamsSampler;
+  [BUCKET_TYPES.RANDOM_SAMPLER]: AggParamsRandomSampler;
   [BUCKET_TYPES.DIVERSIFIED_SAMPLER]: AggParamsDiversifiedSampler;
   [METRIC_TYPES.AVG]: AggParamsAvg;
   [METRIC_TYPES.CARDINALITY]: AggParamsCardinality;
@@ -265,6 +269,7 @@ export interface AggFunctionsMapping {
   aggFilter: ReturnType<typeof aggFilter>;
   aggFilters: ReturnType<typeof aggFilters>;
   aggSignificantTerms: ReturnType<typeof aggSignificantTerms>;
+  aggRandomSampler: ReturnType<typeof aggRandomSampler>;
   aggIpRange: ReturnType<typeof aggIpRange>;
   aggDateRange: ReturnType<typeof aggDateRange>;
   aggRange: ReturnType<typeof aggRange>;

x-pack/plugins/lens/public/editor_frame_service/editor_frame/config_panel/layer_panel.tsx

@@ -319,6 +319,11 @@ export function LayerPanel(
           visualizationState,
           updateVisualization
         ) || []),
+        ...(layerDatasource?.getSupportedActionsForLayer?.(
+          layerId,
+          layerDatasourceState,
+          (newState) => updateDatasource(datasourceId, newState)
+        ) || []),
         ...getSharedActions({
           activeVisualization,
           core,
@@ -333,12 +338,16 @@ export function LayerPanel(
     [
       activeVisualization,
       core,
+      datasourceId,
       isOnlyLayer,
       isTextBasedLanguage,
+      layerDatasource,
+      layerDatasourceState,
       layerId,
       layerIndex,
       onCloneLayer,
       onRemoveLayer,
+      updateDatasource,
       updateVisualization,
       visualizationState,
     ]