Function test_numpy_roundtrip

python/ray/data/tests/datasource/test_numpy.py:82–92

(ray_start_regular_shared, tmp_path)

def test_numpy_roundtrip(ray_start_regular_shared, tmp_path):
    tensor_type = DataContext.get_current().arrow_fixed_shape_tensor_format.to_type()

    ds = ray.data.range_tensor(10, override_num_blocks=2)
    ds.write_numpy(tmp_path, column="data")
    ds = ray.data.read_numpy(tmp_path)
    assert ds.count() == 10
    assert ds.schema() == Schema(pa.schema([("data", tensor_type((1,), pa.int64()))]))
    assert sorted(ds.take_all(), key=lambda row: row["data"]) == [
        {"data": np.array([i])} for i in range(10)
    ]


def test_numpy_read_x(ray_start_regular_shared, tmp_path):