Function test_flat_map

python/ray/data/tests/test_context_propagation.py:107–114

(
    ray_start_regular_shared,
)

def test_flat_map(
    ray_start_regular_shared,
):
    _test_updating_context_after_dataset_creation(
        lambda: ray.data.range(1).flat_map(
            lambda _: [{"id": DataContext.get_current().get_config("foo")}]
        )
    )


def test_map_batches(