[ES|QL] Add st_centroid agg function (elastic#175506)

## Summary Adds new agg function as in elastic/elasticsearch#104218 ### Checklist - [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
fkanout · Mar 4, 2024 · 947cb89 · 947cb89
1 parent 4d5ad06
commit 947cb89
Show file tree

Hide file tree

Showing 3 changed files with 77 additions and 26 deletions.
diff --git a/packages/kbn-monaco/src/esql/lib/ast/definitions/aggs.ts b/packages/kbn-monaco/src/esql/lib/ast/definitions/aggs.ts
@@ -130,4 +130,29 @@ export const statsAggregationFunctionDefinitions: FunctionDefinition[] = [
         },
       ],
     },
+    {
+      name: 'st_centroid',
+      description: i18n.translate('monaco.esql.definitions.stCentroidDoc', {
+        defaultMessage: 'Returns the count of distinct values in a field.',
+      }),
+      supportedCommands: ['stats'],
+      signatures: [
+        {
+          params: [{ name: 'column', type: 'cartesian_point', noNestingFunctions: true }],
+          returnType: 'number',
+          examples: [
+            `from index | stats result = st_centroid(cartesian_field)`,
+            `from index | stats st_centroid(cartesian_field)`,
+          ],
+        },
+        {
+          params: [{ name: 'column', type: 'geo_point', noNestingFunctions: true }],
+          returnType: 'number',
+          examples: [
+            `from index | stats result = st_centroid(geo_field)`,
+            `from index | stats st_centroid(geo_field)`,
+          ],
+        },
+      ],
+    },
   ]);
diff --git a/packages/kbn-monaco/src/esql/lib/ast/validation/validation.test.ts b/packages/kbn-monaco/src/esql/lib/ast/validation/validation.test.ts
@@ -21,6 +21,9 @@ import { chronoLiterals, timeLiterals } from '../definitions/literals';
 import { statsAggregationFunctionDefinitions } from '../definitions/aggs';
 import capitalize from 'lodash/capitalize';
 import { EditorError } from '../../../../types';
+import { camelCase } from 'lodash';
+
+const fieldTypes = ['number', 'date', 'boolean', 'ip', 'string', 'cartesian_point', 'geo_point'];
 
 function getCallbackMocks() {
   return {
@@ -33,11 +36,7 @@ function getCallbackMocks() {
         : /unsupported_index/.test(query)
         ? [{ name: 'unsupported_field', type: 'unsupported' }]
         : [
-            ...['string', 'number', 'date', 'boolean', 'ip'].map((type) => ({
-              name: `${type}Field`,
-              type,
-            })),
-            { name: 'geoPointField', type: 'geo_point' },
+            ...fieldTypes.map((type) => ({ name: `${camelCase(type)}Field`, type })),
             { name: 'any#Char$ field', type: 'number' },
             { name: 'kubernetes.something.something', type: 'number' },
             {
@@ -75,6 +74,10 @@ const toStringSignature = evalFunctionsDefinitions.find(({ name }) => name === '
 const toDateSignature = evalFunctionsDefinitions.find(({ name }) => name === 'to_datetime')!;
 const toBooleanSignature = evalFunctionsDefinitions.find(({ name }) => name === 'to_boolean')!;
 const toIpSignature = evalFunctionsDefinitions.find(({ name }) => name === 'to_ip')!;
+const toGeoPointSignature = evalFunctionsDefinitions.find(({ name }) => name === 'to_geopoint')!;
+const toCartesianPointSignature = evalFunctionsDefinitions.find(
+  ({ name }) => name === 'to_cartesianpoint'
+)!;
 
 const toAvgSignature = statsAggregationFunctionDefinitions.find(({ name }) => name === 'avg')!;
 
@@ -84,6 +87,8 @@ const nestedFunctions = {
   date: prepareNestedFunction(toDateSignature),
   boolean: prepareNestedFunction(toBooleanSignature),
   ip: prepareNestedFunction(toIpSignature),
+  geo_point: prepareNestedFunction(toGeoPointSignature),
+  cartesian_point: prepareNestedFunction(toCartesianPointSignature),
 };
 
 const literals = {
@@ -97,13 +102,15 @@ function getLiteralType(typeString: 'chrono_literal' | 'time_literal') {
   return `1 ${literals[typeString]}`;
 }
 function getFieldName(
-  typeString: 'string' | 'number' | 'date' | 'boolean' | 'ip',
+  typeString: string,
   { useNestedFunction, isStats }: { useNestedFunction: boolean; isStats: boolean }
 ) {
   if (useNestedFunction && isStats) {
     return prepareNestedFunction(toAvgSignature);
   }
-  return useNestedFunction ? nestedFunctions[typeString] : `${typeString}Field`;
+  return useNestedFunction && typeString in nestedFunctions
+    ? nestedFunctions[typeString as keyof typeof nestedFunctions]
+    : `${camelCase(typeString)}Field`;
 }
 
 function getMultiValue(type: 'string[]' | 'number[]' | 'boolean[]' | 'any[]') {
@@ -139,9 +146,9 @@ function getFieldMapping(
 ) {
   return params.map(({ name: _name, type, ...rest }) => {
     const typeString: string = type;
-    if (['string', 'number', 'date', 'boolean', 'ip'].includes(typeString)) {
+    if (fieldTypes.includes(typeString)) {
       return {
-        name: getFieldName(typeString as 'string' | 'number' | 'date' | 'boolean' | 'ip', {
+        name: getFieldName(typeString, {
           useNestedFunction,
           isStats: !useLiterals,
         }),
@@ -394,7 +401,14 @@ describe('validation logic', () => {
 
           const wrongFieldMapping = params.map(({ name: _name, type, ...rest }) => {
             const typeString = type;
-            const canBeFieldButNotString = ['number', 'date', 'boolean', 'ip'].includes(typeString);
+            const canBeFieldButNotString = [
+              'number',
+              'date',
+              'boolean',
+              'ip',
+              'cartesian_point',
+              'geo_point',
+            ].includes(typeString);
             const isLiteralType = /literal$/.test(typeString);
             // pick a field name purposely wrong
             const nameValue = canBeFieldButNotString || isLiteralType ? '"a"' : '5';
@@ -956,7 +970,9 @@ describe('validation logic', () => {
 
           const wrongFieldMapping = params.map(({ name: _name, type, ...rest }) => {
             const typeString = type;
-            const canBeFieldButNotString = ['number', 'date', 'boolean', 'ip'].includes(typeString);
+            const canBeFieldButNotString = fieldTypes
+              .filter((t) => t !== 'string')
+              .includes(typeString);
             const isLiteralType = /literal$/.test(typeString);
             // pick a field name purposely wrong
             const nameValue =
@@ -1198,9 +1214,12 @@ describe('validation logic', () => {
     testErrorsAndWarnings('from a | stats var0 = count(*)', []);
     testErrorsAndWarnings('from a | stats var0 = count()', []);
     testErrorsAndWarnings('from a | stats var0 = avg(numberField), count(*)', []);
+    testErrorsAndWarnings('from a | stats var0 = avg(fn(number)), count(*)', [
+      'Unknown function [fn]',
+    ]);
 
     for (const { name, alias, signatures, ...defRest } of statsAggregationFunctionDefinitions) {
-      for (const { params, returnType } of signatures) {
+      for (const [signatureIndex, { params, returnType }] of Object.entries(signatures)) {
         const fieldMapping = getFieldMapping(params);
         testErrorsAndWarnings(
           `from a | stats var = ${
@@ -1355,7 +1374,9 @@ describe('validation logic', () => {
           // and the message is case of wrong argument type is passed
           const wrongFieldMapping = params.map(({ name: _name, type, ...rest }) => {
             const typeString = type;
-            const canBeFieldButNotString = ['number', 'date', 'boolean', 'ip'].includes(typeString);
+            const canBeFieldButNotString = fieldTypes
+              .filter((t) => t !== 'string')
+              .includes(typeString);
             const isLiteralType = /literal$/.test(typeString);
             // pick a field name purposely wrong
             const nameValue =
@@ -1365,9 +1386,13 @@ describe('validation logic', () => {
 
           const expectedErrors = params.map(
             ({ type }, i) =>
-              `Argument of [${name}] must be [${type}], found value [${
-                wrongFieldMapping[i].name
-              }] type [${wrongFieldMapping[i].name.replace('Field', '')}]`
+              `Argument of [${name}] must be [${
+                // If the function has multiple signatures and all fail, then only
+                // one error will be reported for the first signature type
+                +signatureIndex > 0 ? signatures[0].params[i].type : type
+              }], found value [${wrongFieldMapping[i].name}] type [${wrongFieldMapping[
+                i
+              ].name.replace('Field', '')}]`
           );
           testErrorsAndWarnings(
             `from a | stats ${

diff --git a/packages/kbn-monaco/src/esql/lib/ast/validation/validation.ts b/packages/kbn-monaco/src/esql/lib/ast/validation/validation.ts
@@ -130,7 +130,18 @@ function validateNestedFunctionArg(
     // no need to check the reason here, it is checked already above
     isSupportedFunction(actualArg.name, parentCommand).supported
   ) {
+    // The isSupported check ensure the definition exists
     const argFn = getFunctionDefinition(actualArg.name)!;
+
+    if ('noNestingFunctions' in argDef && argDef.noNestingFunctions) {
+      messages.push(
+        getMessageFromId({
+          messageId: 'noNestedArgumentSupport',
+          values: { name: actualArg.text, argType: argFn.signatures[0].returnType },
+          locations: actualArg.location,
+        })
+      );
+    }
     if (!isEqualType(actualArg, argDef, references, parentCommand)) {
       messages.push(
         getMessageFromId({
@@ -144,16 +155,6 @@ function validateNestedFunctionArg(
           locations: actualArg.location,
         })
       );
-    } else {
-      if ('noNestingFunctions' in argDef && argDef.noNestingFunctions) {
-        messages.push(
-          getMessageFromId({
-            messageId: 'noNestedArgumentSupport',
-            values: { name: actualArg.text, argType: argFn.signatures[0].returnType },
-            locations: actualArg.location,
-          })
-        );
-      }
     }
   }
   return messages;