diff --git a/examples/src/pages/tests/table/treegrid/selection-api.page.tsx b/examples/src/pages/tests/table/treegrid/selection-api.page.tsx
new file mode 100644
index 000000000..f6971f2ff
--- /dev/null
+++ b/examples/src/pages/tests/table/treegrid/selection-api.page.tsx
@@ -0,0 +1,126 @@
+import * as React from 'react';
+
+import {
+  DataSourceApi,
+  InfiniteTableColumn,
+  TreeDataSource,
+  TreeGrid,
+} from '@infinite-table/infinite-react';
+
+export type FileSystemNode = {
+  name: string;
+  type: 'file' | 'folder';
+  children?: FileSystemNode[] | null;
+  sizeKB?: number;
+  id: string;
+  collapsed?: boolean;
+};
+
+export const nodes: FileSystemNode[] = [
+  {
+    name: 'Documents',
+    type: 'folder',
+    id: '1',
+    children: [
+      {
+        name: 'report.doc',
+        type: 'file',
+        sizeKB: 100,
+        id: '2',
+      },
+      {
+        type: 'folder',
+        name: 'pictures',
+        id: '3',
+        collapsed: true,
+        children: [
+          {
+            name: 'mountain.jpg',
+            type: 'file',
+            sizeKB: 302,
+            id: '5',
+          },
+          {
+            name: 'mountain2.jpg',
+            type: 'file',
+            sizeKB: 352,
+            id: '6',
+          },
+        ],
+      },
+      {
+        type: 'folder',
+        name: 'misc',
+        id: '4',
+        collapsed: true,
+        children: [
+          {
+            name: 'beach.jpg',
+            type: 'file',
+            sizeKB: 2024,
+            id: '7',
+          },
+        ],
+      },
+      {
+        type: 'file',
+        name: 'last.txt',
+        id: '8',
+      },
+    ],
+  },
+];
+
+const columns: Record<string, InfiniteTableColumn<FileSystemNode>> = {
+  name: {
+    field: 'name',
+    defaultWidth: 300,
+    renderTreeIcon: true,
+    renderSelectionCheckBox: true,
+    renderValue: ({ value, data }) => {
+      return (
+        <>
+          {value} - {data!.id}
+        </>
+      );
+    },
+  },
+  type: { field: 'type' },
+  sizeKB: { field: 'sizeKB' },
+};
+
+export default function DataTestPage() {
+  const [dataSourceApi, setDataSourceApi] =
+    React.useState<DataSourceApi<FileSystemNode> | null>(null);
+  return (
+    <React.StrictMode>
+      <button onClick={() => dataSourceApi!.treeApi.selectAll()}>
+        Select all
+      </button>
+      <TreeDataSource<FileSystemNode>
+        onReady={setDataSourceApi}
+        data={nodes}
+        primaryKey="id"
+        nodesKey="children"
+        defaultTreeSelection={{
+          defaultSelection: true,
+          deselectedPaths: [['1', '3']],
+          selectedPaths: [['1', '3', '6']],
+        }}
+      >
+        <TreeGrid<FileSystemNode>
+          wrapRowsHorizontally
+          domProps={{
+            style: {
+              margin: '5px',
+              height: 900,
+              border: '1px solid gray',
+              position: 'relative',
+            },
+          }}
+          columns={columns}
+        />
+      </TreeDataSource>
+    </React.StrictMode>
+  );
+}
diff --git a/examples/src/pages/tests/table/treegrid/selection-api.spec.ts b/examples/src/pages/tests/table/treegrid/selection-api.spec.ts
new file mode 100644
index 000000000..2b978a6bb
--- /dev/null
+++ b/examples/src/pages/tests/table/treegrid/selection-api.spec.ts
@@ -0,0 +1,18 @@
+import { test, expect } from '@testing';
+
+export default test.describe('TreeApi', () => {
+  test('getSelectedLeafNodePaths works', async ({ page, apiModel }) => {
+    await page.waitForInfinite();
+
+    const paths = await apiModel.evaluateTreeApi((treeApi) => {
+      return treeApi.getSelectedLeafNodePaths();
+    });
+
+    expect(paths).toEqual([
+      ['1', '2'],
+      ['1', '3', '6'],
+      ['1', '4', '7'],
+      ['1', '8'],
+    ]);
+  });
+});
diff --git a/examples/src/pages/tests/table/treegrid/selection4.page.tsx b/examples/src/pages/tests/table/treegrid/selection4.page.tsx
new file mode 100644
index 000000000..428df3bd4
--- /dev/null
+++ b/examples/src/pages/tests/table/treegrid/selection4.page.tsx
@@ -0,0 +1,162 @@
+import {
+  DataSourceApi,
+  InfiniteTableColumn,
+  TreeDataSource,
+  TreeGrid,
+  TreeSelectionValue,
+} from '@infinite-table/infinite-react';
+import { useState } from 'react';
+
+type FileSystemNode = {
+  id: string;
+  name: string;
+  children?: FileSystemNode[];
+};
+const nodes: FileSystemNode[] = [
+  {
+    id: '1',
+    name: 'Documents',
+    children: [
+      {
+        id: '10',
+        name: 'Private',
+        children: [
+          {
+            id: '100',
+            name: 'Report.docx',
+          },
+          {
+            id: '101',
+            name: 'Vacation.docx',
+          },
+          {
+            id: '102',
+            name: 'CV.pdf',
+          },
+        ],
+      },
+    ],
+  },
+  {
+    id: '2',
+    name: 'Desktop',
+    children: [
+      {
+        id: '20',
+        name: 'unknown.txt',
+      },
+    ],
+  },
+  {
+    id: '3',
+    name: 'Media',
+    children: [
+      {
+        id: '30',
+        name: 'Music',
+      },
+      {
+        id: '31',
+        name: 'Videos',
+        children: [
+          {
+            id: '310',
+            name: 'Vacation.mp4',
+          },
+          {
+            id: '311',
+            name: 'Youtube',
+            children: [
+              {
+                id: '3110',
+                name: 'Infinity',
+              },
+              {
+                id: '3111',
+                name: 'Infinity 2',
+              },
+            ],
+          },
+        ],
+      },
+    ],
+  },
+];
+
+const columns: Record<string, InfiniteTableColumn<FileSystemNode>> = {
+  name: {
+    field: 'name',
+    header: 'Name',
+    defaultWidth: 500,
+    renderValue: ({ value, rowInfo }) => {
+      return (
+        <div style={{ color: 'red', display: 'inline-block' }}>
+          {rowInfo.id} - {value}
+        </div>
+      );
+    },
+    renderTreeIcon: true,
+    renderSelectionCheckBox: true,
+  },
+};
+
+const defaultTreeSelection: TreeSelectionValue = {
+  defaultSelection: false,
+  selectedPaths: [['3']],
+  deselectedPaths: [['3', '31', '311', '3110']],
+};
+
+export default function App() {
+  const [dataSourceApi, setDataSourceApi] =
+    useState<DataSourceApi<FileSystemNode> | null>();
+
+  return (
+    <>
+      <TreeDataSource
+        onReady={setDataSourceApi}
+        nodesKey="children"
+        primaryKey="id"
+        data={dataSource}
+        defaultTreeSelection={defaultTreeSelection}
+        selectionMode="multi-row"
+        onTreeSelectionChange={(e, { treeSelectionState }) => {
+          console.log(
+            'onTreeSelectionChange',
+            e,
+            treeSelectionState.getSelectedLeafNodePaths(),
+          );
+        }}
+      >
+        <div
+          style={{
+            color: 'var(--infinite-cell-color)',
+            padding: 10,
+            display: 'flex',
+            gap: 10,
+          }}
+        >
+          <button
+            onClick={() => {
+              dataSourceApi!.treeApi.selectAll();
+            }}
+          >
+            Select all
+          </button>
+          <button
+            onClick={() => {
+              dataSourceApi!.treeApi.deselectAll();
+            }}
+          >
+            Deselect all
+          </button>
+        </div>
+
+        <TreeGrid columns={columns} domProps={{ style: { height: '100%' } }} />
+      </TreeDataSource>
+    </>
+  );
+}
+
+const dataSource = () => {
+  return Promise.resolve(nodes);
+};
diff --git a/examples/src/pages/tests/table/treegrid/selection4.spec.ts b/examples/src/pages/tests/table/treegrid/selection4.spec.ts
new file mode 100644
index 000000000..631e7462c
--- /dev/null
+++ b/examples/src/pages/tests/table/treegrid/selection4.spec.ts
@@ -0,0 +1,21 @@
+import { test, expect } from '@testing';
+
+export default test.describe('TreeSelectionProp', () => {
+  test('when defined, makes selectionMode default to multi-row', async ({
+    page,
+  }) => {
+    await page.waitForInfinite();
+
+    const headerCheckbox = await page.locator(
+      '.InfiniteHeader input[type="checkbox"]',
+    );
+    expect(
+      await headerCheckbox?.evaluate((el) => {
+        return {
+          checked: (el as HTMLInputElement).checked,
+          indeterminate: (el as HTMLInputElement).indeterminate,
+        };
+      }),
+    ).toEqual({ checked: false, indeterminate: true });
+  });
+});
diff --git a/examples/src/pages/tests/table/treeview/TreeSelectionState.spec.ts b/examples/src/pages/tests/table/treeview/TreeSelectionState.spec.ts
new file mode 100644
index 000000000..57ae5304a
--- /dev/null
+++ b/examples/src/pages/tests/table/treeview/TreeSelectionState.spec.ts
@@ -0,0 +1,137 @@
+import { test, expect } from '@playwright/test';
+import { TreeSelectionState } from '@src/components/DataSource/TreeSelectionState';
+import { tree } from '@src/utils/groupAndPivot';
+
+type FileSystemNode = {
+  id: string;
+  name: string;
+  children?: FileSystemNode[];
+};
+const nodes: FileSystemNode[] = [
+  {
+    id: '1',
+    name: 'Documents',
+    children: [
+      {
+        id: '10',
+        name: 'Private',
+        children: [
+          {
+            id: '100',
+            name: 'Report.docx',
+          },
+          {
+            id: '101',
+            name: 'Vacation.docx',
+          },
+          {
+            id: '102',
+            name: 'CV.pdf',
+          },
+        ],
+      },
+    ],
+  },
+  {
+    id: '2',
+    name: 'Desktop',
+    children: [
+      {
+        id: '20',
+        name: 'unknown.txt',
+      },
+    ],
+  },
+  {
+    id: '3',
+    name: 'Media',
+    children: [
+      {
+        id: '30',
+        name: 'Music',
+      },
+      {
+        id: '31',
+        name: 'Videos',
+        children: [
+          {
+            id: '310',
+            name: 'Vacation.mp4',
+          },
+          {
+            id: '311',
+            name: 'Youtube',
+            children: [
+              {
+                id: '3110',
+                name: 'Infinity',
+              },
+              {
+                id: '3111',
+                name: 'Infinity 2',
+              },
+            ],
+          },
+        ],
+      },
+    ],
+  },
+];
+
+const treeParams = {
+  getNodeChildren: (node: FileSystemNode) => {
+    return node.children || [];
+  },
+  isLeafNode: (node: FileSystemNode) => {
+    return node.children === undefined;
+  },
+  nodesKey: 'children',
+  toKey: (node: FileSystemNode) => node.id,
+};
+
+const { deepMap } = tree(treeParams, nodes);
+
+export default test.describe.parallel('TreeSelectionState', () => {
+  test('should work properly', async () => {
+    const treeSelectionState = new TreeSelectionState(
+      {
+        defaultSelection: false,
+        selectedPaths: [['3']],
+        deselectedPaths: [['3', '311']],
+      },
+      () => ({
+        treeDeepMap: deepMap,
+      }),
+    );
+
+    expect(treeSelectionState.isNodeSelected(['3'])).toBe(null);
+    expect(treeSelectionState.isNodeSelected(['3', '30'])).toBe(true);
+    expect(treeSelectionState.isNodeSelected(['3', '311'])).toBe(false);
+    expect(treeSelectionState.isNodeSelected(['3', '311', '3110'])).toBe(false);
+  });
+
+  test('should work properly - 2', async () => {
+    const treeSelectionState = new TreeSelectionState(
+      {
+        defaultSelection: false,
+        selectedPaths: [['3']],
+        deselectedPaths: [['3', '31', '311', '3110']],
+      },
+      () => ({
+        treeDeepMap: deepMap,
+      }),
+    );
+
+    expect(treeSelectionState.isNodeSelected(['3'])).toBe(null);
+    expect(treeSelectionState.isNodeSelected(['3', '30'])).toBe(true);
+    expect(treeSelectionState.isNodeSelected(['3', '31', '311'])).toBe(null);
+    expect(treeSelectionState.isNodeSelected(['3', '31', '311', '3110'])).toBe(
+      false,
+    );
+    expect(treeSelectionState.isNodeSelected(['3', '31', '311', '3111'])).toBe(
+      true,
+    );
+
+    expect(treeSelectionState.getSelectedCount()).toBe(3);
+  });
+});
diff --git a/examples/src/pages/tests/table/utils/DeepMap.spec.ts b/examples/src/pages/tests/table/utils/DeepMap.spec.ts
index 62aa3dc1b..a65694414 100644
--- a/examples/src/pages/tests/table/utils/DeepMap.spec.ts
+++ b/examples/src/pages/tests/table/utils/DeepMap.spec.ts
@@ -177,7 +177,9 @@ export default test.describe('DeepMap', () => {
       [4, 10],
       [4, 20],
     ]);
-    expect(map.getKeysStartingWith([5], true)).toEqual([[5, 1]]);
+    expect(map.getKeysStartingWith([5], { excludeSelf: true })).toEqual([
+      [5, 1],
+    ]);
     expect(map.getKeysStartingWith([5])).toEqual([[5], [5, 1]]);
   });
 
@@ -194,14 +196,27 @@ export default test.describe('DeepMap', () => {
     // map.set([2, 1, 2], 7);
     // map.set([2, 1, 3], 8);
 
-    expect(map.getKeysStartingWith([], true, 1)).toEqual([[1], [2]]);
-    expect(map.getKeysStartingWith([], false, 1)).toEqual([[], [1], [2]]);
-
-    expect(map.getKeysStartingWith([1], false, 1)).toEqual([[1], [1, 0]]);
-    expect(map.getKeysStartingWith([1], true, 1)).toEqual([[1, 0]]);
-    expect(map.getKeysStartingWith([1, 1, 2], true, 1)).toEqual([]);
-
-    const value = map.getKeysStartingWith([1], true, 2);
+    expect(
+      map.getKeysStartingWith([], { excludeSelf: true, depthLimit: 1 }),
+    ).toEqual([[1], [2]]);
+    expect(
+      map.getKeysStartingWith([], { excludeSelf: false, depthLimit: 1 }),
+    ).toEqual([[], [1], [2]]);
+
+    expect(
+      map.getKeysStartingWith([1], { excludeSelf: false, depthLimit: 1 }),
+    ).toEqual([[1], [1, 0]]);
+    expect(
+      map.getKeysStartingWith([1], { excludeSelf: true, depthLimit: 1 }),
+    ).toEqual([[1, 0]]);
+    expect(
+      map.getKeysStartingWith([1, 1, 2], { excludeSelf: true, depthLimit: 1 }),
+    ).toEqual([]);
+
+    const value = map.getKeysStartingWith([1], {
+      excludeSelf: true,
+      depthLimit: 2,
+    });
 
     expect(value).toEqual([
       [1, 0],
@@ -504,7 +519,7 @@ export default test.describe('DeepMap', () => {
     ]);
   });
 
-  test('getUnnestedKeysStartingWith should work correctly - 1', () => {
+  test('getKeysOfFirstChildOnEachBranchStartingWith should work correctly - 1', () => {
     let map = new DeepMap<string | number, boolean>();
 
     map.set(['3', '31'], true);
@@ -512,8 +527,12 @@ export default test.describe('DeepMap', () => {
     map.set(['1', '10'], true);
     map.set(['3'], true);
 
-    expect(map.getUnnestedKeysStartingWith([], true)).toEqual([['1'], ['3']]);
-    expect(map.getKeysStartingWith([], true)).toEqual([
+    expect(
+      map.getKeysOfFirstChildOnEachBranchStartingWith([], {
+        excludeSelf: true,
+      }),
+    ).toEqual([['1'], ['3']]);
+    expect(map.getKeysStartingWith([], { excludeSelf: true })).toEqual([
       ['3'],
       ['3', '31'],
       ['1'],
@@ -528,10 +547,15 @@ export default test.describe('DeepMap', () => {
     map.set(['3'], true);
     map.set(['1'], true);
 
-    expect(map.getUnnestedKeysStartingWith([], true)).toEqual([['3'], ['1']]);
+    expect(
+      map.getKeysOfFirstChildOnEachBranchStartingWith([], {
+        excludeSelf: true,
+      }),
+    ).toEqual([['3'], ['1']]);
   });
 
-  test('getUnnestedKeysStartingWith should work correctly - 2', () => {
+  // so get unnested means get the first key you find going down the tree on a branch
+  test('getKeysOfFirstChildOnEachBranchStartingWith should work correctly - 2', () => {
     let map = new DeepMap<string | number, boolean>();
 
     map.set(['3', '31', '300'], true);
@@ -549,21 +573,148 @@ export default test.describe('DeepMap', () => {
       ['3', '31', '400'],
       ['1'],
     ];
-    expect(map.getUnnestedKeysStartingWith([], true)).toEqual(expected);
-    expect(map.getUnnestedKeysStartingWith(['3', '31'])).toEqual([
+    expect(
+      map.getKeysOfFirstChildOnEachBranchStartingWith([], {
+        excludeSelf: true,
+      }),
+    ).toEqual(expected);
+    expect(
+      map.getKeysOfFirstChildOnEachBranchStartingWith(['3', '31']),
+    ).toEqual([
       ['3', '31', '300'],
       ['3', '31', '400'],
     ]);
 
     map.set([], true);
 
-    expect(map.getUnnestedKeysStartingWith([])).toEqual([[]]);
-    expect(map.getUnnestedKeysStartingWith([], true)).toEqual(expected);
-    expect(map.getUnnestedKeysStartingWith([], true, 1)).toEqual([['1']]);
-    expect(map.getUnnestedKeysStartingWith([], true, 2)).toEqual([
-      ['3', '30'],
-      ['1'],
+    expect(map.getKeysOfFirstChildOnEachBranchStartingWith([])).toEqual([[]]);
+    expect(
+      map.getKeysOfFirstChildOnEachBranchStartingWith([], {
+        excludeSelf: true,
+      }),
+    ).toEqual(expected);
+    expect(
+      map.getKeysOfFirstChildOnEachBranchStartingWith([], {
+        excludeSelf: true,
+        depthLimit: 1,
+      }),
+    ).toEqual([['1']]);
+    expect(
+      map.getKeysOfFirstChildOnEachBranchStartingWith([], {
+        excludeSelf: true,
+        depthLimit: 2,
+      }),
+    ).toEqual([['3', '30'], ['1']]);
+  });
+
+  test('getKeysOfFirstChildOnEachBranchStartingWith should work correctly - 3', () => {
+    // let map = new DeepMap<string | number, boolean>();
+    // map.set(['1'], true);
+    // map.set(['1', '10'], true);
+    // map.set(['3'], true);
+    // expect(map.getKeysOfFirstChildOnEachBranchStartingWith([], true)).toEqual([
+    //   ['1', '10'],
+    //   ['3'],
+    // ]);
+    let map = new DeepMap<string | number, boolean>();
+    map.set(['1'], true);
+    map.set(['3', '31'], true);
+    map.set(['1', '10'], true);
+    map.set(['3'], true);
+
+    expect(
+      map.getKeysOfFirstChildOnEachBranchStartingWith([], {
+        excludeSelf: true,
+      }),
+    ).toEqual([['1'], ['3']]);
+  });
+
+  test('getKeysForLeafNodesStartingWith should work correctly - with excludeSelf: true', () => {
+    let map = new DeepMap<string | number, boolean>();
+    map.set(['1'], true);
+    map.set(['3', '31'], true);
+    map.set(['1', '10'], true);
+    map.set(['3'], true);
+
+    expect(
+      map.getKeysForLeafNodesStartingWith([], {
+        excludeSelf: true,
+        respectOrder: true,
+      }),
+    ).toEqual([
+      ['3', '31'],
+      ['1', '10'],
     ]);
+
+    map = new DeepMap<string | number, boolean>();
+    map.set(['3', '31'], true);
+    map.set(['1', '10'], true);
+    map.set(['3'], true);
+
+    expect(
+      map.getKeysForLeafNodesStartingWith([], {
+        excludeSelf: true,
+        respectOrder: true,
+      }),
+    ).toEqual([
+      ['3', '31'],
+      ['1', '10'],
+    ]);
+
+    map = new DeepMap<string | number, boolean>();
+    map.set(['1', '10'], true);
+    map.set(['3', '31'], true);
+    map.set(['3'], true);
+
+    expect(
+      map.getKeysForLeafNodesStartingWith(['1', '10'], {
+        excludeSelf: true,
+        respectOrder: true,
+      }),
+    ).toEqual([]);
+    expect(
+      map.getKeysForLeafNodesStartingWith(['1', '10'], {
+        excludeSelf: false,
+        respectOrder: true,
+      }),
+    ).toEqual([['1', '10']]);
+  });
+
+  test('getKeysForLeafNodesStartingWith should work correctly - with depthLimit', () => {
+    let map = new DeepMap<string | number, boolean>();
+    map.set(['1'], true);
+    map.set(['3', '31'], true);
+    map.set(['1', '10'], true);
+    map.set(['3'], true);
+    map.set(['4'], true);
+
+    expect(
+      map.getKeysForLeafNodesStartingWith([], {
+        excludeSelf: true,
+        depthLimit: 1,
+      }),
+    ).toEqual([['4']]);
+    expect(
+      map.getKeysForLeafNodesStartingWith([], {
+        excludeSelf: true,
+        depthLimit: 2,
+        respectOrder: true,
+      }),
+    ).toEqual([['3', '31'], ['1', '10'], ['4']]);
+
+    map = new DeepMap<string | number, boolean>();
+    expect(
+      map.getKeysForLeafNodesStartingWith([], { excludeSelf: false }),
+    ).toEqual([]);
+
+    map = new DeepMap<string | number, boolean>();
+    map.set([], true);
+    expect(
+      map.getKeysForLeafNodesStartingWith([], { excludeSelf: false }),
+    ).toEqual([[]]);
+    expect(
+      map.getKeysForLeafNodesStartingWith([], { excludeSelf: true }),
+    ).toEqual([]);
   });
 
   test('visit depth first, with index', () => {
diff --git a/examples/src/pages/tests/testUtils/InfiniteTableApiModel.ts b/examples/src/pages/tests/testUtils/InfiniteTableApiModel.ts
index 8446940c0..ce7b5503a 100644
--- a/examples/src/pages/tests/testUtils/InfiniteTableApiModel.ts
+++ b/examples/src/pages/tests/testUtils/InfiniteTableApiModel.ts
@@ -1,3 +1,4 @@
+import { TreeApi } from '@src/components/DataSource/TreeApi';
 import { Page } from '@playwright/test';
 import { DataSourceApi } from '@src/components/DataSource/types';
 import { InfiniteTableApi } from '@src/components/InfiniteTable/types';
@@ -35,4 +36,13 @@ export class InfiniteTableApiModel {
       return fn(api);
     }, evaluateFn.toString());
   }
+
+  async evaluateTreeApi(evaluateFn: (treeApi: TreeApi<any>) => void) {
+    return this.page.evaluate((evaluateFn) => {
+      const api = (window as any).DATA_SOURCE_API as DataSourceApi<any>;
+
+      const fn = eval(evaluateFn);
+      return fn(api.treeApi);
+    }, evaluateFn.toString());
+  }
 }
diff --git a/source/src/components/DataSource/RowSelectionState.ts b/source/src/components/DataSource/RowSelectionState.ts
index e944822a6..68d8ff0f4 100644
--- a/source/src/components/DataSource/RowSelectionState.ts
+++ b/source/src/components/DataSource/RowSelectionState.ts
@@ -92,14 +92,23 @@ export class RowSelectionState<T = any> {
 
   getGroupKeysDirectlyInsideGroup(groupKeys: any[]) {
     const { groupDeepMap } = this.getConfig();
-    return groupDeepMap?.getKeysStartingWith(groupKeys, true, 1) || [];
+    return (
+      groupDeepMap?.getKeysStartingWith(groupKeys, {
+        excludeSelf: true,
+        depthLimit: 1,
+      }) || []
+    );
   }
 
   getAllPrimaryKeysInsideGroup(groupKeys: any[]): any[] {
     const { groupDeepMap } = this.getConfig();
 
     if (!groupKeys.length) {
-      const topLevelKeys = groupDeepMap?.getKeysStartingWith([], true, 1) || [];
+      const topLevelKeys =
+        groupDeepMap?.getKeysStartingWith([], {
+          excludeSelf: true,
+          depthLimit: 1,
+        }) || [];
 
       return topLevelKeys
         .map((groupKeys) => this.getAllPrimaryKeysInsideGroup(groupKeys))
@@ -434,7 +443,9 @@ export class RowSelectionState<T = any> {
       return;
     }
     // retrieve any selection under this group
-    const selectedKeys = this.selectedMap.getKeysStartingWith(groupKeys, true);
+    const selectedKeys = this.selectedMap.getKeysStartingWith(groupKeys, {
+      excludeSelf: true,
+    });
 
     // and clean it up
     selectedKeys.forEach((groupKeys) => {
@@ -473,10 +484,9 @@ export class RowSelectionState<T = any> {
       return;
     }
     // retrieve any deselection under this group
-    const deselectedKeys = this.deselectedMap.getKeysStartingWith(
-      groupKeys,
-      true,
-    );
+    const deselectedKeys = this.deselectedMap.getKeysStartingWith(groupKeys, {
+      excludeSelf: true,
+    });
 
     // and clean it up
     deselectedKeys.forEach((groupKeys) => {
@@ -681,18 +691,16 @@ export class RowSelectionState<T = any> {
       const totalCount = this.getGroupCount(groupKeys);
 
       //explicitly selected rows
-      let selectedCount = this.selectedMap.getKeysStartingWith(
-        groupKeys,
-        true,
-        1,
-      ).length;
+      let selectedCount = this.selectedMap.getKeysStartingWith(groupKeys, {
+        excludeSelf: true,
+        depthLimit: 1,
+      }).length;
 
       // explicitly deselected rows
-      let deselectedCount = this.deselectedMap.getKeysStartingWith(
-        groupKeys,
-        true,
-        1,
-      ).length;
+      let deselectedCount = this.deselectedMap.getKeysStartingWith(groupKeys, {
+        excludeSelf: true,
+        depthLimit: 1,
+      }).length;
 
       const notSpecifiedCount = totalCount - (selectedCount + deselectedCount);
 
diff --git a/source/src/components/DataSource/TreeApi.ts b/source/src/components/DataSource/TreeApi.ts
index 80fcf1bb2..16cc421f9 100644
--- a/source/src/components/DataSource/TreeApi.ts
+++ b/source/src/components/DataSource/TreeApi.ts
@@ -1,6 +1,9 @@
 import { DataSourceApi, DataSourceComponentActions } from '.';
 
-import { InfiniteTableRowInfo } from '../InfiniteTable/types';
+import {
+  InfiniteTable_Tree_RowInfoLeafNode,
+  InfiniteTableRowInfo,
+} from '../InfiniteTable/types';
 import { getRowInfoAt } from './dataSourceGetters';
 import { NodePath, TreeExpandState } from './TreeExpandState';
 import {
@@ -41,7 +44,7 @@ export type TreeExpandStateApi<T> = {
   getRowInfoByPath(nodePath: any[]): InfiniteTableRowInfo<T> | null;
 };
 
-type TreeSelectionApi<_T = any> = {
+type TreeSelectionApi<T = any> = {
   get allRowsSelected(): boolean;
   isNodeSelected(nodePath: NodePath): boolean | null;
 
@@ -58,6 +61,19 @@ type TreeSelectionApi<_T = any> = {
   expandAll(): void;
   collapseAll(): void;
   deselectAll(): void;
+
+  getSelectedLeafNodePaths(config?: {
+    rootNodePath?: NodePath;
+    treeSelectionState?: TreeSelectionState<T>;
+  }): NodePath[];
+  getDeselectedLeafNodePaths(config?: {
+    rootNodePath?: NodePath;
+    treeSelectionState?: TreeSelectionState<T>;
+  }): NodePath[];
+  getSelectedLeafRowInfos(config?: {
+    rootNodePath?: NodePath;
+    treeSelectionState?: TreeSelectionState<T>;
+  }): InfiniteTable_Tree_RowInfoLeafNode<T>[];
 };
 
 export type TreeApi<T> = TreeExpandStateApi<T> & TreeSelectionApi<T>;
@@ -94,6 +110,75 @@ export class TreeApiImpl<T> implements TreeApi<T> {
     this.dataSourceApi = param.dataSourceApi;
   }
 
+  getSelectedLeafNodePaths(config?: {
+    rootNodePath?: NodePath;
+    treeSelectionState?: TreeSelectionState<T>;
+  }): NodePath[] {
+    const { treePaths } = this.getState();
+    const treeSelectionState =
+      config?.treeSelectionState || this.getState().treeSelectionState;
+
+    if (!treeSelectionState) {
+      return [];
+    }
+
+    return treeSelectionState.getSelectedLeafNodePaths(
+      config?.rootNodePath,
+      treePaths,
+    );
+  }
+
+  getDeselectedLeafNodePaths(config?: {
+    rootNodePath?: NodePath;
+    treeSelectionState?: TreeSelectionState<T>;
+  }): NodePath[] {
+    const { treePaths } = this.getState();
+    const treeSelectionState =
+      config?.treeSelectionState || this.getState().treeSelectionState;
+
+    if (!treeSelectionState) {
+      return [];
+    }
+
+    return treeSelectionState.getDeselectedLeafNodePaths(
+      config?.rootNodePath,
+      treePaths,
+    );
+  }
+
+  getSelectedLeafRowInfos(config?: {
+    rootNodePath?: NodePath;
+    treeSelectionState?: TreeSelectionState<T>;
+  }): InfiniteTable_Tree_RowInfoLeafNode<T>[] {
+    const { treePaths } = this.getState();
+    const treeSelectionState =
+      config?.treeSelectionState || this.getState().treeSelectionState;
+
+    if (!treePaths || !treeSelectionState) {
+      return [];
+    }
+    const rootNodePath = config?.rootNodePath || [];
+
+    const selectedLeafRowInfos: InfiniteTable_Tree_RowInfoLeafNode<T>[] = [];
+
+    treePaths.getLeafNodesStartingWith(
+      rootNodePath,
+      (pair) => {
+        if (treeSelectionState.isNodeSelected(pair.keys)) {
+          const rowInfo = this.getRowInfoByPath(
+            pair.keys,
+          ) as InfiniteTable_Tree_RowInfoLeafNode<T>;
+          if (rowInfo) {
+            selectedLeafRowInfos.push(rowInfo);
+          }
+        }
+      },
+      { excludeSelf: true },
+    );
+
+    return selectedLeafRowInfos;
+  }
+
   setNodeSelection = (
     nodePath: NodePath,
     selected: boolean,
diff --git a/source/src/components/DataSource/TreeSelectionState.ts b/source/src/components/DataSource/TreeSelectionState.ts
index 5c06a0eaa..2b4793db9 100644
--- a/source/src/components/DataSource/TreeSelectionState.ts
+++ b/source/src/components/DataSource/TreeSelectionState.ts
@@ -49,6 +49,38 @@ export class TreeSelectionState<T = any> {
     return this.getConfig().treeDeepMap;
   }
 
+  getSelectedLeafNodePaths(
+    rootNodePath: NodePath = [],
+    treeDeepMap?: DeepMap<any, any>,
+  ) {
+    treeDeepMap = treeDeepMap || this.getConfig().treeDeepMap;
+    const selectedLeafPaths: NodePath[] = [];
+
+    treeDeepMap.getLeafNodesStartingWith(rootNodePath, (pair) => {
+      if (this.isNodeSelected(pair.keys)) {
+        selectedLeafPaths.push(pair.keys);
+      }
+    });
+
+    return selectedLeafPaths;
+  }
+
+  getDeselectedLeafNodePaths(
+    rootNodePath: NodePath = [],
+    treeDeepMap?: DeepMap<any, any>,
+  ) {
+    treeDeepMap = treeDeepMap || this.getConfig().treeDeepMap;
+    const deselectedLeafPaths: NodePath[] = [];
+
+    treeDeepMap.getLeafNodesStartingWith(rootNodePath, (pair) => {
+      if (!this.isNodeSelected(pair.keys)) {
+        deselectedLeafPaths.push(pair.keys);
+      }
+    });
+
+    return deselectedLeafPaths;
+  }
+
   isLeafNode(nodePath: NodePath) {
     return !this.getTreeDeepMap().has(nodePath);
   }
@@ -193,7 +225,9 @@ export class TreeSelectionState<T = any> {
 
     const childPaths = selectionMap
       // todo this could be replaced with a .hasKeysUnder call (to be implemented in the deep map later)
-      .getUnnestedKeysStartingWith(nodePath, true)
+      .getKeysOfFirstChildOnEachBranchStartingWith(nodePath, {
+        excludeSelf: true,
+      })
       .sort(shortestToLongest);
 
     if (!childPaths.length) {
diff --git a/source/src/components/DataSource/state/getInitialState.ts b/source/src/components/DataSource/state/getInitialState.ts
index f1a9ed2b0..210685ee9 100644
--- a/source/src/components/DataSource/state/getInitialState.ts
+++ b/source/src/components/DataSource/state/getInitialState.ts
@@ -890,15 +890,20 @@ export function getMappedCallbacks<T>() {
           callbackParams,
         };
       }
-      const callbackParams: Parameters<DataSourcePropOnTreeSelectionChange_MultiNode> =
-        [
-          (treeSelection as TreeSelectionState).getState(),
-          {
-            selectionMode: 'multi-row',
-            lastUpdatedNodePath: state.lastSelectionUpdatedNodePathRef.current,
-            dataSourceApi: state.__apiRef.current!,
-          },
-        ];
+      const dataSourceApi = state.__apiRef.current!;
+      const treeApi = dataSourceApi.treeApi;
+      const callbackParams: Parameters<
+        DataSourcePropOnTreeSelectionChange_MultiNode<T>
+      > = [
+        (treeSelection as TreeSelectionState).getState(),
+        {
+          treeSelectionState: treeSelection as TreeSelectionState<T>,
+          selectionMode: 'multi-row',
+          lastUpdatedNodePath: state.lastSelectionUpdatedNodePathRef.current,
+          dataSourceApi,
+          treeApi,
+        },
+      ];
       return {
         callbackParams,
       };
diff --git a/source/src/components/DataSource/types.ts b/source/src/components/DataSource/types.ts
index cfcf848a2..a73a6e1b3 100644
--- a/source/src/components/DataSource/types.ts
+++ b/source/src/components/DataSource/types.ts
@@ -449,11 +449,13 @@ export type DataSourcePropOnRowSelectionChange_MultiRow = (
 ) => void;
 
 export type DataSourcePropOnTreeSelectionChange_MultiNode<T = any> = (
-  treeSelection: TreeSelectionStateObject,
+  treeSelectionStateObject: TreeSelectionStateObject,
   params: {
+    treeSelectionState: TreeSelectionState<T>;
     selectionMode: 'multi-row';
     lastUpdatedNodePath: NodePath | null;
     dataSourceApi: DataSourceApi<T>;
+    treeApi: TreeApi<T>;
   },
 ) => void;
 export type DataSourcePropOnRowSelectionChange_SingleRow = (
diff --git a/source/src/components/HeadlessTable/GridCellManager.ts b/source/src/components/HeadlessTable/GridCellManager.ts
index fd4b7ade4..ec24f5010 100644
--- a/source/src/components/HeadlessTable/GridCellManager.ts
+++ b/source/src/components/HeadlessTable/GridCellManager.ts
@@ -330,7 +330,10 @@ export class GridCellManager<T_ADDITIONAL_CELL_INFO> extends Logger {
   }
 
   getColumnsWithCells() {
-    const positions = this.matrix.getUnnestedKeysStartingWith([], true);
+    const positions = this.matrix.getKeysOfFirstChildOnEachBranchStartingWith(
+      [],
+      { excludeSelf: true },
+    );
 
     // we use a set to remove duplicates
     const columns = new Set(positions.map((pos) => pos[1]));
diff --git a/source/src/components/InfiniteTable/components/draggable/DragList.tsx b/source/src/components/InfiniteTable/components/draggable/DragList.tsx
index a62ed4a16..3dfd98812 100644
--- a/source/src/components/InfiniteTable/components/draggable/DragList.tsx
+++ b/source/src/components/InfiniteTable/components/draggable/DragList.tsx
@@ -222,7 +222,7 @@ function useInteractionTarget(
 }
 
 export const DragList = (props: DragListProps) => {
-  const domRef = React.useRef<HTMLElement>(null);
+  const domRef = React.useRef<HTMLElement | null>(null);
 
   const { dropTargetListId, dragSourceListId } = useDragDropProvider();
 
diff --git a/source/src/utils/DeepMap/index.ts b/source/src/utils/DeepMap/index.ts
index df16a2f8c..3550cc277 100644
--- a/source/src/utils/DeepMap/index.ts
+++ b/source/src/utils/DeepMap/index.ts
@@ -10,6 +10,21 @@ type Pair<KeyType, ValueType> = {
   revision?: number;
 };
 
+type PairWithKeys<KeyType, ValueType> = Pair<KeyType, ValueType> & {
+  keys: KeyType[];
+};
+
+type StartingWithMethodOptions = {
+  excludeSelf?: boolean;
+  depthLimit?: number;
+};
+
+type StartingWithMethodOptionsWithOrder = {
+  excludeSelf?: boolean;
+  depthLimit?: number;
+  respectOrder?: boolean;
+};
+
 const SORT_ASC_REVISION = (p1: Pair<any, any>, p2: Pair<any, any>) =>
   sortAscending(p1.revision!, p2.revision!);
 
@@ -51,8 +66,7 @@ export class DeepMap<KeyType, ValueType> {
 
   getValuesStartingWith(
     keys: KeyType[],
-    excludeSelf?: boolean,
-    depthLimit?: number,
+    options?: StartingWithMethodOptions,
   ): ValueType[] {
     const result: ValueType[] = [];
     this.getStartingWith(
@@ -60,8 +74,7 @@ export class DeepMap<KeyType, ValueType> {
       (_keys, value) => {
         result.push(value);
       },
-      excludeSelf,
-      depthLimit,
+      options,
     );
 
     return result;
@@ -69,8 +82,7 @@ export class DeepMap<KeyType, ValueType> {
 
   getEntriesStartingWith(
     keys: KeyType[],
-    excludeSelf?: boolean,
-    depthLimit?: number,
+    options?: StartingWithMethodOptions,
   ): [KeyType[], ValueType][] {
     const result: [KeyType[], ValueType][] = [];
     this.getStartingWith(
@@ -78,24 +90,82 @@ export class DeepMap<KeyType, ValueType> {
       (keys, value) => {
         result.push([keys, value]);
       },
-      excludeSelf,
-      depthLimit,
+      options,
     );
 
     return result;
   }
 
-  getUnnestedKeysStartingWith(
+  getLeafNodesStartingWith<T>(
     keys: KeyType[],
-    excludeSelf?: boolean,
-    depthLimit?: number,
+    withPair: (pair: PairWithKeys<KeyType, ValueType>) => T,
+    options?: StartingWithMethodOptionsWithOrder,
+  ): T[] {
+    const { respectOrder = false } = options || {};
+    const pairs: PairWithKeys<KeyType, ValueType>[] = [];
+    const result: T[] = [];
+    this.getStartingWith(
+      keys,
+      (_keys, _value, pair) => {
+        if (!pair.map && pair.hasOwnProperty('value')) {
+          if (respectOrder) {
+            pairs.push(pair);
+          } else {
+            result.push(withPair(pair));
+          }
+        }
+      },
+      options,
+    );
+
+    if (respectOrder) {
+      return pairs.sort(SORT_ASC_REVISION).map(withPair);
+    }
+
+    return result;
+  }
+
+  getKeysForLeafNodesStartingWith(
+    keys: KeyType[],
+    options?: StartingWithMethodOptionsWithOrder,
   ): KeyType[][] {
+    return this.getLeafNodesStartingWith(keys, (pair) => pair.keys, options);
+  }
+
+  getKeysAndValuesForLeafNodesStartingWith(
+    keys: KeyType[],
+    options?: StartingWithMethodOptionsWithOrder,
+  ): [KeyType[], ValueType][] {
+    return this.getLeafNodesStartingWith(
+      keys,
+      (pair) => [pair.keys, pair.value!],
+      options,
+    );
+  }
+
+  /**
+   * @param keys - the keys of the node to start from
+   * @param excludeSelf - whether to exclude the start node
+   * @param depthLimit - the depth limit
+   * @returns the keys of the first child on each branch (starting from the node with the given keys)
+   */
+  getKeysOfFirstChildOnEachBranchStartingWith(
+    keys: KeyType[],
+    options?: StartingWithMethodOptionsWithOrder,
+  ): KeyType[][] {
+    const { excludeSelf, depthLimit, respectOrder = true } = options || {};
     const pairs: (Pair<KeyType, ValueType> & { keys: KeyType[] })[] = [];
 
+    const result: KeyType[][] = [];
+
     const fn: (pair: Pair<KeyType, ValueType> & { keys: KeyType[] }) => void = (
       pair,
     ) => {
-      pairs.push(pair);
+      if (respectOrder) {
+        pairs.push(pair);
+      } else {
+        result.push(pair.keys);
+      }
     };
 
     let currentMap = this.map;
@@ -139,7 +209,10 @@ export class DeepMap<KeyType, ValueType> {
       }
     }
     if (stop) {
-      return pairs.sort(SORT_ASC_REVISION).map((pair) => pair.keys);
+      if (respectOrder) {
+        return pairs.sort(SORT_ASC_REVISION).map((pair) => pair.keys);
+      }
+      return result;
     }
 
     this.visitWithNext(
@@ -154,13 +227,15 @@ export class DeepMap<KeyType, ValueType> {
       excludeSelf,
     );
 
-    return pairs.sort(SORT_ASC_REVISION).map((pair) => pair.keys);
+    if (respectOrder) {
+      return pairs.sort(SORT_ASC_REVISION).map((pair) => pair.keys);
+    }
+    return result;
   }
 
   getKeysStartingWith(
     keys: KeyType[],
-    excludeSelf?: boolean,
-    depthLimit?: number,
+    options?: StartingWithMethodOptions,
   ): KeyType[][] {
     const result: KeyType[][] = [];
     this.getStartingWith(
@@ -168,8 +243,7 @@ export class DeepMap<KeyType, ValueType> {
       (keys) => {
         result.push(keys);
       },
-      excludeSelf,
-      depthLimit,
+      options,
     );
 
     return result;
@@ -177,10 +251,14 @@ export class DeepMap<KeyType, ValueType> {
 
   private getStartingWith(
     keys: KeyType[],
-    fn: (key: KeyType[], value: ValueType) => void,
-    excludeSelf?: boolean,
-    depthLimit?: number,
+    fn: (
+      key: KeyType[],
+      value: ValueType,
+      pair: PairWithKeys<KeyType, ValueType>,
+    ) => void,
+    options?: StartingWithMethodOptions,
   ) {
+    const { excludeSelf, depthLimit } = options || {};
     let currentMap = this.map;
     let pair: Pair<KeyType, ValueType> | undefined;
     let stop = false;
@@ -210,7 +288,7 @@ export class DeepMap<KeyType, ValueType> {
 
     if (pair && pair.value !== undefined) {
       if (!excludeSelf) {
-        fn(keys, pair.value!);
+        fn(keys, pair.value!, { ...pair, keys });
       }
     }
     if (stop) {
@@ -219,8 +297,8 @@ export class DeepMap<KeyType, ValueType> {
 
     this.visitWithNext(
       keys,
-      (value, keys, _i, next) => {
-        fn(keys, value);
+      (value, keys, _i, next, pair) => {
+        fn(keys, value, { ...pair, keys });
         next?.();
       },
       false,
@@ -722,3 +800,6 @@ export class DeepMap<KeyType, ValueType> {
   //   return makeIterator();
   // }
 }
+
+// @ts-ignore
+globalThis.DeepMap = DeepMap;
diff --git a/source/src/utils/debugPackage.ts b/source/src/utils/debugPackage.ts
index 39f38d265..73e1d3da3 100644
--- a/source/src/utils/debugPackage.ts
+++ b/source/src/utils/debugPackage.ts
@@ -171,8 +171,9 @@ function isChannelTargeted(channel: string, permissionToken: string) {
     : tokenParts;
 
   if (
-    partsMap.getKeysStartingWith(storagePartsWithoutWildcard, hasWildcard)
-      .length > 0
+    partsMap.getKeysStartingWith(storagePartsWithoutWildcard, {
+      excludeSelf: hasWildcard,
+    }).length > 0
   ) {
     const remainingParts = tokenParts.slice(indexOfToken + 1);
     if (remainingParts.length) {
diff --git a/source/src/utils/groupAndPivot/treeUtils.ts b/source/src/utils/groupAndPivot/treeUtils.ts
index 3a83f4139..8667c3f8e 100644
--- a/source/src/utils/groupAndPivot/treeUtils.ts
+++ b/source/src/utils/groupAndPivot/treeUtils.ts
@@ -41,7 +41,10 @@ export function toTreeDataArray<T = any, RESULT_T = any>(
   }
 
   function traverse(path: string[], arr: RESULT_T[]) {
-    const nextLevelKeys = treeMap.getKeysStartingWith(path, true, 1);
+    const nextLevelKeys = treeMap.getKeysStartingWith(path, {
+      excludeSelf: true,
+      depthLimit: 1,
+    });
 
     for (const nextLevelKey of nextLevelKeys) {
       const p = [...path, nextLevelKey[nextLevelKey.length - 1]];