[ES|QL] Simplify function to retrieve categorize column (#225281)

## Summary

Simplifies the function to extract the categorize function field. The
unit tests we had pass for the new implementation so I am sure we don't
introduce new bugs

src/platform/packages/shared/kbn-esql-utils/src/utils/query_parsing_helpers.ts

@@ -369,47 +369,17 @@ export const getArgsFromRenameFunction = (
  */
 export const getCategorizeField = (esql: string): string[] => {
   const { root } = parse(esql);
-  const statsCommand = root.commands.find(({ name }) => name === 'stats');
-  if (!statsCommand) {
-    return [];
-  }
-  const options: ESQLCommandOption[] = [];
   const columns: string[] = [];
-
+  const functions = Walker.matchAll(root.commands, {
-  walk(statsCommand, {
+    type: 'function',
-    visitCommandOption: (node) => options.push(node),
+    name: 'categorize',
-  });
-
-  const statsByOptions = options.find(({ name }) => name === 'by');
-
-  // categorize is part of the stats by command
-  if (!statsByOptions) {
-    return [];
-  }
-
-  const categorizeOptions = statsByOptions.args.filter((arg) => {
-    return (arg as ESQLFunction).text.toLowerCase().indexOf('categorize') !== -1;
   }) as ESQLFunction[];
 
-  if (categorizeOptions.length) {
+  if (functions.length) {
-    categorizeOptions.forEach((arg) => {
+    functions.forEach((func) => {
-      // ... STATS ... BY CATEGORIZE(field)
+      for (const arg of func.args) if (isColumn(arg)) columns.push(arg.name);
-      if (isFunctionExpression(arg) && arg.name === 'categorize') {
-        walk(arg, {
-          visitColumn: (node) => columns.push(node.name),
-        });
-      } else {
-        // ... STATS ... BY pattern = CATEGORIZE(field)
-        walk(arg, {
-          visitFunction: (node) => {
-            if (node.name === 'categorize') {
-              const columnArgs = node.args.filter((a) => isColumn(a));
-              columnArgs.forEach((c) => columns.push((c as ESQLColumn).name));
-            }
-          },
-        });
-      }
     });
+    return columns;
   }
 
   return columns;