Log the number of examples assigned to each shard in BeamWriter

tomvdw · The TensorFlow Datasets Authors · commit 4aa203e60848 · 2024-12-09T01:44:41.000-08:00
This helps debugging whether the data is distributed correctly

PiperOrigin-RevId: 704190477
diff --git a/tensorflow_datasets/core/writer.py b/tensorflow_datasets/core/writer.py
@@ -612,6 +612,7 @@ def _assign_shard(
         hkey=key, num_buckets=num_shards, max_hkey=largest_key
     )
     self._get_distribution(name="ShardDistribution").update(shard_number)
+    self.inc_counter(f"{self._filename_template.split}.shard_{shard_number}")
     return (shard_number, key_serialized_example)
 
   def _store_split_info(

Original file line number	Diff line number	Diff line change
`@@ -612,6 +612,7 @@ def _assign_shard(`
`612`	`612`	`hkey=key, num_buckets=num_shards, max_hkey=largest_key`
`613`	`613`	`)`
`614`	`614`	`self._get_distribution(name="ShardDistribution").update(shard_number)`
	`615`	`+ self.inc_counter(f"{self._filename_template.split}.shard_{shard_number}")`
`615`	`616`	`return (shard_number, key_serialized_example)`
`616`	`617`
`617`	`618`	`def _store_split_info(`