Allow SelectGroupSlices for QP

app/packages/state/src/recoil/queryPerformance.test.ts

@@ -3,7 +3,10 @@ import { describe, expect, it, vi } from "vitest";
 vi.mock("recoil");
 vi.mock("recoil-relay");
 
-import type { TestSelectorFamily } from "../../../../__mocks__/recoil";
+import type {
+  TestSelector,
+  TestSelectorFamily,
+} from "../../../../__mocks__/recoil";
 import { setMockAtoms } from "../../../../__mocks__/recoil";
 import * as queryPerformance from "./queryPerformance";
 
@@ -23,4 +26,38 @@ describe("tests query performance selectors", () => {
 
     expect(test()).toEqual(new Set(["ground_truth.id", "ground_truth.label"]));
   });
+
+  it("resolves query performant views", () => {
+    const test = <TestSelector<typeof queryPerformance.isQueryPerformantView>>(
+      (<unknown>queryPerformance.isQueryPerformantView)
+    );
+
+    setMockAtoms({
+      _view__setter: [],
+    });
+    expect(test()).toBe(true);
+
+    setMockAtoms({
+      _view__setter: [
+        {
+          _cls: "fiftyone.core.stages.SelectGroupSlices",
+        },
+      ],
+    });
+    expect(test()).toBe(true);
+
+    setMockAtoms({
+      _view__setter: [
+        {
+          _cls: "unsupported",
+        },
+      ],
+    });
+    expect(test()).toBe(false);
+
+    setMockAtoms({
+      _view__setter: [{}, {}],
+    });
+    expect(test()).toBe(false);
+  });
 });

app/packages/state/src/recoil/queryPerformance.ts

@@ -19,6 +19,9 @@ import { datasetId, datasetName } from "./selectors";
 import { State } from "./types";
 import { view } from "./view";
 
+const SELECT_GROUP_SLICES = "fiftyone.core.stages.SelectGroupSlices";
+const VALID_QP_STAGES = new Set([SELECT_GROUP_SLICES]);
+
 export const lightningQuery = graphQLSelectorFamily<
   foq.lightningQuery$variables,
   foq.LightningInput["paths"],
@@ -188,6 +191,22 @@ export const indexedPaths = selectorFamily<Set<string>, string>({
     },
 });
 
+export const isQueryPerformantView = selector({
+  key: "isQueryPerformantView",
+  get: ({ get }) => {
+    const stages = get(view);
+    if (!stages?.length) {
+      return true;
+    }
+
+    if (stages.length === 1) {
+      return VALID_QP_STAGES.has(stages[0]._cls);
+    }
+
+    return false;
+  },
+});
+
 export const enableQueryPerformanceConfig = selector({
   key: "enableQueryPerformanceConfig",
   get: ({ get }) => get(config).enableQueryPerformance,
@@ -198,24 +217,15 @@ export const defaultQueryPerformanceConfig = selector({
   get: ({ get }) => get(config).defaultQueryPerformance,
 });
 
-const queryPerformanceStore = atomFamily<boolean, string>({
-  key: "queryPerformanceStore",
-  default: undefined,
-  effects: (datasetId) => [
-    getBrowserStorageEffectForKey(`queryPerformance-${datasetId}`, {
-      sessionStorage: true,
-      valueClass: "boolean",
-    }),
-  ],
-});
-
 export const queryPerformance = selector<boolean>({
   key: "queryPerformance",
-  get: ({ get }) => {
-    if (get(view).length) {
-      return false;
-    }
+  get: ({ get }) => get(queryPerformanceSetting) && get(isQueryPerformantView),
+  set: ({ set }, value) => set(queryPerformanceSetting, value),
+});
 
+export const queryPerformanceSetting = selector<boolean>({
+  key: "queryPerformanceSetting",
+  get: ({ get }) => {
     if (!get(enableQueryPerformanceConfig)) {
       return false;
     }
@@ -234,3 +244,14 @@ export const queryPerformance = selector<boolean>({
     );
   },
 });
+
+const queryPerformanceStore = atomFamily<boolean, string>({
+  key: "queryPerformanceStore",
+  default: undefined,
+  effects: (datasetId) => [
+    getBrowserStorageEffectForKey(`queryPerformance-${datasetId}`, {
+      sessionStorage: true,
+      valueClass: "boolean",
+    }),
+  ],
+});