def _test_sharding(self, row_count, cpu_count, expected): df = pd.DataFrame({_DUMMY_COL: list(range(row_count))}) with fixed_core_count(cpu_count): shards = list(file_io.iter_shard_dataframe(df, _ROWS_PER_CORE)) result = [[j[_DUMMY_COL].tolist() for j in i] for i in shards] self.assertAllEqual(expected, result)