From fad7bc2d4029452c5d94a4b0c814ffe7abb68754 Mon Sep 17 00:00:00 2001
From: Martin Mark <martin@martinmark.com>
Date: Mon, 23 Dec 2024 12:22:48 -0500
Subject: [PATCH] Adding operator group headers

---
 .../Home/Browse3/filters/FilterRow.tsx        |  4 +-
 .../Home/Browse3/filters/SelectOperator.tsx   | 16 ++++--
 .../Home/Browse3/filters/common.ts            | 52 +++++++++++++++++++
 3 files changed, 66 insertions(+), 6 deletions(-)

diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/filters/FilterRow.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/filters/FilterRow.tsx
index 8b1fb7b6022..5a20a8ede48 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/filters/FilterRow.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/filters/FilterRow.tsx
@@ -6,7 +6,7 @@ import {GridFilterItem} from '@mui/x-data-grid-pro';
 import React, {useMemo} from 'react';
 
 import {Button} from '../../../../Button';
-import {FilterId, getFieldType, getOperatorOptions, isWeaveRef} from './common';
+import {FilterId, getFieldType, getGroupedOperatorOptions, isWeaveRef} from './common';
 import {SelectField, SelectFieldOption} from './SelectField';
 import {SelectOperator} from './SelectOperator';
 import {SelectValue} from './SelectValue';
@@ -36,7 +36,7 @@ export const FilterRow = ({
   };
 
   const operatorOptions = useMemo(
-    () => getOperatorOptions(item.field),
+    () => getGroupedOperatorOptions(item.field),
     [item.field]
   );
 
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/filters/SelectOperator.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/filters/SelectOperator.tsx
index 01966b1fa7c..8e26758252c 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/filters/SelectOperator.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/filters/SelectOperator.tsx
@@ -5,10 +5,11 @@ import {Select} from '@wandb/weave/components/Form/Select';
 import React from 'react';
 
 import {Tooltip} from '../../../../Tooltip';
-import {SelectOperatorOption} from './common';
+import {OperatorGroupedOption, SelectOperatorOption} from './common';
+import {components, GroupHeadingProps} from 'react-select';
 
 type SelectOperatorProps = {
-  options: SelectOperatorOption[];
+  options: OperatorGroupedOption[];
   value: string;
   onSelectOperator: (value: string) => void;
   isDisabled?: boolean;
@@ -24,13 +25,19 @@ const OptionLabel = (props: SelectOperatorOption) => {
   );
 };
 
+const GroupHeading = (props: GroupHeadingProps<SelectOperatorOption, false, OperatorGroupedOption>) => {
+  return <components.GroupHeading {...props} />;
+};
+
 export const SelectOperator = ({
   options,
   value,
   onSelectOperator,
   isDisabled,
 }: SelectOperatorProps) => {
-  const selectedOption = options.find(o => o.value === value) ?? options[0];
+  // Find the operator from the grouped selection:
+  const flattenedOptions = options.flatMap(group => group.options);
+  const selectedOption = flattenedOptions.find(o => o.value === value) ?? flattenedOptions[0];
 
   const onReactSelectChange = (option: SelectOperatorOption | null) => {
     if (option) {
@@ -39,13 +46,14 @@ export const SelectOperator = ({
   };
 
   return (
-    <Select<SelectOperatorOption>
+    <Select<SelectOperatorOption, false, OperatorGroupedOption>
       options={options}
       value={selectedOption}
       placeholder="Select operator"
       onChange={onReactSelectChange}
       formatOptionLabel={OptionLabel}
       isDisabled={isDisabled}
+      components={{GroupHeading}}
     />
   );
 };
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/filters/common.ts b/weave-js/src/components/PagePanelComponents/Home/Browse3/filters/common.ts
index 17cfbf45ac4..328f6a0d0d9 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/filters/common.ts
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/filters/common.ts
@@ -301,3 +301,55 @@ export const upsertFilter = (
     items,
   };
 };
+
+export type OperatorGroupedOption = {
+  label: string;
+  options: SelectOperatorOption[];
+};
+
+/** 
+ * Return grouped operators by type (string, number, bool, etc.). 
+ * Customize the group labels & operators as needed. 
+ */
+export function getGroupedOperatorOptions(field: string): OperatorGroupedOption[] {
+  const stringOperators: SelectOperatorOption[] = [
+    {value: '(string): contains', label: 'contains'},
+    {value: '(string): equals', label: 'equals'},
+    {value: '(string): in', label: 'in'},
+  ];
+  const numberOperators: SelectOperatorOption[] = [
+    {value: '(number): =', label: '='},
+    {value: '(number): !=', label: '≠'},
+    {value: '(number): <', label: '<'},
+    {value: '(number): <=', label: '≤'},
+    {value: '(number): >', label: '>'},
+    {value: '(number): >=', label: '≥'},
+  ];
+  const dateOperators: SelectOperatorOption[] = [
+    {value: '(date): after', label: 'after'},
+    {value: '(date): before', label: 'before'},
+  ];
+  const boolOperators: SelectOperatorOption[] = [
+    {value: '(bool): is', label: 'is'},
+  ];
+  const anyOperators: SelectOperatorOption[] = [
+    {value: '(any): isEmpty', label: 'is empty'},
+    {value: '(any): isNotEmpty', label: 'is not empty'},
+  ];
+
+  // Select which groups/ops to show based on the field type
+  const fieldType = getFieldType(field);
+  if (fieldType === 'datetime') {
+    return [{label: 'Date Operators', options: dateOperators}];
+  } else if (fieldType === 'user') {
+    return [{label: 'String Operators', options: stringOperators}];
+  }
+  // Fallback: show all grouped
+  return [
+    {label: 'String', options: stringOperators},
+    {label: 'Number', options: numberOperators},
+    {label: 'Boolean', options: boolOperators},
+    {label: 'Date', options: dateOperators},
+    {label: 'Other', options: anyOperators},
+  ];
+}