[ML] Fix sampling probability when switching between automatic and manual

x-pack/plugins/data_visualizer/public/application/common/components/document_count_content/document_count_content.tsx

@@ -15,36 +15,58 @@ import {
   EuiPanel,
   EuiSpacer,
   EuiCallOut,
-  EuiRange,
   EuiSelect,
   EuiFormRow,
 } from '@elastic/eui';
 import { i18n } from '@kbn/i18n';
-import { debounce, sortedIndex } from 'lodash';
 import { FormattedMessage } from '@kbn/i18n-react';
 import { isDefined } from '@kbn/ml-is-defined';
 import type { DocumentCountChartPoint } from './document_count_chart';
 import {
-  RANDOM_SAMPLER_STEP,
-  RANDOM_SAMPLER_PROBABILITIES,
   RandomSamplerOption,
   RANDOM_SAMPLER_SELECT_OPTIONS,
   RANDOM_SAMPLER_OPTION,
 } from '../../../index_data_visualizer/constants/random_sampler';
 import { TotalCountHeader } from './total_count_header';
 import type { DocumentCountStats } from '../../../../../common/types/field_stats';
 import { DocumentCountChart } from './document_count_chart';
+import { RandomSamplerRangeSlider } from './random_sampler_range_slider';
 
 export interface Props {
   documentCountStats?: DocumentCountStats;
   totalCount: number;
   samplingProbability?: number | null;
-  setSamplingProbability?: (value: number) => void;
+  setSamplingProbability?: (value: number | null) => void;
   randomSamplerPreference?: RandomSamplerOption;
   setRandomSamplerPreference: (value: RandomSamplerOption) => void;
   loading: boolean;
 }
 
+const ProbabilityUsedMessage = ({ samplingProbability }: Pick<Props, 'samplingProbability'>) => {
+  return isDefined(samplingProbability) ? (
+    <div data-test-subj="dvRandomSamplerProbabilityUsedMsg">
+      <EuiSpacer size="m" />
+
+      <FormattedMessage
+        id="xpack.dataVisualizer.randomSamplerSettingsPopUp.probabilityLabel"
+        defaultMessage="Probability used: {samplingProbability}%"
+        values={{ samplingProbability: samplingProbability * 100 }}
+      />
+    </div>
+  ) : null;
+};
+
+const CalculatingProbabilityMessage = (
+  <div data-test-subj="dvRandomSamplerCalculatingProbabilityMsg">
+    <EuiSpacer size="m" />
+
+    <FormattedMessage
+      id="xpack.dataVisualizer.randomSamplerSettingsPopUp.calculatingProbabilityLabel"
+      defaultMessage="Calculating the optimal probability"
+    />
+  </div>
+);
+
 export const DocumentCountContent: FC<Props> = ({
   documentCountStats,
   totalCount,
@@ -64,24 +86,6 @@ export const DocumentCountContent: FC<Props> = ({
     setShowSamplingOptionsPopover(false);
   }, [setShowSamplingOptionsPopover]);
 
-  // eslint-disable-next-line react-hooks/exhaustive-deps
-  const updateSamplingProbability = useCallback(
-    debounce((newProbability: number) => {
-      if (setSamplingProbability) {
-        const idx = sortedIndex(RANDOM_SAMPLER_PROBABILITIES, newProbability);
-        const closestPrev = RANDOM_SAMPLER_PROBABILITIES[idx - 1];
-        const closestNext = RANDOM_SAMPLER_PROBABILITIES[idx];
-        const closestProbability =
-          Math.abs(closestPrev - newProbability) < Math.abs(closestNext - newProbability)
-            ? closestPrev
-            : closestNext;
-
-        setSamplingProbability(closestProbability / 100);
-      }
-    }, 100),
-    [setSamplingProbability]
-  );
-
   const calloutInfoMessage = useMemo(() => {
     switch (randomSamplerPreference) {
       case RANDOM_SAMPLER_OPTION.OFF:
@@ -126,19 +130,6 @@ export const DocumentCountContent: FC<Props> = ({
 
   const approximate = documentCountStats.randomlySampled === true;
 
-  const ProbabilityUsed =
-    randomSamplerPreference !== RANDOM_SAMPLER_OPTION.OFF && isDefined(samplingProbability) ? (
-      <div data-test-subj="dvRandomSamplerAutomaticProbabilityMsg">
-        <EuiSpacer size="m" />
-
-        <FormattedMessage
-          id="xpack.dataVisualizer.randomSamplerSettingsPopUp.probabilityLabel"
-          defaultMessage="Probability used: {samplingProbability}%"
-          values={{ samplingProbability: samplingProbability * 100 }}
-        />
-      </div>
-    ) : null;
-
   return (
     <>
       <EuiFlexGroup alignItems="center" gutterSize="xs">
@@ -195,37 +186,19 @@ export const DocumentCountContent: FC<Props> = ({
               </EuiFormRow>
 
               {randomSamplerPreference === RANDOM_SAMPLER_OPTION.ON_MANUAL ? (
-                <EuiFlexItem grow={true}>
-                  <EuiSpacer size="m" />
-                  <EuiFormRow
-                    label={i18n.translate(
-                      'xpack.dataVisualizer.randomSamplerSettingsPopUp.randomSamplerPercentageRowLabel',
-                      {
-                        defaultMessage: 'Sampling percentage',
-                      }
-                    )}
-                  >
-                    <EuiRange
-                      fullWidth
-                      showValue
-                      showTicks
-                      showRange={false}
-                      min={RANDOM_SAMPLER_STEP}
-                      max={0.5 * 100}
-                      value={(samplingProbability ?? 1) * 100}
-                      ticks={RANDOM_SAMPLER_PROBABILITIES.map((d) => ({
-                        value: d,
-                        label: d === 0.001 || d >= 5 ? `${d}%` : '',
-                      }))}
-                      onChange={(e) => updateSamplingProbability(Number(e.currentTarget.value))}
-                      step={RANDOM_SAMPLER_STEP}
-                      data-test-subj="dvRandomSamplerProbabilityRange"
-                    />
-                  </EuiFormRow>
-                </EuiFlexItem>
-              ) : (
-                ProbabilityUsed
-              )}
+                <RandomSamplerRangeSlider
+                  samplingProbability={samplingProbability}
+                  setSamplingProbability={setSamplingProbability}
+                />
+              ) : null}
+
+              {randomSamplerPreference === RANDOM_SAMPLER_OPTION.ON_AUTOMATIC ? (
+                loading ? (
+                  CalculatingProbabilityMessage
+                ) : (
+                  <ProbabilityUsedMessage samplingProbability={samplingProbability} />
+                )
+              ) : null}
             </EuiPanel>
           </EuiPopover>
           <EuiFlexItem />

x-pack/plugins/data_visualizer/public/application/common/components/document_count_content/random_sampler_range_slider.tsx

@@ -0,0 +1,112 @@
+/*
+ * 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; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { EuiButton, EuiFlexItem, EuiFormRow, EuiRange, EuiSpacer } from '@elastic/eui';
+import { i18n } from '@kbn/i18n';
+import { isDefined } from '@kbn/ml-is-defined';
+import { FormattedMessage } from '@kbn/i18n-react';
+import React, { useState } from 'react';
+import { roundToDecimalPlace } from '../utils';
+import {
+  MIN_SAMPLER_PROBABILITY,
+  RANDOM_SAMPLER_PROBABILITIES,
+  RANDOM_SAMPLER_STEP,
+} from '../../../index_data_visualizer/constants/random_sampler';
+
+export const RandomSamplerRangeSlider = ({
+  samplingProbability,
+  setSamplingProbability,
+}: {
+  samplingProbability?: number | null;
+  setSamplingProbability?: (value: number | null) => void;
+}) => {
+  // Keep track of the input in sampling probability slider when mode is on - manual
+  // before 'Apply' is clicked
+  const [samplingProbabilityInput, setSamplingProbabilityInput] = useState(samplingProbability);
+
+  const isInvalidSamplingProbabilityInput =
+    !isDefined(samplingProbabilityInput) ||
+    isNaN(samplingProbabilityInput) ||
+    samplingProbabilityInput < MIN_SAMPLER_PROBABILITY ||
+    samplingProbabilityInput > 0.5;
+
+  const inputValue = (samplingProbabilityInput ?? MIN_SAMPLER_PROBABILITY) * 100;
+
+  return (
+    <EuiFlexItem grow={true}>
+      <EuiSpacer size="m" />
+      <EuiFormRow
+        fullWidth
+        label={i18n.translate(
+          'xpack.dataVisualizer.randomSamplerSettingsPopUp.randomSamplerPercentageRowLabel',
+          {
+            defaultMessage: 'Sampling percentage',
+          }
+        )}
+        helpText={i18n.translate(
+          'xpack.dataVisualizer.randomSamplerSettingsPopUp.randomSamplerPercentageRowHelpText',
+          {
+            defaultMessage: 'Choose a value between 0.001% and 50% to randomly sample data.',
+          }
+        )}
+      >
+        <EuiRange
+          fullWidth
+          showValue
+          showRange
+          showLabels
+          showInput="inputWithPopover"
+          min={RANDOM_SAMPLER_STEP}
+          max={50}
+          // Rounding to 0 decimal place because sometimes js results in weird number when multiplying fractions
+          // e.g. 0.07 * 100 yields 7.000000000000001
+          value={
+            inputValue >= 1
+              ? roundToDecimalPlace(inputValue, 0)
+              : roundToDecimalPlace(inputValue, 3)
+          }
+          ticks={RANDOM_SAMPLER_PROBABILITIES.map((d) => ({
+            value: d,
+            label: d === 0.001 || d >= 5 ? `${d}` : '',
+          }))}
+          isInvalid={isInvalidSamplingProbabilityInput}
+          onChange={(e) => {
+            const value = parseFloat((e.target as HTMLInputElement).value);
+            const prevValue = samplingProbabilityInput ? samplingProbabilityInput * 100 : value;
+
+            if (value > 0 && value <= 1) {
+              setSamplingProbabilityInput(value / 100);
+            } else {
+              // Because the incremental step is very small (0.0001),
+              // everytime user clicks the ^/∨ in the numerical input
+              // we need to make sure it rounds up or down to the next whole number
+              const nearestInt = value > prevValue ? Math.ceil(value) : Math.floor(value);
+              setSamplingProbabilityInput(nearestInt / 100);
+            }
+          }}
+          step={RANDOM_SAMPLER_STEP}
+          data-test-subj="dvRandomSamplerProbabilityRange"
+          append={
+            <EuiButton
+              disabled={isInvalidSamplingProbabilityInput}
+              onClick={() => {
+                if (setSamplingProbability && isDefined(samplingProbabilityInput)) {
+                  setSamplingProbability(samplingProbabilityInput);
+                }
+              }}
+            >
+              <FormattedMessage
+                id="xpack.dataVisualizer.randomSamplerSettingsPopUp.randomSamplerPercentageApply"
+                defaultMessage="Apply"
+              />
+            </EuiButton>
+          }
+        />
+      </EuiFormRow>
+    </EuiFlexItem>
+  );
+};

x-pack/plugins/data_visualizer/public/application/index_data_visualizer/components/index_data_visualizer_view/index_data_visualizer_view.tsx

@@ -41,6 +41,7 @@ import { useCurrentEuiTheme } from '../../../common/hooks/use_current_eui_theme'
 import {
   DV_FROZEN_TIER_PREFERENCE,
   DV_RANDOM_SAMPLER_PREFERENCE,
+  DV_RANDOM_SAMPLER_P_VALUE,
   type DVKey,
   type DVStorageMapped,
 } from '../../types/storage';
@@ -71,7 +72,11 @@ import { DataVisualizerDataViewManagement } from '../data_view_management';
 import { GetAdditionalLinks } from '../../../common/components/results_links';
 import { useDataVisualizerGridData } from '../../hooks/use_data_visualizer_grid_data';
 import { DataVisualizerGridInput } from '../../embeddables/grid_embeddable/grid_embeddable';
-import { RANDOM_SAMPLER_OPTION } from '../../constants/random_sampler';
+import {
+  MIN_SAMPLER_PROBABILITY,
+  RANDOM_SAMPLER_OPTION,
+  RandomSamplerOption,
+} from '../../constants/random_sampler';
 
 interface DataVisualizerPageState {
   overallStats: OverallStats;
@@ -143,6 +148,11 @@ export const IndexDataVisualizerView: FC<IndexDataVisualizerViewProps> = (dataVi
     DVStorageMapped<typeof DV_RANDOM_SAMPLER_PREFERENCE>
   >(DV_RANDOM_SAMPLER_PREFERENCE, RANDOM_SAMPLER_OPTION.ON_AUTOMATIC);
 
+  const [savedRandomSamplerProbability, saveRandomSamplerProbability] = useStorage<
+    DVKey,
+    DVStorageMapped<typeof DV_RANDOM_SAMPLER_P_VALUE>
+  >(DV_RANDOM_SAMPLER_P_VALUE, MIN_SAMPLER_PROBABILITY);
+
   const [frozenDataPreference, setFrozenDataPreference] = useStorage<
     DVKey,
     DVStorageMapped<typeof DV_FROZEN_TIER_PREFERENCE>
@@ -156,6 +166,7 @@ export const IndexDataVisualizerView: FC<IndexDataVisualizerViewProps> = (dataVi
     () =>
       getDefaultDataVisualizerListState({
         rndSamplerPref: savedRandomSamplerPreference,
+        probability: savedRandomSamplerProbability,
       }),
     // We just need to load the saved preference when the page is first loaded
     // eslint-disable-next-line react-hooks/exhaustive-deps
@@ -316,9 +327,38 @@ export const IndexDataVisualizerView: FC<IndexDataVisualizerViewProps> = (dataVi
     ]
   );
 
-  const setSamplingProbability = (value: number | null) => {
-    setDataVisualizerListState({ ...dataVisualizerListState, probability: value });
-  };
+  const setSamplingProbability = useCallback(
+    (value: number | null) => {
+      if (savedRandomSamplerPreference === RANDOM_SAMPLER_OPTION.ON_MANUAL && value !== null) {
+        saveRandomSamplerProbability(value);
+      }
+      setDataVisualizerListState({ ...dataVisualizerListState, probability: value });
+    },
+    [
+      dataVisualizerListState,
+      saveRandomSamplerProbability,
+      savedRandomSamplerPreference,
+      setDataVisualizerListState,
+    ]
+  );
+
+  const setRandomSamplerPreference = useCallback(
+    (nextPref: RandomSamplerOption) => {
+      if (nextPref === RANDOM_SAMPLER_OPTION.ON_MANUAL) {
+        // By default, when switching to manual, restore previously chosen probability
+        // else, default to 0.001%
+        setSamplingProbability(
+          savedRandomSamplerProbability &&
+            savedRandomSamplerProbability > 0 &&
+            savedRandomSamplerProbability <= 0.5
+            ? savedRandomSamplerProbability
+            : MIN_SAMPLER_PROBABILITY
+        );
+      }
+      saveRandomSamplerPreference(nextPref);
+    },
+    [savedRandomSamplerProbability, setSamplingProbability, saveRandomSamplerPreference]
+  );
 
   useEffect(
     function clearFiltersOnLeave() {
@@ -565,7 +605,7 @@ export const IndexDataVisualizerView: FC<IndexDataVisualizerViewProps> = (dataVi
                       }
                       loading={overallStatsProgress.loaded < 100}
                       randomSamplerPreference={savedRandomSamplerPreference}
-                      setRandomSamplerPreference={saveRandomSamplerPreference}
+                      setRandomSamplerPreference={setRandomSamplerPreference}
                     />
                   </EuiFlexGroup>
                 </>

x-pack/plugins/data_visualizer/public/application/index_data_visualizer/constants/random_sampler.ts

@@ -8,13 +8,11 @@
 import { i18n } from '@kbn/i18n';
 
 export const RANDOM_SAMPLER_PROBABILITIES = [
-  0.5, 0.25, 0.1, 0.05, 0.025, 0.01, 0.005, 0.0025, 0.001, 0.0005, 0.00025, 0.0001, 0.00005,
-  0.00001,
-]
-  .reverse()
-  .map((n) => n * 100);
+  0.00001, 0.00005, 0.0001, 0.0005, 0.001, 0.005, 0.01, 0.05, 0.1, 0.2, 0.3, 0.4, 0.5,
+].map((n) => n * 100);
 
-export const RANDOM_SAMPLER_STEP = 0.00001 * 100;
+export const MIN_SAMPLER_PROBABILITY = 0.00001;
+export const RANDOM_SAMPLER_STEP = MIN_SAMPLER_PROBABILITY * 100;
 
 export const RANDOM_SAMPLER_OPTION = {
   ON_AUTOMATIC: 'on_automatic',

x-pack/plugins/data_visualizer/public/application/index_data_visualizer/types/storage.ts

@@ -16,7 +16,7 @@ export const DV_RANDOM_SAMPLER_P_VALUE = 'dataVisualizer.randomSamplerPValue';
 export type DV = Partial<{
   [DV_FROZEN_TIER_PREFERENCE]: FrozenTierPreference;
   [DV_RANDOM_SAMPLER_PREFERENCE]: RandomSamplerOption;
-  [DV_RANDOM_SAMPLER_P_VALUE]: number;
+  [DV_RANDOM_SAMPLER_P_VALUE]: null | number;
 }> | null;
 
 export type DVKey = keyof Exclude<DV, null>;