def test_with_tf_data_api(synthetic_dataset):
    """Verify that WeightedSamplingReader is compatible with make_petastorm_dataset"""

    np.random.seed(42)

    fields_to_read = ['id.*', 'image_png']

    # Use cur_shard=0, shard_count=2 to get only half samples from the second reader.
    readers = [make_reader(synthetic_dataset.url, schema_fields=fields_to_read, workers_count=1),
               make_reader(synthetic_dataset.url, schema_fields=fields_to_read, workers_count=1,
                           cur_shard=0, shard_count=2)]

    with WeightedSamplingReader(readers, [0.5, 0.5]) as mixer:
        dataset = make_petastorm_dataset(mixer)
        iterator = dataset.make_one_shot_iterator()
        tensor = iterator.get_next()
        rows_count = 0
        with tf.Session() as sess:
            while True:
                try:
                    sess.run(tensor)
                    rows_count += 1
                except tf.errors.OutOfRangeError:
                    break

        # We expect iterations to finish once the second read has exhausted its samples. For each sample in the
        # second reaader we read approximately 1 sample from the first.
        expected_rows_approx = len(synthetic_dataset.data)
        np.testing.assert_allclose(rows_count, expected_rows_approx, atol=20)
Verify that WeightedSamplingReader is compatible with make_petastorm_dataset
petastorm/tests/test_weighted_sampling_reader.py