Add support for top values

src/common/schema.ts

@@ -21,7 +21,13 @@
  * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
  */
 
-import {Explore, Field, JoinRelationship} from '@malloydata/malloy';
+import {
+  Explore,
+  Field,
+  JoinRelationship,
+  ModelDef,
+  isSourceDef,
+} from '@malloydata/malloy';
 
 export function isFieldAggregate(field: Field) {
   return field.isAtomicField() && field.isCalculation();
@@ -114,6 +120,22 @@ export const quoteIfNecessary = (element: string) => {
   return element;
 };
 
+/**
+ * Retrieve a source from a model safely
+ *
+ * @param modelDef Model definition
+ * @param sourceName Source name
+ * @returns SourceDef for given name, or throws if not a source
+ */
+
+export const getSourceDef = (modelDef: ModelDef, sourceName: string) => {
+  const result = modelDef.contents[sourceName];
+  if (isSourceDef(result)) {
+    return result;
+  }
+  throw new Error(`Not a source: ${sourceName}`);
+};
+
 const RESERVED: string[] = [
   'ALL',
   'AND',

src/common/types/message_types.ts

@@ -287,6 +287,7 @@ export enum ComposerMessageType {
   NewModel = 'new-model',
   ResultSuccess = 'result-success',
   ResultError = 'result-error',
+  SearchIndex = 'search-index',
 }
 
 export interface ComposerMessageNewModel {
@@ -309,10 +310,16 @@ export interface ComposerMessageResultError {
   error: string;
 }
 
+export interface ComposerMessageSearchIndex {
+  type: ComposerMessageType.SearchIndex;
+  result: ResultJSON;
+}
+
 export type ComposerMessage =
   | ComposerMessageNewModel
   | ComposerMessageResultSuccess
-  | ComposerMessageResultError;
+  | ComposerMessageResultError
+  | ComposerMessageSearchIndex;
 
 export enum ComposerPageMessageType {
   Ready = 'ready',

src/extension/commands/open_composer.ts

@@ -7,21 +7,11 @@
 
 import * as vscode from 'vscode';
 import {getWebviewHtml} from '../webviews';
-import {
-  ComposerMessage,
-  ComposerMessageType,
-  ComposerPageMessage,
-  ComposerPageMessageType,
-} from '../../common/types/message_types';
-import {
-  getActiveDocumentMetadata,
-  runMalloyQuery,
-} from './utils/run_query_utils';
+import {getActiveDocumentMetadata} from './utils/run_query_utils';
 import {Utils} from 'vscode-uri';
 import {MALLOY_EXTENSION_STATE} from '../state';
 import {WorkerConnection} from '../worker_connection';
-import {WebviewMessageManager} from '../webview_message_manager';
-import {QuerySpec} from '../../common/types/query_spec';
+import {ComposerMessageManager} from './utils/composer_message_manager';
 
 const icon = 'turtle.svg';
 
@@ -59,74 +49,17 @@ export async function openComposer(
 
     sourceName ??= Object.keys(modelDef.contents)[0];
 
-    const messages = new WebviewMessageManager<
-      ComposerMessage,
-      ComposerPageMessage
-    >(composerPanel);
-    messages.postMessage({
-      type: ComposerMessageType.NewModel,
+    const messageManager = new ComposerMessageManager(
+      worker,
+      composerPanel,
       documentMeta,
       modelDef,
       sourceName,
-      viewName,
-    });
-    messages.onReceiveMessage(message => {
-      switch (message.type) {
-        case ComposerPageMessageType.RunQuery:
-          {
-            const {id, query, queryName} = message;
-            vscode.window
-              .withProgress(
-                {
-                  location: vscode.ProgressLocation.Notification,
-                  title: `Running ${queryName}`,
-                  cancellable: true,
-                },
-                async (progress, cancellationToken) => {
-                  const querySpec: QuerySpec = {
-                    type: 'string',
-                    text: query,
-                    documentMeta,
-                  };
-                  const result = await runMalloyQuery(
-                    worker,
-                    querySpec,
-                    id,
-                    queryName,
-                    {withWebview: false},
-                    cancellationToken,
-                    progress
-                  );
-                  return result;
-                }
-              )
-              .then(
-                result => {
-                  if (result) {
-                    messages.postMessage({
-                      type: ComposerMessageType.ResultSuccess,
-                      id,
-                      result,
-                    });
-                  } else {
-                    messages.postMessage({
-                      type: ComposerMessageType.ResultError,
-                      id,
-                      error: 'No results',
-                    });
-                  }
-                },
-                error => {
-                  messages.postMessage({
-                    type: ComposerMessageType.ResultError,
-                    id,
-                    error: error instanceof Error ? error.message : `${error}`,
-                  });
-                }
-              );
-          }
-          break;
-      }
+      viewName
+    );
+
+    composerPanel.onDidDispose(() => {
+      messageManager.dispose();
     });
   }
 }

src/extension/commands/utils/composer_message_manager.ts

@@ -0,0 +1,151 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ *  LICENSE file in the root directory of this source tree.
+ */
+
+import * as vscode from 'vscode';
+import {v1 as uuid} from 'uuid';
+import {WebviewMessageManager} from '../../webview_message_manager';
+import {
+  ComposerMessage,
+  ComposerMessageType,
+  ComposerPageMessage,
+  ComposerPageMessageRunQuery,
+  ComposerPageMessageType,
+} from '../../../common/types/message_types';
+import {DocumentMetadata, QuerySpec} from '../../../common/types/query_spec';
+import {ModelDef, SearchValueMapResult} from '@malloydata/malloy';
+import {runMalloyQuery} from './run_query_utils';
+import {WorkerConnection} from '../../worker_connection';
+import {getSourceDef} from '../../../common/schema';
+
+export class ComposerMessageManager
+  extends WebviewMessageManager<ComposerMessage, ComposerPageMessage>
+  implements vscode.Disposable
+{
+  constructor(
+    private worker: WorkerConnection,
+    composerPanel: vscode.WebviewPanel,
+    private documentMeta: DocumentMetadata,
+    private modelDef: ModelDef,
+    private sourceName: string,
+    viewName?: string
+  ) {
+    super(composerPanel);
+
+    this.postMessage({
+      type: ComposerMessageType.NewModel,
+      documentMeta,
+      modelDef,
+      sourceName,
+      viewName,
+    });
+
+    this.onReceiveMessage(async message => {
+      switch (message.type) {
+        case ComposerPageMessageType.RunQuery:
+          await this.runQuery(message);
+          break;
+      }
+    });
+
+    void this.initializeIndex();
+  }
+
+  async runQuery(message: ComposerPageMessageRunQuery) {
+    {
+      const {id, query, queryName} = message;
+      try {
+        const result = await this.runQueryWithProgress(id, queryName, query);
+        if (result) {
+          this.postMessage({
+            type: ComposerMessageType.ResultSuccess,
+            id,
+            result,
+          });
+        } else {
+          this.postMessage({
+            type: ComposerMessageType.ResultError,
+            id,
+            error: 'No results',
+          });
+        }
+      } catch (error) {
+        this.postMessage({
+          type: ComposerMessageType.ResultError,
+          id,
+          error: error instanceof Error ? error.message : `${error}`,
+        });
+      }
+    }
+  }
+
+  async initializeIndex(): Promise<SearchValueMapResult[] | undefined> {
+    const sourceDef = getSourceDef(this.modelDef, this.sourceName);
+    const indexQuery = sourceDef.fields.find(
+      ({name, as}) => (as || name) === 'search_index'
+    );
+    const limit = 10;
+
+    if (indexQuery) {
+      const searchMapMalloy = `
+        run: ${this.sourceName}
+          -> ${indexQuery.as || indexQuery.name}
+          -> {
+            where: fieldType = 'string'
+            group_by: fieldName
+            aggregate: cardinality is count(fieldValue)
+            nest: values is {
+              select: fieldValue, weight
+              order_by: weight desc
+              limit: ${limit}
+            }
+            limit: 1000
+          }
+      `;
+      const result = await this.runQueryWithProgress(
+        uuid(),
+        `Search index for ${this.sourceName}`,
+        searchMapMalloy
+      );
+      if (result) {
+        this.postMessage({
+          type: ComposerMessageType.SearchIndex,
+          result,
+        });
+      }
+    }
+    return undefined;
+  }
+
+  async runQueryWithProgress(id: string, queryName: string, query: string) {
+    return vscode.window.withProgress(
+      {
+        location: vscode.ProgressLocation.Notification,
+        title: queryName,
+        cancellable: true,
+      },
+      async (progress, cancellationToken) => {
+        const querySpec: QuerySpec = {
+          type: 'string',
+          text: query,
+          documentMeta: this.documentMeta,
+        };
+        const result = await runMalloyQuery(
+          this.worker,
+          querySpec,
+          id,
+          queryName,
+          {withWebview: false},
+          cancellationToken,
+          progress
+        );
+        return result;
+      }
+    );
+  }
+
+  dispose() {}
+}

src/extension/webviews/composer_page/App.tsx

@@ -16,7 +16,7 @@ import {
 } from '../../../common/types/message_types';
 import {Composer} from './Composer';
 import {DocumentMetadata} from '../../../common/types/query_spec';
-import {ModelDef, Result} from '@malloydata/malloy';
+import {ModelDef, Result, SearchValueMapResult} from '@malloydata/malloy';
 import {RunQuery} from '@malloydata/query-composer';
 
 export interface AppProps {
@@ -35,6 +35,7 @@ export const App: React.FC<AppProps> = ({vscode}) => {
   const [modelDef, setModelDef] = React.useState<ModelDef>();
   const [sourceName, setSourceName] = React.useState<string>();
   const [viewName, setViewName] = React.useState<string>();
+  const [topValues, setTopValues] = React.useState<SearchValueMapResult[]>();
 
   React.useEffect(() => {
     const messageHandler = ({data}: MessageEvent<ComposerMessage>) => {
@@ -59,13 +60,23 @@ export const App: React.FC<AppProps> = ({vscode}) => {
             }
           }
           break;
-        case ComposerMessageType.ResultError: {
-          const {id, error} = data;
-          if (QueriesInFlight[id]) {
-            QueriesInFlight[id]?.reject(new Error(error));
-            delete QueriesInFlight[id];
+        case ComposerMessageType.ResultError:
+          {
+            const {id, error} = data;
+            if (QueriesInFlight[id]) {
+              QueriesInFlight[id]?.reject(new Error(error));
+              delete QueriesInFlight[id];
+            }
           }
-        }
+          break;
+        case ComposerMessageType.SearchIndex:
+          {
+            const result = Result.fromJSON(data.result);
+            setTopValues(
+              result._queryResult.result as unknown as SearchValueMapResult[]
+            );
+          }
+          break;
       }
     };
     window.addEventListener('message', messageHandler);
@@ -105,6 +116,7 @@ export const App: React.FC<AppProps> = ({vscode}) => {
           sourceName={sourceName}
           viewName={viewName}
           runQuery={runQuery}
+          topValues={topValues}
         />
       </div>
     );

src/extension/webviews/composer_page/Composer.tsx

@@ -24,10 +24,9 @@ export interface ComposerProps {
   sourceName: string;
   viewName?: string;
   runQuery: RunQuery;
+  topValues: SearchValueMapResult[] | undefined;
 }
 
-const topValues: SearchValueMapResult[] = [];
-
 const nullUpdateQueryInUrl = () => {};
 
 export const Composer: React.FC<ComposerProps> = ({
@@ -36,6 +35,7 @@ export const Composer: React.FC<ComposerProps> = ({
   sourceName,
   viewName,
   runQuery: runQueryImp,
+  topValues,
 }) => {
   const {queryMalloy, queryName, queryModifiers, querySummary} =
     useQueryBuilder(