Category count query fix

change/@itwin-tree-widget-react-e0e85de3-ea1f-48ff-b8da-1cb2f72be783.json

@@ -0,0 +1,7 @@
+{
+  "type": "patch",
+  "comment": "Fix category count query throwing error on large iModels",
+  "packageName": "@itwin/tree-widget-react",
+  "email": "100586436+JonasDov@users.noreply.github.com",
+  "dependentChangeType": "patch"
+}

packages/itwin/tree-widget/src/test/trees/models-tree/internal/ModelsTreeIdsCache.test.ts

@@ -22,7 +22,7 @@ describe("ModelsTreeIdsCache", () => {
     const categoryId = "0x2";
     const elementIds = ["0x10", "0x20", "0x30"];
     const stub = sinon.fake((query: string) => {
-      if (query.includes(`WHERE Parent.Id IS NULL AND (Model.Id = ${modelId} AND Category.Id = ${categoryId})`)) {
+      if (query.includes(`WHERE Parent.Id IS NULL AND (Model.Id = ${modelId} AND Category.Id IN (${categoryId}))`)) {
         return [{ modelId, categoryId, elementsCount: elementIds.length }];
       }
       throw new Error(`Unexpected query: ${query}`);

packages/itwin/tree-widget/src/test/trees/models-tree/internal/ModelsTreeVisibilityHandler.test.ts

@@ -2196,6 +2196,35 @@ describe("ModelsTreeVisibilityHandler", () => {
       });
     });
 
+    it("validates visibility for large iModel", async function () {
+      await using buildIModelResult = await buildIModel(this, async (builder) => {
+        const modelsToTurnOn = new Array<string>();
+        let elementId = "";
+        for (let i = 0; i < 3; ++i) {
+          const model = insertPhysicalModelWithPartition({ builder, partitionParentId: IModel.rootSubjectId, codeValue: `model${i}` }).id;
+          modelsToTurnOn.push(model);
+          for (let j = 0; j < 1000; ++j) {
+            const categoryId = insertSpatialCategory({ builder, codeValue: `category${i}-${j}` }).id;
+            elementId = insertPhysicalElement({ builder, modelId: model, categoryId }).id;
+          }
+        }
+        return { modelsToTurnOn, elementId };
+      });
+
+      const { imodel, ...ids } = buildIModelResult;
+      using visibilityTestData = createVisibilityTestData({ imodel });
+      const { handler, provider, viewport } = visibilityTestData;
+      await Promise.all(ids.modelsToTurnOn.map(async (modelId) => handler.changeVisibility(createModelHierarchyNode(modelId), true)));
+      viewport.setAlwaysDrawn(new Set([ids.elementId]));
+      viewport.renderFrame();
+      await validateHierarchyVisibility({
+        provider,
+        handler,
+        viewport,
+        visibilityExpectations: VisibilityExpectations.all("visible"),
+      });
+    });
+
     it("showing model makes it, all its categories and elements visible and doesn't affect other models", async function () {
       await using buildIModelResult = await buildIModel(this, async (builder) => {
         const categoryId = insertSpatialCategory({ builder, codeValue: "category" }).id;

packages/itwin/tree-widget/src/tree-widget-react/components/trees/models-tree/internal/ModelsTreeIdsCache.ts

@@ -4,9 +4,10 @@
  *--------------------------------------------------------------------------------------------*/
 
 import type { Subscription } from "rxjs";
-import { bufferTime, filter, firstValueFrom, mergeAll, mergeMap, ReplaySubject, Subject } from "rxjs";
+import { bufferCount, bufferTime, filter, firstValueFrom, from, mergeAll, mergeMap, ReplaySubject, Subject } from "rxjs";
 import { assert, Id64 } from "@itwin/core-bentley";
 import { IModel } from "@itwin/core-common";
+import { collect } from "../../common/Rxjs.js";
 import { pushToMap } from "../../common/Utils.js";
 
 import type { InstanceKey } from "@itwin/presentation-shared";
@@ -388,41 +389,63 @@ export class ModelsTreeIdsCache {
   private async queryCategoryElementCounts(
     input: Array<{ modelId: Id64String; categoryId: Id64String }>,
   ): Promise<Array<{ modelId: number; categoryId: number; elementsCount: number }>> {
-    const reader = this._queryExecutor.createQueryReader(
-      {
-        ctes: [
-          /* sql */ `
-            CategoryElements(id, modelId, categoryId) AS (
-              SELECT ECInstanceId, Model.Id, Category.Id
-              FROM ${this._hierarchyConfig.elementClassSpecification}
-              WHERE
-                Parent.Id IS NULL
-                AND (
-                  ${input.map(({ modelId, categoryId }) => `Model.Id = ${modelId} AND Category.Id = ${categoryId}`).join(" OR ")}
-                )
-
-              UNION ALL
-
-              SELECT c.ECInstanceId, p.modelId, p.categoryId
-              FROM ${this._hierarchyConfig.elementClassSpecification} c
-              JOIN CategoryElements p ON c.Parent.Id = p.id
-            )
-          `,
-        ],
-        ecsql: `
-          SELECT modelId, categoryId, COUNT(id) elementsCount
-          FROM CategoryElements
-          GROUP BY modelId, categoryId
-        `,
-      },
-      { rowFormat: "ECSqlPropertyNames", limit: "unbounded" },
-    );
-
-    const result = new Array<{ modelId: number; categoryId: number; elementsCount: number }>();
-    for await (const row of reader) {
-      result.push({ modelId: row.modelId, categoryId: row.categoryId, elementsCount: row.elementsCount });
+    const modelCategoryMap = new Map<Id64String, Id64Set>();
+    for (const { modelId, categoryId } of input) {
+      const entry = modelCategoryMap.get(modelId);
+      if (!entry) {
+        modelCategoryMap.set(modelId, new Set([categoryId]));
+      } else {
+        entry.add(categoryId);
+      }
     }
-    return result;
+    const modelCategoryWhereClauses = new Array<string>();
+    for (const [modelId, categoryIds] of modelCategoryMap) {
+      modelCategoryWhereClauses.push(`Model.Id = ${modelId} AND Category.Id IN (${[...categoryIds].join(", ")})`);
+    }
+
+    return collect(
+      from(modelCategoryWhereClauses).pipe(
+        bufferCount(Math.ceil(modelCategoryWhereClauses.length / Math.ceil(modelCategoryWhereClauses.length / 2900))),
+        mergeMap(async (whereClauses) => {
+          const reader = this._queryExecutor.createQueryReader(
+            {
+              ctes: [
+                `
+                  CategoryElements(id, modelId, categoryId) AS (
+                    SELECT ECInstanceId, Model.Id, Category.Id
+                    FROM ${this._hierarchyConfig.elementClassSpecification}
+                    WHERE
+                      Parent.Id IS NULL
+                      AND (
+                        ${whereClauses.join(" OR ")}
+                      )
+
+                    UNION ALL
+
+                    SELECT c.ECInstanceId, p.modelId, p.categoryId
+                    FROM ${this._hierarchyConfig.elementClassSpecification} c
+                    JOIN CategoryElements p ON c.Parent.Id = p.id
+                  )
+                `,
+              ],
+              ecsql: `
+                SELECT modelId, categoryId, COUNT(id) elementsCount
+                FROM CategoryElements
+                GROUP BY modelId, categoryId
+              `,
+            },
+            { rowFormat: "ECSqlPropertyNames", limit: "unbounded" },
+          );
+
+          const result = new Array<{ modelId: number; categoryId: number; elementsCount: number }>();
+          for await (const row of reader) {
+            result.push({ modelId: row.modelId, categoryId: row.categoryId, elementsCount: row.elementsCount });
+          }
+          return result;
+        }),
+        mergeAll(),
+      ),
+    );
   }
 
   public async getCategoryElementsCount(modelId: Id64String, categoryId: Id64String): Promise<number> {