Merge "Refactor restore to deal with heterogeneous datasets"