Function test_numpy_write

python/ray/data/tests/datasource/test_numpy.py:109–117

(ray_start_regular_shared, tmp_path)

def test_numpy_write(ray_start_regular_shared, tmp_path):
    ds = ray.data.range_tensor(1)

    ds.write_numpy(tmp_path, column="data")

    actual_array = np.concatenate(
        [np.load(os.path.join(tmp_path, filename)) for filename in os.listdir(tmp_path)]
    )
    assert actual_array == np.array((0,))


if __name__ == "__main__":