Fix PPL head command to show simple result count (opensearch-project#11405)

* Fix PPL histogram to exclude head clause from total hit count

Signed-off-by: Kai Huang <ahkcs@amazon.com>

* Simplify head command display: show only row count without total

When PPL query ends with head command, display "X results" instead of
"X of Y results". Remove the separate total count query infrastructure
since it is no longer needed.

- Replace stripHeadFromQuery with queryEndsWithHead for detection
- Remove buildPPLTotalCountQuery, executeTotalCountQuery, and related logic
- Delete use_total_count_results.ts hook
- Update action_bar to skip total hits when head is detected

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Kai Huang <ahkcs@amazon.com>

---------

Signed-off-by: Kai Huang <ahkcs@amazon.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: ahkcs

src/plugins/explore/public/application/utils/state_management/actions/query_actions.ts

@@ -260,12 +260,12 @@ export const executeQueries = createAsyncThunk<
   const dataTableCacheKey = defaultCacheKey;
   const breakdownField = state.queryEditor.breakdownField;
   const histogramCacheKey = prepareHistogramCacheKey(query, !!breakdownField);
+  const queryString = defaultPrepareQueryString(query);
 
   // Check what needs execution for core queries
   // If results exist but query status is UNINITIALIZED (after cancel), we need to re-execute
   const dataTableQueryStatus = state.queryEditor.queryStatusMap[dataTableCacheKey];
   const histogramQueryStatus = state.queryEditor.queryStatusMap[histogramCacheKey];
-
   // Early exit if query should be skipped
   if (shouldSkipQueryExecution(query)) {
     return;
@@ -278,7 +278,6 @@ export const executeQueries = createAsyncThunk<
     query.language !== 'PROMQL' &&
     (!results[histogramCacheKey] ||
       histogramQueryStatus?.status === QueryExecutionStatus.UNINITIALIZED);
-
   const promises = [];
   // Execute query without aggregations
   if (needsDataTableQuery) {
@@ -287,7 +286,7 @@ export const executeQueries = createAsyncThunk<
         executeDataTableQuery({
           services,
           cacheKey: dataTableCacheKey,
-          queryString: defaultPrepareQueryString(query),
+          queryString,
         })
       )
     );
@@ -300,7 +299,7 @@ export const executeQueries = createAsyncThunk<
       executeHistogramQuery({
         services,
         cacheKey: histogramCacheKey,
-        queryString: defaultPrepareQueryString(query),
+        queryString,
         interval,
       })
     );
@@ -538,13 +537,15 @@ const executeQueryBase = async (
       histogramConfig = createHistogramConfigWithInterval(dataView, interval, services, getState);
     }
 
+    let effectiveQuery = queryString;
+    if (query.language === 'PPL' && histogramConfig && isHistogramQuery) {
+      effectiveQuery = buildPPLHistogramQuery(queryString, histogramConfig);
+    }
+
     const preparedQueryObject = {
       ...query,
       dataset,
-      query:
-        query.language === 'PPL' && isHistogramQuery && histogramConfig
-          ? buildPPLHistogramQuery(queryString, histogramConfig)
-          : queryString,
+      query: effectiveQuery,
     };
 
     let searchSource;

src/plugins/explore/public/application/utils/state_management/actions/utils.test.ts

@@ -113,6 +113,46 @@ describe('Utils - Histogram Breakdown Support', () => {
     });
   });
 
+  describe('queryEndsWithHead', () => {
+    it('should detect head at end of query', () => {
+      expect(utils.queryEndsWithHead('source=logs | head 200')).toBe(true);
+    });
+
+    it('should detect head with from clause', () => {
+      expect(utils.queryEndsWithHead('source=logs | head 200 from 10')).toBe(true);
+    });
+
+    it('should detect head without explicit count', () => {
+      expect(utils.queryEndsWithHead('source=logs | head')).toBe(true);
+    });
+
+    it('should detect head followed by where clause', () => {
+      expect(utils.queryEndsWithHead('source=logs | head 200 | where status = 200')).toBe(true);
+    });
+
+    it('should return false when no head is present', () => {
+      expect(utils.queryEndsWithHead('source=logs | WHERE status = 200')).toBe(false);
+    });
+
+    it('should return false when head is only inside subquery', () => {
+      expect(utils.queryEndsWithHead('source=logs | where id in [source=other | head 10]')).toBe(
+        false
+      );
+    });
+
+    it('should detect main query head even with subquery head', () => {
+      expect(
+        utils.queryEndsWithHead('source=logs | where id in [source=other | head 10] | head 200')
+      ).toBe(true);
+    });
+
+    it('should return false when head is in the middle of the query', () => {
+      expect(utils.queryEndsWithHead('source=logs | head 200 | stats count() by status')).toBe(
+        false
+      );
+    });
+  });
+
   describe('buildPPLHistogramQuery', () => {
     it('should return original query when aggs is missing', () => {
       const query = 'source=logs';
@@ -144,6 +184,16 @@ describe('Utils - Histogram Breakdown Support', () => {
       const result = utils.buildPPLHistogramQuery(query, histogramConfig);
       expect(result).toBe('source=logs | stats count() by span(@timestamp, 1h)');
     });
+
+    it('should preserve head clause in histogram query', () => {
+      const query = 'source=logs | head 200';
+      const histogramConfig = createBaseHistogramConfig({
+        aggs: { 2: { date_histogram: {} } },
+      });
+
+      const result = utils.buildPPLHistogramQuery(query, histogramConfig);
+      expect(result).toBe('source=logs | head 200 | stats count() by span(@timestamp, 1h)');
+    });
   });
 
   describe('processRawResultsForHistogram', () => {

src/plugins/explore/public/application/utils/state_management/actions/utils.ts

@@ -103,6 +103,15 @@ export function fillMissingTimestamps(
   return filledSeriesMap;
 }
 
+/**
+ * Checks if the main query ends with a head command (optionally followed by `from N` or `| where`).
+ * Subquery brackets [...] are masked so that head inside subqueries is ignored.
+ */
+export const queryEndsWithHead = (queryString: string): boolean => {
+  const masked = queryString.replace(/\[.*?\]/g, (match) => '\0'.repeat(match.length));
+  return /\|\s*head\b(\s+\d+)?(\s+from\s+\d+)?\s*(\|\s*where\b.*)?\s*$/i.test(masked);
+};
+
 export const buildPPLHistogramQuery = (query: string, histogramConfig: HistogramConfig): string => {
   const { aggs, finalInterval, timeFieldName, breakdownField } = histogramConfig;
 
@@ -223,7 +232,7 @@ export const processRawResultsForHistogram = (
 
     Object.entries(responseAggs).forEach(([id, value]) => {
       if (aggsConfig && aggsConfig.date_histogram) {
-        let totalHits = rawResults.hits.total;
+        let totalHits = 0;
         const buckets = value as Array<{ key: string; value: number }>;
         tempResult.aggregations[id] = {
           buckets: buckets.map((bucket) => {

src/plugins/explore/public/components/tabs/action_bar/action_bar.test.tsx

@@ -87,6 +87,19 @@ jest.mock('../../../application/utils/hooks/use_histogram_results', () => ({
   }),
 }));
 
+jest.mock('react-redux', () => ({
+  ...jest.requireActual('react-redux'),
+  useSelector: () => ({ language: 'PPL', query: 'source=logs', dataset: { id: 'test' } }),
+}));
+
+jest.mock('../../../application/utils/state_management/actions/query_actions', () => ({
+  defaultPrepareQueryString: () => 'source=logs',
+}));
+
+jest.mock('../../../application/utils/state_management/actions/utils', () => ({
+  queryEndsWithHead: () => false,
+}));
+
 describe('ActionBar', () => {
   test('should render the action bar component', () => {
     render(<ActionBar />);

src/plugins/explore/public/components/tabs/action_bar/action_bar.tsx

@@ -5,6 +5,7 @@
 import React, { memo, useMemo } from 'react';
 import { useObservable } from 'react-use';
 import { of } from 'rxjs';
+import { useSelector as useReduxSelector } from 'react-redux';
 import { DiscoverResultsActionBar } from './results_action_bar/results_action_bar';
 import { ExploreServices } from '../../../types';
 import { useOpenSearchDashboards } from '../../../../../opensearch_dashboards_react/public';
@@ -14,6 +15,9 @@ import { useDatasetContext } from '../../../application/context';
 import { ExploreFlavor } from '../../../../common';
 import { useTabResults } from '../../../application/utils/hooks/use_tab_results';
 import { useHistogramResults } from '../../../application/utils/hooks/use_histogram_results';
+import { RootState } from '../../../application/utils/state_management/store';
+import { defaultPrepareQueryString } from '../../../application/utils/state_management/actions/query_actions';
+import { queryEndsWithHead } from '../../../application/utils/state_management/actions/utils';
 
 interface ActionBarProps {
   filteredRowsCount?: number;
@@ -28,6 +32,7 @@ const ActionBarComponent = ({ filteredRowsCount }: ActionBarProps = {}) => {
   const { dataset } = useDatasetContext();
   const { results } = useTabResults();
   const { results: histogramResults } = useHistogramResults();
+  const query = useReduxSelector((state: RootState) => state.query);
   const { core, inspector, inspectorAdapters, slotRegistry } = services;
   const savedSearch = useSelector(selectSavedSearch);
 
@@ -45,7 +50,10 @@ const ActionBarComponent = ({ filteredRowsCount }: ActionBarProps = {}) => {
   };
 
   const rows = results?.hits?.hits || [];
-  const totalHits = histogramResults?.hits.total;
+  // When query has head command, just show row count (no "X of Y" total)
+  const queryString = query.language === 'PPL' ? defaultPrepareQueryString(query) : '';
+  const hasHead = query.language === 'PPL' && queryEndsWithHead(queryString);
+  const totalHits = hasHead ? undefined : histogramResults?.hits.total;
   const elapsedMs = results?.elapsedMs;
 
   return (