def test_transform_function_with_predicate(synthetic_dataset, reader_factory):
    """Make sure we apply transform only after we apply the predicate"""

    with reader_factory(synthetic_dataset.url, schema_fields=[TestSchema.id, TestSchema.id2],
                        predicate=in_lambda(['id2'], lambda id2: id2 == 1),
                        transform_spec=TransformSpec(removed_fields=['id2'])) as reader:
        rows = list(reader)
        assert 'id2' not in rows[0]._fields
        actual_ids = np.asarray(list(row.id for row in rows))
        assert actual_ids.size > 0
        # In the test data id2 = id % 2, which means we expect only odd ids to remain after
        # we apply lambda id2: id2 == 1 predicate.
        assert np.all(actual_ids % 2 == 1)
Make sure we apply transform only after we apply the predicate
petastorm/tests/test_end_to_end.py