import uuid

import pyarrow as pa

import geneva
from geneva.runners.ray.raycluster import RayCluster


# use a random version to force checkpoint to invalidate
@geneva.udf(num_cpus=1, version=uuid.uuid4().hex)
def plus_one(a: int) -> int:
    return a + 1


def test_ray_add_column_pipeline(
    geneva_test_bucket: str,
    standard_cluster: RayCluster,
) -> None:
    conn = geneva.connect(geneva_test_bucket)
    table_name = uuid.uuid4().hex
    table = conn.create_table(
        table_name,
        pa.Table.from_pydict({"a": pa.array(range(10240))}),
    )
    with standard_cluster:
        table.add_columns(
            {"b": plus_one},
            materialize=True,
            batch_size=32,
            concurrency=32,
        )

    assert table.to_arrow() == pa.Table.from_pydict(
        {"a": pa.array(range(10240)), "b": pa.array(range(1, 10241))}
    )
    conn.drop_table(table_name)
