Expand TreeGrid

Hi Guys,

In my application I am using TreeGrid and I have trouble in expanding it. I want to expand TreeGrid, when the user clicks on Save button. As the save button gets clicked, I am setting new data to TreeGrid from the backend service and loading it on UI. In this case, when I use the expand feature of TreeGrid, it is unable to find the object, as it is loaded with new data. Below is the sample code. Please provide me the solution in expanding the TreeGrid in this scenario. Right now I am using Vaadin 8.3.2 version.

save = new Button("Save", click -> {
			TreeNode node;
			if (nodeTree.getSelectedItems().size() > 0) {
				node = nodeTree.getSelectedItems().iterator().next();
					// Calling the service
					personalizationService.updateEntity(node, frequency);
					nodeTree.setTreeData(personalizationService.getTreeData());	
				}
				nodeTree.expand(node);
			}

Hey Guys, Any update on this??