Update beam_utils_test to support numpy2

pierrot0 · The TensorFlow Datasets Authors · commit 33be83521779 · 2024-08-26T02:27:34.000-07:00
PiperOrigin-RevId: 667507374
diff --git a/tensorflow_datasets/core/beam_utils_test.py b/tensorflow_datasets/core/beam_utils_test.py
@@ -51,6 +51,9 @@ def test_read_from_tfds(
             dummy_dataset, split=split, workers_per_shard=workers_per_shard
         )
         | beam.Map(dataset_utils.as_numpy)
+        # Post numpy2, we don't get `{'id': 0}` but
+        # `{'id': np.int64(0)}`
+        | beam.Map(lambda x: {'id': int(x['id'])})
         | beam.io.WriteToText(os.fspath(tmp_path / 'out.txt'))
     )
 

Original file line number	Diff line number	Diff line change
`@@ -51,6 +51,9 @@ def test_read_from_tfds(`
`51`	`51`	`dummy_dataset, split=split, workers_per_shard=workers_per_shard`
`52`	`52`	`)`
`53`	`53`	`\| beam.Map(dataset_utils.as_numpy)`
	`54`	+ # Post numpy2, we don't get `{'id': 0}` but
	`55`	+ # `{'id': np.int64(0)}`
	`56`	`+ \| beam.Map(lambda x: {'id': int(x['id'])})`
`54`	`57`	`\| beam.io.WriteToText(os.fspath(tmp_path / 'out.txt'))`
`55`	`58`	`)`
`56`	`59`