-
Notifications
You must be signed in to change notification settings - Fork 199
Expand file tree
/
Copy pathtest_datasource.py
More file actions
61 lines (42 loc) · 2.52 KB
/
Copy pathtest_datasource.py
File metadata and controls
61 lines (42 loc) · 2.52 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright the Vortex contributors
import pyarrow as pa
import pytest
import ray
from ray.data import read_datasource # pyright: ignore[reportUnknownVariableType]
import vortex as vx
from vortex.ray.datasource import VortexDatasource, partition
from .test_file import record
@pytest.fixture(scope="module")
def ray_init():
# Ray's uv_runtime_env_hook would auto-upload the working directory to
# workers, but vortex-python's compiled _lib extension exceeds Ray's
# 512 MiB upload limit. Disable the hook for these local-mode tests.
# (Ray 2.55 added a string-type validation that broke the previous
# `working_dir: None` workaround from ray-project/ray#53848.)
import ray._private.ray_constants as ray_constants
ray_constants.RAY_ENABLE_UV_RUN_RUNTIME_ENV = False
_ = ray.init() # pyright: ignore[reportUnknownMemberType]
yield None
ray.shutdown() # pyright: ignore[reportUnknownMemberType]
def test_partition():
assert partition(1, []) == [[]]
assert partition(1, [1]) == [[1]]
assert partition(1, [1, 2, 3]) == [[1, 2, 3]]
assert partition(2, [1, 2, 3]) == [[1, 2], [3]]
assert partition(3, [1, 2, 3]) == [[1], [2], [3]]
assert partition(2, list(range(9))) == [[0, 1, 2, 3, 4], [5, 6, 7, 8]]
assert partition(3, list(range(9))) == [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
assert partition(3, list(range(11))) == [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10]]
def test_vortex_datasource(ray_init, tmpdir_factory): # pyright: ignore[reportUnknownParameterType, reportMissingParameterType, reportUnusedParameter]
folder = tmpdir_factory.mktemp("data") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
arr1 = vx.array([record(x) for x in range(5)])
vx.io.write(arr1, str(folder / "01.vortex")) # pyright: ignore[reportUnknownArgumentType]
arr2 = vx.array([record(x) for x in range(5, 10)])
vx.io.write(arr2, str(folder / "02.vortex")) # pyright: ignore[reportUnknownArgumentType]
ds = read_datasource(VortexDatasource(url=str(folder))) # pyright: ignore[reportUnknownArgumentType]
# Without an explicit sort, Ray may reorder rows *even within a single record batch*.
ds = ds.sort("index")
tbl = pa.concat_tables(pa.Table.from_pydict(x) for x in ds.iter_batches()) # pyright: ignore[reportArgumentType, reportUnknownMemberType, reportUnknownVariableType]
expected = pa.Table.from_pylist([record(x) for x in range(0, 10)], schema=tbl.schema)
assert tbl == expected