| 32 | |
| 33 | |
| 34 | def test_pyarrow(ray_start_regular_shared): |
| 35 | ds = ray.data.range(5) |
| 36 | assert ds.map(lambda x: {"b": x["id"] + 2}).take() == [ |
| 37 | {"b": 2}, |
| 38 | {"b": 3}, |
| 39 | {"b": 4}, |
| 40 | {"b": 5}, |
| 41 | {"b": 6}, |
| 42 | ] |
| 43 | assert ds.map(lambda x: {"b": x["id"] + 2}).filter( |
| 44 | lambda x: x["b"] % 2 == 0 |
| 45 | ).take() == [{"b": 2}, {"b": 4}, {"b": 6}] |
| 46 | assert ds.filter(lambda x: x["id"] == 0).flat_map( |
| 47 | lambda x: [{"b": x["id"] + 2}, {"b": x["id"] + 20}] |
| 48 | ).take() == [{"b": 2}, {"b": 20}] |
| 49 | |
| 50 | |
| 51 | def _create_dataset(op, data): |