(ray_start_regular_shared)
| 54 | |
| 55 | |
| 56 | def test_to_numpy_refs(ray_start_regular_shared): |
| 57 | # Tensor Dataset |
| 58 | ds = ray.data.range_tensor(10, override_num_blocks=2) |
| 59 | arr = np.concatenate(extract_values("data", ray.get(ds.to_numpy_refs()))) |
| 60 | np.testing.assert_equal(arr, np.expand_dims(np.arange(0, 10), 1)) |
| 61 | |
| 62 | # Table Dataset |
| 63 | ds = ray.data.range(10) |
| 64 | arr = np.concatenate([t["id"] for t in ray.get(ds.to_numpy_refs())]) |
| 65 | np.testing.assert_equal(arr, np.arange(0, 10)) |
| 66 | |
| 67 | # Test multi-column Arrow dataset. |
| 68 | ds = ray.data.from_arrow(pa.table({"a": [1, 2, 3], "b": [4, 5, 6]})) |
| 69 | arrs = ray.get(ds.to_numpy_refs()) |
| 70 | np.testing.assert_equal( |
| 71 | arrs, [{"a": np.array([1, 2, 3]), "b": np.array([4, 5, 6])}] |
| 72 | ) |
| 73 | |
| 74 | # Test multi-column Pandas dataset. |
| 75 | ds = ray.data.from_pandas(pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})) |
| 76 | arrs = ray.get(ds.to_numpy_refs()) |
| 77 | np.testing.assert_equal( |
| 78 | arrs, [{"a": np.array([1, 2, 3]), "b": np.array([4, 5, 6])}] |
| 79 | ) |
| 80 | |
| 81 | |
| 82 | def test_numpy_roundtrip(ray_start_regular_shared, tmp_path): |
nothing calls this directly
no test coverage detected
searching dependent graphs…