state persistence - cells not expanding upon setting state

I am working on adding state persistence to my application and have come across an issue where the cells are not expanding when they should be.

  <igx-tree-grid igxGridState #disciplineSchedulesGrid (columnInit)="onDisciplineScheduleColumnInit($event)" [data]="disciplineList" [isLoading]="!disciplineList" [autoGenerate]="false"
  [allowFiltering]="true" [emptyGridTemplate]="emptyGrid" childDataKey="schedules" expansionDepth="0"
  [rowSelection]="'single'" [hideRowSelectors]="true"
  (rowSelectionChanging)="selectScheduleGridRow($event)">
    <igx-grid-toolbar>
      <igx-grid-toolbar-title>Rate Review</igx-grid-toolbar-title>

      <div *ngIf="bidderRates" style="display:flex">
        <igx-select #select name="bidderSelect" [placeholder]="'Please select Bidder'" (selectionChanging)="selectBidderId($event)" id="bidderSelect" class="vertical-centre">
          <igx-select-item *ngIf="!bidderList" [value]="null">
            No Bidders created
          </igx-select-item>
          <igx-select-item *ngFor="let bidder of bidderList" [value]="bidder.id">
            {{ bidder.bidderId }} <span *ngIf="bidder.name"> - {{ bidder.name }}</span>
          </igx-select-item>
        </igx-select>
      </div>

      <button type="button" id="btnAdd" igxButton="contained" igxRipple igxRippleCentered="true" (click)="gridStateComponent.showAddDialog(gridId, rateGridState)">Add State</button>

      <igx-select *ngIf="gridStateComponent" #select name="gridSelect" [(ngModel)]="gridStateComponent.selectedGridStateId" (selectionChanging)="gridStateComponent.selectGridStates($event)" id="gridSelect" class="vertical-centre" style="padding-left: 0;">
        <igx-select-item [value]="'default'" (click)="restoreGrid()">
          DEFAULT
        </igx-select-item>
        <igx-select-item *ngFor="let gridState of gridStateComponent.filteredGridStates" [value]="gridState.id">
          {{ gridState.stateName }}
        </igx-select-item>
      </igx-select>

      <button *ngIf="canDownload" #btnDownload id="btnDownload" igxButton="contained" igxRipple igxRippleCentered="true" (click)="startExcelDownload()" [disabled]="downloadRunning || (bidderRates && !bidderId)">
        <igx-icon name="download" style="padding-right: 1rem;"></igx-icon>Download Rates
      </button>
    </igx-grid-toolbar>
    <igx-column field="id" header="id" [editable]="false" [sortable]="false" [groupable]="false" [filterable]="false" [resizable]="true" [hidden]="true"></igx-column>
    <igx-column field="code" header="Disciplines & Schedules" [sortable]="true" [filterable]="true" [resizable]="true">
      <ng-template #disciplineScheduleCodeTemplate igxCell let-cell="cell">
        <span class="igx-grid__td" style="width: 100%">{{ cell.value | SORCode: disciplineCodeFormat }} - {{ cell.row.data.description }}</span>
      </ng-template>
    </igx-column>
  </igx-tree-grid>

the method for selecting the state:

 selectGridStates(event: ISelectionEventArgs) {
    if (event.newSelection.value === "default") {
      this.selectedGridStateId = 'default';
      return;
    }

    this.selectedGridStateId = this.gridStates.filter(gs => gs.id == event.newSelection.value)[0].id;

    const stateName = this.filteredGridStates.filter(gs => gs.id === this.selectedGridStateId)[0].stateName;

    this.gridIds.forEach(gridId => {
      const state = this.gridStates.filter(gs => gs.gridId === gridId && gs.stateName === stateName)[0].state;
      this.gridStateDirectives[gridId].setState(JSON.parse(state));
      this.cdr.detectChanges();

          // Manually apply cell expansion state
      const expansionState = JSON.parse(state).expansion;
      if (expansionState) {
        expansionState.forEach(exp => {
          if (exp[1]) {
            const row = this.gridStateDirectives[gridId]?.grid?._dataView?.find(
              r => r.key?.id === exp[0].id);
           row.expanded = true;
            this.gridStateDirectives[gridId].grid.expandRow(row.id);
          }
        });
      }
    });
   }

I thought just using state would do this but it doesn't seem and when trying to manually get them to expand by doing:

        this.gridStateDirectives[gridId].grid.expandRow(row.id);

I am getting the error:

RROR TypeError: Cannot read properties of undefined (reading 'expanded') at _IgxTreeGridAPIService.allow_expansion_state_change (infragistics-igniteui-angular.mjs:111010:13)

I'm not sure where I'm going wrong, as far as I can tell it is storing the rows that are selected and expanded

The cell selection is specified in the cellSelection array:

"cellSelection": [{ "rowStart": 8, "rowEnd": 8, "columnStart": 0, "columnEnd": 0 }] This indicates that a single cell is selected at row 8, column 0.

The expansion information is shown in the expansion array:

"expansion": [[{ "code": 20, "description": "STEELWORK", "libraryId": "aff67ced-12aa-4b7b-a3af-08db8f460e7f", "jointDisciplineId": null, "schedules": [...] }]] This shows that a row with code "20" and description "STEELWORK" is expanded.

so I'm just confused as to why these changes are not taking place when setting the state?

I figured out the issue, the tree grid was missing a primary key which just can be set by simply adding:

primaryKey="id"

then once this is done I can do:

  const expansionState = JSON.parse(state).expansion;
  if (expansionState) {
    expansionState.forEach(exp => {
      if (exp[1]) {
        this.gridStateDirectives[gridId].grid.expandRow(exp[0].id);
      }
    });
  }

and that will allow to expand the necessary row.

Edit: Figured out now that doing this:

  const expansionState = JSON.parse(state).expansion;
  if (expansionState) {
    expansionState.forEach(exp => {
      if (exp[1]) {
        this.gridStateDirectives[gridId].grid.expandRow(exp[0].id);
      }
    });
  }

isn't actually necessary. When you set the primary key as the id, the states are saved differently. exp[0] will just contain an id now, but since the primary keys are set, when you use the set state method, it will just expand the rows and set the selected rows automatically