feat(customHomePage/hierarchy): add support of pagination for root level nodes (#14485)

Co-authored-by: Chris Collins <chriscollins3456@gmail.com>

datahub-web-react/src/app/homeV3/modules/hierarchyViewModule/childrenLoader/context/ChildrenLoaderContext.ts

@@ -1,13 +1,13 @@
 import React from 'react';
 
 import { ChildrenLoaderContextType } from '@app/homeV3/modules/hierarchyViewModule/childrenLoader/types';
-import { DEFAULT_NUMBER_OF_CHILDREN_TO_LOAD } from '@app/homeV3/modules/hierarchyViewModule/treeView/constants';
+import { DEFAULT_LOAD_BATCH_SIZE } from '@app/homeV3/modules/hierarchyViewModule/treeView/constants';
 
 const DEFAULT_CONTEXT_STATE: ChildrenLoaderContextType = {
     get: () => undefined,
     upsert: () => {},
     onLoad: () => {},
-    maxNumberOfChildrenToLoad: DEFAULT_NUMBER_OF_CHILDREN_TO_LOAD,
+    maxNumberOfChildrenToLoad: DEFAULT_LOAD_BATCH_SIZE,
 };
 
 const ChildrenLoaderContext = React.createContext<ChildrenLoaderContextType>(DEFAULT_CONTEXT_STATE);

datahub-web-react/src/app/homeV3/modules/hierarchyViewModule/childrenLoader/context/ChildrenLoaderProvider.tsx

@@ -2,7 +2,7 @@ import React, { useCallback, useState } from 'react';
 
 import ChildrenLoaderContext from '@app/homeV3/modules/hierarchyViewModule/childrenLoader/context/ChildrenLoaderContext';
 import { ChildrenLoaderMetadata, MetadataMap } from '@app/homeV3/modules/hierarchyViewModule/childrenLoader/types';
-import { DEFAULT_NUMBER_OF_CHILDREN_TO_LOAD } from '@app/homeV3/modules/hierarchyViewModule/treeView/constants';
+import { DEFAULT_LOAD_BATCH_SIZE } from '@app/homeV3/modules/hierarchyViewModule/treeView/constants';
 import { TreeNode } from '@app/homeV3/modules/hierarchyViewModule/treeView/types';
 
 interface Props {
@@ -13,7 +13,7 @@ interface Props {
 export function ChildrenLoaderProvider({
     children,
     onLoadFinished,
-    maxNumberOfChildrenToLoad = DEFAULT_NUMBER_OF_CHILDREN_TO_LOAD,
+    maxNumberOfChildrenToLoad = DEFAULT_LOAD_BATCH_SIZE,
 }: React.PropsWithChildren<Props>) {
     const [metadataMap, setMetadataMap] = useState<MetadataMap>({});
 

datahub-web-react/src/app/homeV3/modules/hierarchyViewModule/components/domains/DomainsSelectableTreeView.tsx

@@ -12,10 +12,15 @@ import useSelectableDomainTree from '@app/homeV3/modules/hierarchyViewModule/com
 import { useHierarchyFormContext } from '@app/homeV3/modules/hierarchyViewModule/components/form/HierarchyFormContext';
 import TreeView from '@app/homeV3/modules/hierarchyViewModule/treeView/TreeView';
 import { TreeNode } from '@app/homeV3/modules/hierarchyViewModule/treeView/types';
-import { getTopLevelSelectedValuesFromTree } from '@app/homeV3/modules/hierarchyViewModule/treeView/utils';
+import {
+    getOutOfTreeSelectedValues,
+    getTopLevelSelectedValuesFromTree,
+} from '@app/homeV3/modules/hierarchyViewModule/treeView/utils';
 
 const Wrapper = styled.div``;
 
+const LOAD_BATCH_SIZE = 25;
+
 export default function DomainsSelectableTreeView() {
     const form = Form.useFormInstance();
     const {
@@ -24,12 +29,23 @@ export default function DomainsSelectableTreeView() {
 
     const { parentValues, addParentValue, removeParentValue } = useParentValuesToLoadChildren();
 
-    const { tree, selectedValues, setSelectedValues, loading } = useSelectableDomainTree(initialSelectedValues);
+    const {
+        tree,
+        selectedValues,
+        setSelectedValues,
+        loading,
+        loadMoreRootNodes,
+        rootDomainsMoreLoading,
+        rootNodesTotal,
+    } = useSelectableDomainTree(initialSelectedValues, LOAD_BATCH_SIZE);
 
     const updateSelectedValues = useCallback(
         (newSelectedValues: string[]) => {
             const topLevelSelectedValues = getTopLevelSelectedValuesFromTree(newSelectedValues, tree.nodes);
-            form.setFieldValue('domainAssets', topLevelSelectedValues);
+            // add out of tree selected values as some root nodes could not be loaded yet
+            const outOfTreeSelectedValues = getOutOfTreeSelectedValues(newSelectedValues, tree.nodes);
+            const newProcessedSelectedValues = new Set([...topLevelSelectedValues, ...outOfTreeSelectedValues]);
+            form.setFieldValue('domainAssets', Array.from(newProcessedSelectedValues));
             setSelectedValues(newSelectedValues);
         },
         [form, setSelectedValues, tree],
@@ -57,7 +73,7 @@ export default function DomainsSelectableTreeView() {
 
     return (
         <Wrapper>
-            <ChildrenLoaderProvider onLoadFinished={onLoadFinished}>
+            <ChildrenLoaderProvider onLoadFinished={onLoadFinished} maxNumberOfChildrenToLoad={LOAD_BATCH_SIZE}>
                 <ChildrenLoader parentValues={parentValues} loadChildren={useChildrenDomainsLoader} />
 
                 <TreeView
@@ -70,6 +86,11 @@ export default function DomainsSelectableTreeView() {
                     updateSelectedValues={updateSelectedValues}
                     loadChildren={startLoadingOfChildren}
                     renderNodeLabel={(nodeProps) => <DomainSelectableTreeNodeRenderer {...nodeProps} />}
+                    loadRootNodes={loadMoreRootNodes}
+                    rootNodesLoading={rootDomainsMoreLoading}
+                    loadingTriggerType="infiniteScroll"
+                    rootNodesTotal={rootNodesTotal}
+                    loadBatchSize={LOAD_BATCH_SIZE}
                 />
             </ChildrenLoaderProvider>
         </Wrapper>

datahub-web-react/src/app/homeV3/modules/hierarchyViewModule/components/domains/hooks/__tests__/useSelectableDomainTree.test.ts

@@ -1,8 +1,10 @@
+import { ApolloQueryResult, NetworkStatus } from '@apollo/client';
 import { act } from '@testing-library/react';
 import { renderHook } from '@testing-library/react-hooks';
 import { beforeEach, describe, expect, it, vi } from 'vitest';
 
 import useInitialDomains from '@app/homeV3/modules/hierarchyViewModule/components/domains/hooks/useInitialDomains';
+import useLoadMoreRootDomains from '@app/homeV3/modules/hierarchyViewModule/components/domains/hooks/useLoadMoreRootDomains';
 import useRootDomains from '@app/homeV3/modules/hierarchyViewModule/components/domains/hooks/useRootDomains';
 import useSelectableDomainTree from '@app/homeV3/modules/hierarchyViewModule/components/domains/hooks/useSelectableDomainTree';
 import useTreeNodesFromFlatDomains from '@app/homeV3/modules/hierarchyViewModule/components/domains/hooks/useTreeNodesFromFlatDomains';
@@ -10,11 +12,13 @@ import useTreeNodesFromDomains from '@app/homeV3/modules/hierarchyViewModule/com
 import useTree from '@app/homeV3/modules/hierarchyViewModule/treeView/useTree';
 import { mergeTrees } from '@app/homeV3/modules/hierarchyViewModule/treeView/utils';
 
-import { EntityType } from '@types';
+import { GetSearchResultsForMultipleQuery } from '@graphql/search.generated';
+import { Domain, EntityType } from '@types';
 
 // Mock all the dependent hooks
 vi.mock('@app/homeV3/modules/hierarchyViewModule/components/domains/hooks/useInitialDomains');
 vi.mock('@app/homeV3/modules/hierarchyViewModule/components/domains/hooks/useRootDomains');
+vi.mock('@app/homeV3/modules/hierarchyViewModule/components/domains/hooks/useLoadMoreRootDomains');
 vi.mock('@app/homeV3/modules/hierarchyViewModule/components/domains/hooks/useTreeNodesFromFlatDomains');
 vi.mock('@app/homeV3/modules/hierarchyViewModule/components/domains/hooks/useTreeNodesFromListDomains');
 vi.mock('@app/homeV3/modules/hierarchyViewModule/treeView/useTree');
@@ -22,11 +26,18 @@ vi.mock('@app/homeV3/modules/hierarchyViewModule/treeView/utils');
 
 const mockUseInitialDomains = vi.mocked(useInitialDomains);
 const mockUseRootDomains = vi.mocked(useRootDomains);
+const mockUseLoadMoreRootDomains = vi.mocked(useLoadMoreRootDomains);
 const mockUseTreeNodesFromFlatDomains = vi.mocked(useTreeNodesFromFlatDomains);
 const mockUseTreeNodesFromDomains = vi.mocked(useTreeNodesFromDomains);
 const mockUseTree = vi.mocked(useTree);
 const mockMergeTrees = vi.mocked(mergeTrees);
 
+const mockRefetch = (): Promise<ApolloQueryResult<GetSearchResultsForMultipleQuery>> => {
+    return new Promise<ApolloQueryResult<GetSearchResultsForMultipleQuery>>((resolve) => {
+        resolve({ data: { searchAcrossEntities: undefined }, networkStatus: NetworkStatus.refetch, loading: false });
+    });
+};
+
 // Mock tree object with dynamic nodes
 const mockTreeMethods = {
     nodes: [] as any[],
@@ -75,7 +86,20 @@ describe('useSelectableDomainTree hook', () => {
 
         // Setup default mock implementations
         mockUseInitialDomains.mockReturnValue({ data: undefined, domains: [], loading: false });
-        mockUseRootDomains.mockReturnValue({ data: undefined, domains: [], total: 0, loading: false });
+        mockUseRootDomains.mockReturnValue({
+            data: undefined,
+            domains: [],
+            total: 0,
+            loading: false,
+            refetch: mockRefetch,
+        });
+        mockUseLoadMoreRootDomains.mockReturnValue({
+            loading: false,
+            loadMoreRootDomains: () =>
+                new Promise<Domain[]>((resolve) => {
+                    resolve([]);
+                }),
+        });
         mockUseTreeNodesFromFlatDomains.mockReturnValue(undefined);
         mockUseTreeNodesFromDomains.mockReturnValue(undefined);
         mockUseTree.mockReturnValue(mockTreeMethods);
@@ -115,7 +139,13 @@ describe('useSelectableDomainTree hook', () => {
 
     describe('tree initialization', () => {
         it('should not initialize tree when root domains are loading', () => {
-            mockUseRootDomains.mockReturnValue({ data: undefined, domains: [], total: 0, loading: true });
+            mockUseRootDomains.mockReturnValue({
+                data: undefined,
+                domains: [],
+                total: 0,
+                loading: true,
+                refetch: mockRefetch,
+            });
 
             renderHook(() => useSelectableDomainTree([]));
 
@@ -123,7 +153,13 @@ describe('useSelectableDomainTree hook', () => {
         });
 
         it('should not initialize tree when root tree nodes are undefined', () => {
-            mockUseRootDomains.mockReturnValue({ data: undefined, domains: [mockDomain1], total: 1, loading: false });
+            mockUseRootDomains.mockReturnValue({
+                data: undefined,
+                domains: [mockDomain1],
+                total: 1,
+                loading: false,
+                refetch: mockRefetch,
+            });
             mockUseTreeNodesFromDomains.mockReturnValue(undefined);
 
             renderHook(() => useSelectableDomainTree([]));
@@ -132,7 +168,13 @@ describe('useSelectableDomainTree hook', () => {
         });
 
         it('should not initialize tree when initial selected tree nodes are undefined', () => {
-            mockUseRootDomains.mockReturnValue({ data: undefined, domains: [mockDomain1], total: 1, loading: false });
+            mockUseRootDomains.mockReturnValue({
+                data: undefined,
+                domains: [mockDomain1],
+                total: 1,
+                loading: false,
+                refetch: mockRefetch,
+            });
             mockUseTreeNodesFromDomains.mockReturnValue([mockTreeNode1]);
             mockUseTreeNodesFromFlatDomains.mockReturnValue(undefined);
 
@@ -146,21 +188,33 @@ describe('useSelectableDomainTree hook', () => {
             const initialSelectedTreeNodes = [mockTreeNode2];
             const mergedTreeNodes = [mockTreeNode1, mockTreeNode2];
 
-            mockUseRootDomains.mockReturnValue({ data: undefined, domains: [mockDomain1], total: 1, loading: false });
+            mockUseRootDomains.mockReturnValue({
+                data: undefined,
+                domains: [mockDomain1],
+                total: 1,
+                loading: false,
+                refetch: mockRefetch,
+            });
             mockUseTreeNodesFromDomains.mockReturnValue(rootTreeNodes);
             mockUseTreeNodesFromFlatDomains.mockReturnValue(initialSelectedTreeNodes);
             mockMergeTrees.mockReturnValue(mergedTreeNodes);
 
             renderHook(() => useSelectableDomainTree([]));
 
-            expect(mockMergeTrees).toHaveBeenCalledWith(rootTreeNodes, initialSelectedTreeNodes);
+            expect(mockMergeTrees).toHaveBeenCalledWith(rootTreeNodes, []);
             expect(mockTreeMethods.replace).toHaveBeenCalledWith(mergedTreeNodes);
         });
     });
 
     describe('selectedValues management', () => {
         beforeEach(() => {
-            mockUseRootDomains.mockReturnValue({ data: undefined, domains: [mockDomain1], total: 1, loading: false });
+            mockUseRootDomains.mockReturnValue({
+                data: undefined,
+                domains: [mockDomain1],
+                total: 1,
+                loading: false,
+                refetch: mockRefetch,
+            });
             mockUseTreeNodesFromDomains.mockReturnValue([mockTreeNode1]);
             mockUseTreeNodesFromFlatDomains.mockReturnValue([mockTreeNode2]);
         });
@@ -185,7 +239,13 @@ describe('useSelectableDomainTree hook', () => {
         });
 
         it('should handle loading state correctly when root domains are loading', () => {
-            mockUseRootDomains.mockReturnValue({ data: undefined, domains: [], total: 0, loading: true });
+            mockUseRootDomains.mockReturnValue({
+                data: undefined,
+                domains: [],
+                total: 0,
+                loading: true,
+                refetch: mockRefetch,
+            });
 
             const { result } = renderHook(() => useSelectableDomainTree([]));
 
@@ -193,7 +253,13 @@ describe('useSelectableDomainTree hook', () => {
         });
 
         it('should handle ready state when all data is loaded', () => {
-            mockUseRootDomains.mockReturnValue({ data: undefined, domains: [mockDomain1], total: 1, loading: false });
+            mockUseRootDomains.mockReturnValue({
+                data: undefined,
+                domains: [mockDomain1],
+                total: 1,
+                loading: false,
+                refetch: mockRefetch,
+            });
             mockUseTreeNodesFromDomains.mockReturnValue([mockTreeNode1]);
 
             const { result } = renderHook(() => useSelectableDomainTree([]));
@@ -204,7 +270,13 @@ describe('useSelectableDomainTree hook', () => {
 
     describe('tree method delegation', () => {
         beforeEach(() => {
-            mockUseRootDomains.mockReturnValue({ data: undefined, domains: [mockDomain1], total: 1, loading: false });
+            mockUseRootDomains.mockReturnValue({
+                data: undefined,
+                domains: [mockDomain1],
+                total: 1,
+                loading: false,
+                refetch: mockRefetch,
+            });
             mockUseTreeNodesFromDomains.mockReturnValue([mockTreeNode1]);
             mockUseTreeNodesFromFlatDomains.mockReturnValue([mockTreeNode2]);
         });
@@ -254,7 +326,13 @@ describe('useSelectableDomainTree hook', () => {
         it('should use domains from root domains hook', () => {
             const rootDomains = [mockDomain1];
 
-            mockUseRootDomains.mockReturnValue({ data: undefined, domains: rootDomains, total: 1, loading: false });
+            mockUseRootDomains.mockReturnValue({
+                data: undefined,
+                domains: rootDomains,
+                total: 1,
+                loading: false,
+                refetch: mockRefetch,
+            });
 
             renderHook(() => useSelectableDomainTree([]));
 
@@ -262,14 +340,20 @@ describe('useSelectableDomainTree hook', () => {
         });
 
         it('should handle mixed states correctly', () => {
-            mockUseRootDomains.mockReturnValue({ data: undefined, domains: [mockDomain1], total: 1, loading: false });
+            mockUseRootDomains.mockReturnValue({
+                data: undefined,
+                domains: [mockDomain1],
+                total: 1,
+                loading: false,
+                refetch: mockRefetch,
+            });
             mockUseTreeNodesFromDomains.mockReturnValue([mockTreeNode1]);
             mockUseTreeNodesFromFlatDomains.mockReturnValue([mockTreeNode2]);
 
             const { result } = renderHook(() => useSelectableDomainTree([]));
 
             expect(result.current.loading).toBe(false);
-            expect(mockMergeTrees).toHaveBeenCalledWith([mockTreeNode1], [mockTreeNode2]);
+            expect(mockMergeTrees).toHaveBeenCalledWith([mockTreeNode1], []);
         });
 
         it('should handle empty initial domains', () => {
@@ -281,7 +365,13 @@ describe('useSelectableDomainTree hook', () => {
         });
 
         it('should handle empty root domains', () => {
-            mockUseRootDomains.mockReturnValue({ data: undefined, domains: [mockDomain1], total: 1, loading: false });
+            mockUseRootDomains.mockReturnValue({
+                data: undefined,
+                domains: [mockDomain1],
+                total: 1,
+                loading: false,
+                refetch: mockRefetch,
+            });
             mockUseTreeNodesFromDomains.mockReturnValue([mockTreeNode1]);
             mockUseTreeNodesFromFlatDomains.mockReturnValue([mockTreeNode2]);
 

datahub-web-react/src/app/homeV3/modules/hierarchyViewModule/components/domains/hooks/useDomains.ts

@@ -1,14 +1,28 @@
 import { useMemo } from 'react';
 
-import { useListDomainsQuery } from '@graphql/domain.generated';
+import { isDomain } from '@app/entityV2/domain/utils';
+import { ENTITY_NAME_FIELD } from '@app/searchV2/context/constants';
 
-export default function useDomains(parentDomainUrns: string | undefined, start: number, count: number, skip: boolean) {
+import { useGetSearchResultsForMultipleQuery } from '@graphql/search.generated';
-    const { data, loading } = useListDomainsQuery({
+import { EntityType, FilterOperator, SortOrder } from '@types';
+
+export default function useDomains(parentDomainUrn: string | undefined, start: number, count: number, skip: boolean) {
+    const { data, loading, refetch } = useGetSearchResultsForMultipleQuery({
         variables: {
             input: {
                 start,
                 count,
-                parentDomain: parentDomainUrns,
+                query: '*',
+                types: [EntityType.Domain],
+                orFilters: parentDomainUrn
+                    ? [{ and: [{ field: 'parentDomain', values: [parentDomainUrn] }] }]
+                    : [{ and: [{ field: 'parentDomain', condition: FilterOperator.Exists, negated: true }] }],
+                sortInput: {
+                    sortCriteria: [{ field: ENTITY_NAME_FIELD, sortOrder: SortOrder.Ascending }],
+                },
+                searchFlags: {
+                    skipCache: true,
+                },
             },
         },
         skip,
@@ -16,10 +30,10 @@ export default function useDomains(parentDomainUrns: string | undefined, start:
 
     const domains = useMemo(() => {
         if (skip) return [];
-        return data?.listDomains?.domains;
+        return data?.searchAcrossEntities?.searchResults.map((result) => result.entity).filter(isDomain);
     }, [data, skip]);
 
-    const total = useMemo(() => data?.listDomains?.total, [data]);
+    const total = useMemo(() => data?.searchAcrossEntities?.total, [data]);
 
-    return { data, domains, total, loading };
+    return { data, domains, total, loading, refetch };
 }

datahub-web-react/src/app/homeV3/modules/hierarchyViewModule/components/domains/hooks/useLoadMoreRootDomains.ts

@@ -0,0 +1,66 @@
+import { useCallback, useState } from 'react';
+
+import { isDomain } from '@app/entityV2/domain/utils';
+import { ENTITY_NAME_FIELD } from '@app/searchV2/context/constants';
+
+import { useGetSearchResultsForMultipleQuery } from '@graphql/search.generated';
+import { EntityType, FilterOperator, SortOrder } from '@types';
+
+export default function useLoadMoreRootDomains() {
+    const [loading, setLoading] = useState<boolean>(false);
+
+    const { refetch } = useGetSearchResultsForMultipleQuery({
+        variables: {
+            input: {
+                query: '*',
+                types: [EntityType.Domain],
+                start: 0,
+                count: 0,
+                sortInput: {
+                    sortCriteria: [{ field: ENTITY_NAME_FIELD, sortOrder: SortOrder.Ascending }],
+                },
+            },
+        },
+        skip: true,
+    });
+
+    const loadMoreRootDomains = useCallback(
+        async (start: number, pageSize: number) => {
+            setLoading(true);
+
+            try {
+                const response = await refetch({
+                    input: {
+                        start,
+                        types: [EntityType.Domain],
+                        query: '*',
+                        count: pageSize,
+                        orFilters: [
+                            { and: [{ field: 'parentDomain', condition: FilterOperator.Exists, negated: true }] },
+                        ],
+                        sortInput: {
+                            sortCriteria: [{ field: ENTITY_NAME_FIELD, sortOrder: SortOrder.Ascending }],
+                        },
+                    },
+                });
+
+                return (
+                    response.data.searchAcrossEntities?.searchResults
+                        ?.map((searchResult) => searchResult.entity)
+                        .filter(isDomain) ?? []
+                );
+            } catch (e) {
+                console.error('Something went wrong during fetching root domains', e);
+                return [];
+            } finally {
+                setLoading(false);
+            }
+        },
+        [refetch],
+    );
+
+    return {
+        loadMoreRootDomains,
+        loading,
+    };
+}

datahub-web-react/src/app/homeV3/modules/hierarchyViewModule/components/domains/hooks/useRootDomains.ts

@@ -1,7 +1,5 @@
 import useDomains from '@app/homeV3/modules/hierarchyViewModule/components/domains/hooks/useDomains';
 
-const MAX_ROOT_DOMAINS = 1000;
+export default function useRootDomains(count: number) {
-
+    return useDomains(undefined, 0, count, false);
-export default function useRootDomains() {
-    return useDomains(undefined, 0, MAX_ROOT_DOMAINS, false);
 }

datahub-web-react/src/app/homeV3/modules/hierarchyViewModule/components/domains/hooks/useSelectableDomainTree.ts

@@ -1,25 +1,64 @@
-import { useEffect, useState } from 'react';
+import { useCallback, useEffect, useState } from 'react';
 
 import useDomainTreeNodesSorter from '@app/homeV3/modules/hierarchyViewModule/components/domains/hooks/useDomainTreeNodesSorter';
 import useInitialDomains from '@app/homeV3/modules/hierarchyViewModule/components/domains/hooks/useInitialDomains';
+import useLoadMoreRootDomains from '@app/homeV3/modules/hierarchyViewModule/components/domains/hooks/useLoadMoreRootDomains';
 import useRootDomains from '@app/homeV3/modules/hierarchyViewModule/components/domains/hooks/useRootDomains';
 import useTreeNodesFromFlatDomains from '@app/homeV3/modules/hierarchyViewModule/components/domains/hooks/useTreeNodesFromFlatDomains';
 import useTreeNodesFromDomains from '@app/homeV3/modules/hierarchyViewModule/components/domains/hooks/useTreeNodesFromListDomains';
+import { convertDomainToTreeNode } from '@app/homeV3/modules/hierarchyViewModule/components/domains/utils';
+import { DEFAULT_LOAD_BATCH_SIZE } from '@app/homeV3/modules/hierarchyViewModule/treeView/constants';
+import { TreeNode } from '@app/homeV3/modules/hierarchyViewModule/treeView/types';
 import useTree from '@app/homeV3/modules/hierarchyViewModule/treeView/useTree';
 import { mergeTrees } from '@app/homeV3/modules/hierarchyViewModule/treeView/utils';
 
-export default function useSelectableDomainTree(initialSelectedDomainUrns: string[] | undefined) {
+export default function useSelectableDomainTree(
+    initialSelectedDomainUrns: string[] | undefined,
+    loadBatchSize = DEFAULT_LOAD_BATCH_SIZE,
+) {
     const [isInitialized, setIsInitialized] = useState<boolean>(false);
     const [selectedValues, setSelectedValues] = useState<string[]>(initialSelectedDomainUrns ?? []);
     const nodesSorter = useDomainTreeNodesSorter();
     const tree = useTree(undefined, nodesSorter);
 
+    // FYI: We get and extract parents from initially selected domains
+    // to have possibility to detect if some node has any selected nested nodes
     const { domains: initialDomains } = useInitialDomains(initialSelectedDomainUrns ?? []);
     const initialSelectedTreeNodes = useTreeNodesFromFlatDomains(initialDomains);
 
-    const { domains: rootDomains, loading: rootDomainsLoading } = useRootDomains();
+    const {
+        domains: rootDomains,
+        loading: rootDomainsLoading,
+        total: rootDomainsTotal,
+    } = useRootDomains(loadBatchSize);
     const rootTreeNodes = useTreeNodesFromDomains(rootDomains, false);
 
+    const { loadMoreRootDomains, loading: rootDomainsMoreLoading } = useLoadMoreRootDomains();
+
+    const preprocessRootNodes = useCallback(
+        (rootNodes: TreeNode[]) => {
+            if (!initialSelectedTreeNodes) return rootNodes;
+
+            // Merge initial selected nodes with the same root nodes
+            const initialSelectedTreeNodesToMerge = initialSelectedTreeNodes.filter((initialSelectedTreeNode) =>
+                rootNodes.some((rootNode) => rootNode.value === initialSelectedTreeNode.value),
+            );
+            return mergeTrees(rootNodes, initialSelectedTreeNodesToMerge);
+        },
+        [initialSelectedTreeNodes],
+    );
+
+    const loadMoreRootNodes = useCallback(async () => {
+        const hasMoreRootNodes = (rootDomainsTotal ?? 0) > tree.nodes.length;
+        if (!rootDomainsMoreLoading && hasMoreRootNodes) {
+            const domains = await loadMoreRootDomains(tree.nodes.length, loadBatchSize);
+            if (domains) {
+                const treeNodes = domains.map((domain) => convertDomainToTreeNode(domain));
+                tree.merge(preprocessRootNodes(treeNodes));
+            }
+        }
+    }, [tree, loadMoreRootDomains, rootDomainsTotal, rootDomainsMoreLoading, preprocessRootNodes, loadBatchSize]);
+
     useEffect(() => {
         if (
             !rootDomainsLoading &&
@@ -27,15 +66,18 @@ export default function useSelectableDomainTree(initialSelectedDomainUrns: strin
             rootTreeNodes !== undefined &&
             initialSelectedTreeNodes !== undefined
         ) {
-            tree.replace(mergeTrees(rootTreeNodes, initialSelectedTreeNodes));
+            tree.replace(preprocessRootNodes(rootTreeNodes));
             setIsInitialized(true);
         }
-    }, [tree, isInitialized, rootTreeNodes, initialSelectedTreeNodes, rootDomainsLoading]);
+    }, [tree, isInitialized, rootTreeNodes, initialSelectedTreeNodes, rootDomainsLoading, preprocessRootNodes]);
 
     return {
         tree,
         loading: !isInitialized,
         selectedValues,
         setSelectedValues,
+        loadMoreRootNodes,
+        rootDomainsMoreLoading,
+        rootNodesTotal: rootDomainsTotal,
     };
 }

datahub-web-react/src/app/homeV3/modules/hierarchyViewModule/treeView/TreeNodeRenderer.tsx

@@ -1,27 +1,20 @@
-import { Button, Checkbox, Loader } from '@components';
+import { Checkbox } from '@components';
 import React, { useEffect, useMemo } from 'react';
 import styled from 'styled-components';
 
 import ChildrenLoader from '@app/homeV3/modules/hierarchyViewModule/treeView/ChildrenLoader';
 import DepthMargin from '@app/homeV3/modules/hierarchyViewModule/treeView/DepthMargin';
 import ExpandToggler from '@app/homeV3/modules/hierarchyViewModule/treeView/ExpandToggler';
+import Row from '@app/homeV3/modules/hierarchyViewModule/treeView/components/Row';
+import TreeNodesViewLoader from '@app/homeV3/modules/hierarchyViewModule/treeView/components/TreeNodesViewLoader';
+import NodesLoaderWrapper from '@app/homeV3/modules/hierarchyViewModule/treeView/components/itemsLoaderWrapper/NodesLoaderWrapper';
 import useTreeViewContext from '@app/homeV3/modules/hierarchyViewModule/treeView/context/useTreeViewContext';
 import { TreeNode } from '@app/homeV3/modules/hierarchyViewModule/treeView/types';
 
-const Row = styled.div`
-    display: flex;
-    flex-direction: row;
-    align-items: center;
-    gap: 8px;
-    width: 100%;
-`;
-
 const SpaceFiller = styled.div`
     flex-grow: 1;
 `;
 
-const LoaderWrapper = styled.div``;
-
 interface Props {
     node: TreeNode;
     depth: number;
@@ -29,8 +22,11 @@ interface Props {
 
 export default function TreeNodeRenderer({ node, depth }: Props) {
     const {
+        loadingTriggerType,
         getHasParentNode,
         getIsRootNode,
+        getChildrenTotal,
+        getChildrenLength,
         renderNodeLabel,
         getIsExpandable,
         getIsExpanded,
@@ -45,9 +41,8 @@ export default function TreeNodeRenderer({ node, depth }: Props) {
         enableIntermediateSelectState,
         toggleSelected,
         getIsChildrenLoading,
-        getNumberOfNotLoadedChildren,
         loadChildren,
-        numberOfChildrenToLoad,
+        loadBatchSize: numberOfChildrenToLoad,
     } = useTreeViewContext();
 
     const isExpandable = useMemo(() => getIsExpandable(node), [node, getIsExpandable]);
@@ -73,13 +68,8 @@ export default function TreeNodeRenderer({ node, depth }: Props) {
 
     const isRootNode = useMemo(() => getIsRootNode(node), [node, getIsRootNode]);
 
-    const numberOfNotLoadedChildren = useMemo(
+    const childrenTotal = useMemo(() => getChildrenTotal(node), [node, getChildrenTotal]);
-        () => getNumberOfNotLoadedChildren(node),
+    const childrenLength = useMemo(() => getChildrenLength(node), [node, getChildrenLength]);
-        [node, getNumberOfNotLoadedChildren],
-    );
-
-    const maxNumberOfChildrenToLoad = Math.min(numberOfNotLoadedChildren, numberOfChildrenToLoad);
-    const shouldShowLoadMoreButton = isExpanded && !isChildrenLoading && numberOfNotLoadedChildren > 0;
 
     // Automatically select the current node if parent is selected if explicitlySelectChildren is not enabled
     // FYI: required to get loaded children selected if parent is selected
@@ -90,7 +80,16 @@ export default function TreeNodeRenderer({ node, depth }: Props) {
     }, [explicitlySelectChildren, hasParentNode, isParentSelected, select, node, isSelected, isSelectable]);
 
     return (
-        <>
+        <NodesLoaderWrapper
+            trigger={loadingTriggerType}
+            onLoad={() => loadChildren(node)}
+            total={childrenTotal}
+            current={childrenLength}
+            pageSize={numberOfChildrenToLoad}
+            depth={depth + 1}
+            enabled={isExpanded}
+            loading={isChildrenLoading}
+        >
             <Row>
                 <DepthMargin depth={depth} />
 
@@ -127,27 +126,8 @@ export default function TreeNodeRenderer({ node, depth }: Props) {
             {/* Run loading on expand */}
             {isExpanded && <ChildrenLoader node={node} />}
 
-            {/* Show more button */}
-            {shouldShowLoadMoreButton && (
-                <Row>
-                    <DepthMargin depth={depth + 1} />
-                    <ExpandToggler expandable={false} />
-                    <Button onClick={() => loadChildren(node)} variant="link" color="gray">
-                        Show {maxNumberOfChildrenToLoad} more
-                    </Button>
-                </Row>
-            )}
-
             {/* Loading indicator */}
-            {isExpanded && isChildrenLoading && (
+            {isExpanded && isChildrenLoading && <TreeNodesViewLoader depth={depth + 1} />}
-                <Row>
+        </NodesLoaderWrapper>
-                    <DepthMargin depth={depth + 1} />
-                    <ExpandToggler expandable={false} />
-                    <LoaderWrapper>
-                        <Loader size="xs" />
-                    </LoaderWrapper>
-                </Row>
-            )}
-        </>
     );
 }

datahub-web-react/src/app/homeV3/modules/hierarchyViewModule/treeView/TreeNodesRenderer.tsx

@@ -1,8 +1,11 @@
-import React from 'react';
+import React, { useCallback } from 'react';
 import styled from 'styled-components';
 
 import TreeNodeRenderer from '@app/homeV3/modules/hierarchyViewModule/treeView/TreeNodeRenderer';
+import TreeNodesViewLoader from '@app/homeV3/modules/hierarchyViewModule/treeView/components/TreeNodesViewLoader';
+import NodesLoaderWrapper from '@app/homeV3/modules/hierarchyViewModule/treeView/components/itemsLoaderWrapper/NodesLoaderWrapper';
 import useTreeViewContext from '@app/homeV3/modules/hierarchyViewModule/treeView/context/useTreeViewContext';
+import { TreeNode } from '@app/homeV3/modules/hierarchyViewModule/treeView/types';
 
 const Wrapper = styled.div`
     display: flex;
@@ -17,14 +20,38 @@ const InlineBlockWrapper = styled.div<{ $hasExpanded: boolean }>`
 `;
 
 export default function TreeNodesRenderer() {
-    const { nodes, hasAnyExpanded } = useTreeViewContext();
+    const {
+        nodes,
+        hasAnyExpanded,
+        loadBatchSize: numberOfChildrenToLoad,
+        loadRootNodes,
+        rootNodesLength,
+        rootNodesTotal,
+        rootNodesLoading,
+        loadingTriggerType,
+    } = useTreeViewContext();
+
+    const renderNode = useCallback((node: TreeNode) => <TreeNodeRenderer node={node} depth={0} key={node.value} />, []);
 
     return (
         <InlineBlockWrapper $hasExpanded={hasAnyExpanded}>
             <Wrapper>
-                {nodes.map((node) => (
+                {loadRootNodes ? (
-                    <TreeNodeRenderer node={node} depth={0} key={node.value} />
+                    <NodesLoaderWrapper
-                ))}
+                        trigger={loadingTriggerType}
+                        total={rootNodesTotal}
+                        current={rootNodesLength}
+                        enabled={!rootNodesLoading}
+                        depth={0}
+                        onLoad={loadRootNodes}
+                        pageSize={numberOfChildrenToLoad}
+                    >
+                        {nodes.map(renderNode)}
+                    </NodesLoaderWrapper>
+                ) : (
+                    nodes.map(renderNode)
+                )}
+                {rootNodesLoading && <TreeNodesViewLoader depth={0} />}
             </Wrapper>
         </InlineBlockWrapper>
     );

datahub-web-react/src/app/homeV3/modules/hierarchyViewModule/treeView/__tests__/utils.test.ts

@@ -6,6 +6,7 @@ import {
     flattenTreeNodes,
     getAllParentValues,
     getAllValues,
+    getOutOfTreeSelectedValues,
     getTopLevelSelectedValuesFromTree,
     getValueToTreeNodeMapping,
     mergeTrees,
@@ -418,4 +419,56 @@ describe('treeView utils', () => {
             expect(result).toEqual(['parent1', 'child2']);
         });
     });
+
+    describe('getOutOfTreeSelectedValues', () => {
+        const child1 = createTreeNode('child1');
+        const child2 = createTreeNode('child2');
+        const parent1 = createTreeNode('parent1', 'parent1', [child1]);
+        const parent2 = createTreeNode('parent2', 'parent2', [child2]);
+        const tree = [parent1, parent2];
+
+        it('should return only values out of tree', () => {
+            const selectedValues = ['parent1', 'child2', 'parent3'];
+
+            const result = getOutOfTreeSelectedValues(selectedValues, tree);
+
+            expect(result).toEqual(['parent3']);
+        });
+
+        it('should handle empty tree', () => {
+            const selectedValues = ['parent1', 'child2', 'parent3'];
+
+            const result = getOutOfTreeSelectedValues(selectedValues, []);
+
+            expect(result).toEqual(selectedValues);
+        });
+
+        it('should handle empty selected values', () => {
+            const result = getOutOfTreeSelectedValues([], tree);
+
+            expect(result).toEqual([]);
+        });
+
+        it('should handle empty selected values and tree', () => {
+            const result = getOutOfTreeSelectedValues([], []);
+
+            expect(result).toEqual([]);
+        });
+
+        it('should return empty array when all values included in tree', () => {
+            const selectedValues = ['parent1', 'child2'];
+
+            const result = getOutOfTreeSelectedValues(selectedValues, tree);
+
+            expect(result).toEqual([]);
+        });
+
+        it('should return empty array when all values not included in tree', () => {
+            const selectedValues = ['parent3'];
+
+            const result = getOutOfTreeSelectedValues(selectedValues, tree);
+
+            expect(result).toEqual(['parent3']);
+        });
+    });
 });

datahub-web-react/src/app/homeV3/modules/hierarchyViewModule/treeView/components/Row.tsx

@@ -0,0 +1,11 @@
+import styled from 'styled-components';
+
+const Row = styled.div`
+    display: flex;
+    flex-direction: row;
+    align-items: center;
+    gap: 8px;
+    width: 100%;
+`;
+
+export default Row;

datahub-web-react/src/app/homeV3/modules/hierarchyViewModule/treeView/components/TreeNodesViewLoader.tsx

@@ -0,0 +1,25 @@
+import { Loader } from '@components';
+import React from 'react';
+import styled from 'styled-components';
+
+import DepthMargin from '@app/homeV3/modules/hierarchyViewModule/treeView/DepthMargin';
+import ExpandToggler from '@app/homeV3/modules/hierarchyViewModule/treeView/ExpandToggler';
+import Row from '@app/homeV3/modules/hierarchyViewModule/treeView/components/Row';
+
+const LoaderWrapper = styled.div``;
+
+interface Props {
+    depth: number;
+}
+
+export default function TreeNodesViewLoader({ depth }: Props) {
+    return (
+        <Row>
+            <DepthMargin depth={depth + 1} />
+            <ExpandToggler expandable={false} />
+            <LoaderWrapper>
+                <Loader size="xs" />
+            </LoaderWrapper>
+        </Row>
+    );
+}

datahub-web-react/src/app/homeV3/modules/hierarchyViewModule/treeView/components/itemsLoaderWrapper/NodesLoaderByButtonWrapper.tsx

@@ -0,0 +1,36 @@
+import { Button } from '@components';
+import React, { useMemo } from 'react';
+
+import DepthMargin from '@app/homeV3/modules/hierarchyViewModule/treeView/DepthMargin';
+import ExpandToggler from '@app/homeV3/modules/hierarchyViewModule/treeView/ExpandToggler';
+import Row from '@app/homeV3/modules/hierarchyViewModule/treeView/components/Row';
+import { NodesLoaderWrapperProps } from '@app/homeV3/modules/hierarchyViewModule/treeView/components/itemsLoaderWrapper/types';
+
+export default function NodesLoaderByButtonWrapper({
+    children,
+    total,
+    current,
+    pageSize,
+    depth,
+    enabled,
+    loading,
+    onLoad,
+}: React.PropsWithChildren<NodesLoaderWrapperProps>) {
+    const loadMoreNumber = useMemo(() => Math.min(total - current, pageSize), [total, current, pageSize]);
+
+    return (
+        <>
+            {children}
+
+            {enabled && !loading && loadMoreNumber > 0 && (
+                <Row>
+                    <DepthMargin depth={depth + 1} />
+                    <ExpandToggler expandable={false} />
+                    <Button onClick={onLoad} variant="link" color="gray">
+                        Show {loadMoreNumber} more
+                    </Button>
+                </Row>
+            )}
+        </>
+    );
+}

datahub-web-react/src/app/homeV3/modules/hierarchyViewModule/treeView/components/itemsLoaderWrapper/NodesLoaderInfiniteScrollWrapper.tsx

@@ -0,0 +1,45 @@
+import React, { useEffect, useMemo, useState } from 'react';
+import { useInView } from 'react-intersection-observer';
+import styled from 'styled-components';
+
+import { NodesLoaderWrapperProps } from '@app/homeV3/modules/hierarchyViewModule/treeView/components/itemsLoaderWrapper/types';
+
+export const ObserverContainer = styled.div`
+    height: 1px;
+    margin-top: 1px;
+`;
+
+export default function NodesLoaderInfiniteScrollWrapper({
+    children,
+    total,
+    current,
+    enabled,
+    loading,
+    onLoad,
+}: React.PropsWithChildren<NodesLoaderWrapperProps>) {
+    const [scrollRef, inView] = useInView();
+    // FYI: additional flag to prevent `onLoad` calling
+    // when infinite scroll triggered but real loading hasn't started yet
+    const [shouldPreventLoading, setShouldPreventLoading] = useState<boolean>(false);
+    // reset flag when observer container is out of view (it's hidden while real loading because `loading` is true)
+    useEffect(() => {
+        if (!inView) setShouldPreventLoading(false);
+    }, [inView]);
+
+    const hasMoreNodes = useMemo(() => total - current > 0, [total, current]);
+
+    useEffect(() => {
+        if (enabled && inView && hasMoreNodes && !shouldPreventLoading) {
+            setShouldPreventLoading(true);
+            onLoad();
+        }
+    }, [shouldPreventLoading, inView, enabled, hasMoreNodes, onLoad]);
+
+    return (
+        <>
+            {children}
+
+            {enabled && !loading && hasMoreNodes && <ObserverContainer ref={scrollRef} />}
+        </>
+    );
+}

datahub-web-react/src/app/homeV3/modules/hierarchyViewModule/treeView/components/itemsLoaderWrapper/NodesLoaderWrapper.tsx

@@ -0,0 +1,26 @@
+import React, { useMemo } from 'react';
+
+import NodesLoaderByButtonWrapper from '@app/homeV3/modules/hierarchyViewModule/treeView/components/itemsLoaderWrapper/NodesLoaderByButtonWrapper';
+import NodesLoaderInfiniteScrollWrapper from '@app/homeV3/modules/hierarchyViewModule/treeView/components/itemsLoaderWrapper/NodesLoaderInfiniteScrollWrapper';
+import { NodesLoaderWrapperProps } from '@app/homeV3/modules/hierarchyViewModule/treeView/components/itemsLoaderWrapper/types';
+import { LoadingTriggerType } from '@app/homeV3/modules/hierarchyViewModule/treeView/types';
+
+interface Props extends NodesLoaderWrapperProps {
+    trigger?: LoadingTriggerType;
+    loading?: boolean;
+}
+
+export default function NodesLoaderWrapper({ children, trigger, ...props }: React.PropsWithChildren<Props>) {
+    const NodesLoaderWrapperComponent = useMemo(() => {
+        switch (trigger) {
+            case 'button':
+                return NodesLoaderByButtonWrapper;
+            case 'infiniteScroll':
+                return NodesLoaderInfiniteScrollWrapper;
+            default:
+                return NodesLoaderByButtonWrapper;
+        }
+    }, [trigger]);
+
+    return <NodesLoaderWrapperComponent {...props}>{children}</NodesLoaderWrapperComponent>;
+}

datahub-web-react/src/app/homeV3/modules/hierarchyViewModule/treeView/components/itemsLoaderWrapper/types.ts

@@ -0,0 +1,10 @@
+export interface NodesLoaderWrapperProps {
+    total: number;
+    // how many items already loaded
+    current: number;
+    pageSize: number;
+    depth: number;
+    enabled?: boolean;
+    loading?: boolean;
+    onLoad: () => void;
+}

datahub-web-react/src/app/homeV3/modules/hierarchyViewModule/treeView/constants.ts

@@ -1 +1 @@
-export const DEFAULT_NUMBER_OF_CHILDREN_TO_LOAD = 5;
+export const DEFAULT_LOAD_BATCH_SIZE = 5;

datahub-web-react/src/app/homeV3/modules/hierarchyViewModule/treeView/context/TreeViewContext.tsx

@@ -1,6 +1,6 @@
 import React from 'react';
 
-import { DEFAULT_NUMBER_OF_CHILDREN_TO_LOAD } from '@app/homeV3/modules/hierarchyViewModule/treeView/constants';
+import { DEFAULT_LOAD_BATCH_SIZE } from '@app/homeV3/modules/hierarchyViewModule/treeView/constants';
 import { TreeViewContextType } from '@app/homeV3/modules/hierarchyViewModule/treeView/types';
 
 const DEFAULT_TREE_VIEW_CONTEXT: TreeViewContextType = {
@@ -8,6 +8,10 @@ const DEFAULT_TREE_VIEW_CONTEXT: TreeViewContextType = {
 
     getHasParentNode: () => false,
     getIsRootNode: () => false,
+    rootNodesLength: 0,
+    rootNodesTotal: 0,
+    getChildrenLength: () => 0,
+    getChildrenTotal: () => 0,
 
     getIsExpandable: () => false,
     getIsExpanded: () => false,
@@ -26,9 +30,8 @@ const DEFAULT_TREE_VIEW_CONTEXT: TreeViewContextType = {
     toggleSelected: () => {},
 
     getIsChildrenLoading: () => false,
-    getNumberOfNotLoadedChildren: () => 0,
     loadChildren: () => {},
-    numberOfChildrenToLoad: DEFAULT_NUMBER_OF_CHILDREN_TO_LOAD,
+    loadBatchSize: DEFAULT_LOAD_BATCH_SIZE,
 
     explicitlySelectChildren: false,
     explicitlyUnselectChildren: false,

datahub-web-react/src/app/homeV3/modules/hierarchyViewModule/treeView/context/TreeViewContextProvider.tsx

@@ -1,6 +1,6 @@
 import React, { useCallback, useEffect, useMemo, useState } from 'react';
 
-import { DEFAULT_NUMBER_OF_CHILDREN_TO_LOAD } from '@app/homeV3/modules/hierarchyViewModule/treeView/constants';
+import { DEFAULT_LOAD_BATCH_SIZE } from '@app/homeV3/modules/hierarchyViewModule/treeView/constants';
 import TreeViewContext from '@app/homeV3/modules/hierarchyViewModule/treeView/context/TreeViewContext';
 import { TreeNode, TreeViewContextProviderProps } from '@app/homeV3/modules/hierarchyViewModule/treeView/types';
 import {
@@ -22,6 +22,10 @@ export default function TreeViewContextProvider({
     selectable,
     updateSelectedValues,
     expandParentNodesOfInitialSelectedValues,
+    loadingTriggerType = 'button',
+    rootNodesTotal: rootNodesTotalProperty,
+    loadRootNodes,
+    rootNodesLoading,
     loadChildren: loadAsyncChildren,
     renderNodeLabel,
     explicitlySelectChildren,
@@ -29,7 +33,7 @@ export default function TreeViewContextProvider({
     explicitlySelectParent,
     explicitlyUnselectParent,
     enableIntermediateSelectState,
-    numberOfChildrenToLoad = DEFAULT_NUMBER_OF_CHILDREN_TO_LOAD,
+    loadBatchSize = DEFAULT_LOAD_BATCH_SIZE,
 }: React.PropsWithChildren<TreeViewContextProviderProps>) {
     const [internalExpandedValues, setInternalExpandedValues] = useState<string[]>(expandedValues ?? []);
     const [isExpandedValuesInitialized, setIsExpandedValuesInitialized] = useState<boolean>(false);
@@ -92,6 +96,18 @@ export default function TreeViewContextProvider({
         [getHasParentNode, getRootNodes, valueToTreeNodeMap],
     );
 
+    const rootNodesLength = useMemo(() => nodes.length, [nodes]);
+    const rootNodesTotal = useMemo(
+        () => rootNodesTotalProperty ?? rootNodesLength,
+        [rootNodesTotalProperty, rootNodesLength],
+    );
+
+    const getChildrenLength = useCallback((node: TreeNode) => node.children?.length ?? 0, []);
+    const getChildrenTotal = useCallback(
+        (node: TreeNode) => node.totalChildren ?? getChildrenLength(node),
+        [getChildrenLength],
+    );
+
     // Expanding
 
     const getIsExpandable = useCallback((node: TreeNode) => !!node.children?.length || !!node.hasAsyncChildren, []);
@@ -258,17 +274,18 @@ export default function TreeViewContextProvider({
     // Loading of children
     const getIsChildrenLoading = useCallback((node: TreeNode) => !!node.isChildrenLoading, []);
 
-    const getNumberOfNotLoadedChildren = useCallback(
-        (node: TreeNode) => (node.totalChildren ? node.totalChildren - (node.children?.length ?? 0) : 0),
-        [],
-    );
-
     return (
         <TreeViewContext.Provider
             value={{
                 nodes: preprocessedNodes,
+
+                // Node utils
                 getHasParentNode,
                 getIsRootNode,
+                rootNodesLength,
+                rootNodesTotal,
+                getChildrenLength,
+                getChildrenTotal,
 
                 // Expanding
                 getIsExpandable,
@@ -293,11 +310,13 @@ export default function TreeViewContextProvider({
                 explicitlyUnselectParent,
                 enableIntermediateSelectState,
 
+                loadingTriggerType,
+                loadRootNodes,
+                rootNodesLoading,
                 // Async loading of children
                 getIsChildrenLoading,
-                getNumberOfNotLoadedChildren,
                 loadChildren,
-                numberOfChildrenToLoad,
+                loadBatchSize,
 
                 renderNodeLabel,
             }}

datahub-web-react/src/app/homeV3/modules/hierarchyViewModule/treeView/types.ts

@@ -24,11 +24,17 @@ export interface TreeNodeProps {
     depth: number;
 }
 
-export interface TreeViewContextType {
+export type LoadingTriggerType = 'button' | 'infiniteScroll';
-    nodes: TreeNode[];
 
+export interface TreeViewContextType {
+    // Tree/node utils and params
+    nodes: TreeNode[];
     getHasParentNode: (node: TreeNode) => boolean;
     getIsRootNode: (node: TreeNode) => boolean;
+    rootNodesLength: number;
+    rootNodesTotal: number;
+    getChildrenLength: (node: TreeNode) => number;
+    getChildrenTotal: (node: TreeNode) => number;
 
     // Expand
     getIsExpandable: (node: TreeNode) => boolean;
@@ -54,13 +60,18 @@ export interface TreeViewContextType {
     explicitlyUnselectParent?: boolean;
     enableIntermediateSelectState?: boolean;
 
-    // Async loading of children
+    // Async loading
     // -------------------------------------------------
+    loadingTriggerType?: LoadingTriggerType;
+    // Optional loading of root nodes
+    loadRootNodes?: () => void;
+    hasMoreRootNodes?: boolean;
+    rootNodesLoading?: boolean;
+
     getIsChildrenLoading: (node: TreeNode) => boolean;
-    getNumberOfNotLoadedChildren: (node: TreeNode) => number;
     loadChildren: (node: TreeNode) => void;
     // Max number of children to load per each loadChildren call
-    numberOfChildrenToLoad: number;
+    loadBatchSize: number;
 
     // Optional custom node label renderer
     renderNodeLabel?: (props: TreeNodeProps) => React.ReactNode;
@@ -104,7 +115,13 @@ export interface TreeViewContextProviderProps {
     renderNodeLabel?: (props: TreeNodeProps) => React.ReactNode;
 
     // Async
+    // Optional pagination/loading of root nodes
+    loadingTriggerType?: LoadingTriggerType;
+    rootNodesTotal?: number;
+    loadRootNodes?: () => void;
+    hasMoreRootNodes?: boolean;
+    rootNodesLoading?: boolean;
     // Callback to load children of a specific node
     loadChildren?: (node: TreeNode) => void;
-    numberOfChildrenToLoad?: number;
+    loadBatchSize?: number;
 }

datahub-web-react/src/app/homeV3/modules/hierarchyViewModule/treeView/utils.ts

@@ -192,3 +192,8 @@ export function sortTree(tree: TreeNode[], nodesSorter: (nodes: TreeNode[]) => T
 
     return nodesSorter(tree).map(traverse);
 }
+
+export function getOutOfTreeSelectedValues(selectedValues: string[], tree: TreeNode[]): string[] {
+    const allTreeValues = getAllValues(tree);
+    return selectedValues.filter((value) => !allTreeValues.includes(value));
+}