From 41cc82b8d032129d2570507c204a0a0f2e8fa955 Mon Sep 17 00:00:00 2001
From: Josh Hunt <joshhunt@users.noreply.github.com>
Date: Mon, 24 Mar 2025 15:57:00 +0000
Subject: [PATCH] Combobox: Refactor sortByGroup for performance (#102664)

* first pass at improving perf of sortByGroup

* polish up
---
 .../components/Combobox/useOptions.test.ts    | 55 ++++++++++++++++-
 .../src/components/Combobox/useOptions.ts     | 60 ++++++++++++-------
 2 files changed, 94 insertions(+), 21 deletions(-)

diff --git a/packages/grafana-ui/src/components/Combobox/useOptions.test.ts b/packages/grafana-ui/src/components/Combobox/useOptions.test.ts
index ff5f5ae61b0..a80d6f01168 100644
--- a/packages/grafana-ui/src/components/Combobox/useOptions.test.ts
+++ b/packages/grafana-ui/src/components/Combobox/useOptions.test.ts
@@ -1,6 +1,6 @@
 import { renderHook, act, waitFor } from '@testing-library/react';
 
-import { useOptions } from './useOptions';
+import { sortByGroup, useOptions } from './useOptions';
 
 describe('useOptions', () => {
   it('should handle a large number of synchronous options without throwing an error', () => {
@@ -101,3 +101,56 @@ describe('useOptions', () => {
     expect(result.current.asyncError).toBe(true);
   });
 });
+
+describe('sortByGroup', () => {
+  it('should return original array when no groups exist', () => {
+    const options = [
+      { label: 'Apple', value: 'apple' },
+      { label: 'Banana', value: 'banana' },
+      { label: 'Carrot', value: 'carrot' },
+    ];
+
+    const { options: sortedOptions, groupStartIndices } = sortByGroup(options);
+
+    expect(sortedOptions).toBe(options); // Check reference equality
+    expect(groupStartIndices.size).toBe(0);
+  });
+
+  it('should return original array when only one group exists', () => {
+    const options = [
+      { label: 'Apple', value: 'apple', group: 'fruits' },
+      { label: 'Banana', value: 'banana', group: 'fruits' },
+      { label: 'Tomato', value: 'tomato', group: 'fruits' },
+    ];
+
+    const { options: sortedOptions, groupStartIndices } = sortByGroup(options);
+
+    expect(sortedOptions).toEqual(options);
+    expect(groupStartIndices.size).toBe(1);
+    expect(groupStartIndices.get('fruits')).toBe(0);
+  });
+
+  it('should group options and track group start indices', () => {
+    const options = [
+      { label: 'Apple', value: 'apple', group: 'fruits' },
+      { label: 'Carrot', value: 'carrot', group: 'vegetables' },
+      { label: 'Banana', value: 'banana', group: 'fruits' },
+      { label: 'Celery', value: 'celery', group: 'vegetables' },
+      { label: 'Other', value: 'other' }, // Ungrouped
+    ];
+
+    const { options: sortedOptions, groupStartIndices } = sortByGroup(options);
+
+    expect(sortedOptions).toEqual([
+      { label: 'Apple', value: 'apple', group: 'fruits' },
+      { label: 'Banana', value: 'banana', group: 'fruits' },
+      { label: 'Carrot', value: 'carrot', group: 'vegetables' },
+      { label: 'Celery', value: 'celery', group: 'vegetables' },
+      { label: 'Other', value: 'other' },
+    ]);
+
+    expect(groupStartIndices.size).toBe(2);
+    expect(groupStartIndices.get('fruits')).toBe(0);
+    expect(groupStartIndices.get('vegetables')).toBe(2);
+  });
+});
diff --git a/packages/grafana-ui/src/components/Combobox/useOptions.ts b/packages/grafana-ui/src/components/Combobox/useOptions.ts
index 59ff783786b..bedd6dc1b8b 100644
--- a/packages/grafana-ui/src/components/Combobox/useOptions.ts
+++ b/packages/grafana-ui/src/components/Combobox/useOptions.ts
@@ -115,41 +115,61 @@ export function useOptions<T extends string | number>(rawOptions: AsyncOptions<T
   return { options: finalOptions, groupStartIndices, updateOptions, asyncLoading, asyncError };
 }
 
-function sortByGroup<T extends string | number>(options: Array<ComboboxOption<T>>) {
+/**
+ * Sorts options by group and returns the sorted options and the starting index of each group
+ */
+export function sortByGroup<T extends string | number>(options: Array<ComboboxOption<T>>) {
+  // Group options by their group
   const groupedOptions = new Map<string | undefined, Array<ComboboxOption<T>>>();
+  const groupStartIndices = new Map<string | undefined, number>();
+
   for (const option of options) {
-    const groupExists = groupedOptions.has(option.group);
-    if (groupExists) {
-      groupedOptions.get(option.group)?.push(option);
+    const group = option.group;
+    const existing = groupedOptions.get(group);
+    if (existing) {
+      existing.push(option);
     } else {
-      groupedOptions.set(option.group, [option]);
+      groupedOptions.set(group, [option]);
     }
   }
 
-  // Create a map to track the starting index of each group
-  const groupStartIndices = new Map<string, number>();
+  // If we only have one group (either the undefined group, or a single group), return the original array
+  if (groupedOptions.size <= 1) {
+    if (options[0]?.group) {
+      groupStartIndices.set(options[0]?.group, 0);
+    }
+
+    return {
+      options,
+      groupStartIndices,
+    };
+  }
+
+  // 'Preallocate' result array with same size as input - very minor optimization
+  const result: Array<ComboboxOption<T>> = new Array(options.length);
+
   let currentIndex = 0;
 
-  // Reorganize options to have groups first, then undefined group
-  let reorganizeOptions: Array<ComboboxOption<T>> = [];
+  // Fill result array with grouped options
   for (const [group, groupOptions] of groupedOptions) {
-    if (!group) {
-      continue;
+    if (group) {
+      groupStartIndices.set(group, currentIndex);
+      for (const option of groupOptions) {
+        result[currentIndex++] = option;
+      }
     }
-
-    groupStartIndices.set(group, currentIndex);
-    reorganizeOptions = reorganizeOptions.concat(groupOptions);
-    currentIndex += groupOptions.length;
   }
 
-  const undefinedGroupOptions = groupedOptions.get(undefined);
-  if (undefinedGroupOptions) {
-    groupStartIndices.set('undefined', currentIndex);
-    reorganizeOptions = reorganizeOptions.concat(undefinedGroupOptions);
+  // Add ungrouped options at the end
+  const ungrouped = groupedOptions.get(undefined);
+  if (ungrouped) {
+    for (const option of ungrouped) {
+      result[currentIndex++] = option;
+    }
   }
 
   return {
-    options: reorganizeOptions,
+    options: result,
     groupStartIndices,
   };
 }