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)
Пример #2
0
 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)