wehub-resource-sync
391 lines
12 KiB
Python
391 lines
12 KiB
Python
import os
|
|
|
|
import lance
|
|
import pyarrow as pa
|
|
import pytest
|
|
from packaging.version import Version
|
|
from pytest_lazy_fixtures import lf as lazy_fixture
|
|
|
|
import ray
|
|
from ray._common.test_utils import wait_for_condition
|
|
from ray.data import Schema
|
|
from ray.data._internal.datasource.lance_datasink import (
|
|
_WRITE_LANCE_FRAGMENTS_DESCRIPTION,
|
|
LanceDatasink,
|
|
_write_fragment,
|
|
)
|
|
from ray.data.datasource import SaveMode
|
|
from ray.data.datasource.path_util import _unwrap_protocol
|
|
|
|
# Skip tests for older pylance versions (<=0.3.19) due to incompatible lance API changes with pyarrow v9.0.0
|
|
pytestmark = pytest.mark.skipif(
|
|
Version(lance.__version__) <= Version("0.3.19"),
|
|
reason=f"pylance {lance.__version__} <= 0.3.19; API incompatible",
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"fs,data_path",
|
|
[
|
|
(None, lazy_fixture("local_path")),
|
|
(lazy_fixture("local_fs"), lazy_fixture("local_path")),
|
|
(lazy_fixture("s3_fs"), lazy_fixture("s3_path")),
|
|
(
|
|
lazy_fixture("s3_fs_with_space"),
|
|
lazy_fixture("s3_path_with_space"),
|
|
), # Path contains space.
|
|
(
|
|
lazy_fixture("s3_fs_with_anonymous_crendential"),
|
|
lazy_fixture("s3_path_with_anonymous_crendential"),
|
|
),
|
|
],
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"batch_size",
|
|
[None, 100],
|
|
)
|
|
def test_lance_read_basic(fs, data_path, batch_size):
|
|
df1 = pa.table({"one": [2, 1, 3, 4, 6, 5], "two": ["b", "a", "c", "e", "g", "f"]})
|
|
setup_data_path = _unwrap_protocol(data_path)
|
|
path = os.path.join(setup_data_path, "test.lance")
|
|
lance.write_dataset(df1, path)
|
|
|
|
ds_lance = lance.dataset(path)
|
|
assert ds_lance is not None
|
|
df2 = pa.table(
|
|
{
|
|
"one": [1, 2, 3, 4, 5, 6],
|
|
"three": [4, 5, 8, 9, 12, 13],
|
|
"four": ["u", "v", "w", "x", "y", "z"],
|
|
}
|
|
)
|
|
ds_lance.merge(df2, "one")
|
|
|
|
if batch_size is None:
|
|
ds = ray.data.read_lance(path)
|
|
else:
|
|
ds = ray.data.read_lance(path, scanner_options={"batch_size": batch_size})
|
|
|
|
# Test metadata-only ops.
|
|
assert ds.count() == 6
|
|
assert ds.schema() == Schema(
|
|
pa.schema(
|
|
{
|
|
"one": pa.int64(),
|
|
"two": pa.string(),
|
|
"three": pa.int64(),
|
|
"four": pa.string(),
|
|
}
|
|
)
|
|
)
|
|
|
|
# Test read.
|
|
values = [[s["one"], s["two"]] for s in ds.take_all()]
|
|
assert sorted(values) == [
|
|
[1, "a"],
|
|
[2, "b"],
|
|
[3, "c"],
|
|
[4, "e"],
|
|
[5, "f"],
|
|
[6, "g"],
|
|
]
|
|
|
|
# Test column projection.
|
|
ds = ray.data.read_lance(path, columns=["one"])
|
|
values = [s["one"] for s in ds.take_all()]
|
|
assert sorted(values) == [1, 2, 3, 4, 5, 6]
|
|
assert ds.schema().names == ["one"]
|
|
|
|
|
|
@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
|
|
def test_lance_read_with_scanner_fragments(data_path):
|
|
table = pa.table({"one": [2, 1, 3, 4, 6, 5], "two": ["b", "a", "c", "e", "g", "f"]})
|
|
setup_data_path = _unwrap_protocol(data_path)
|
|
path = os.path.join(setup_data_path, "test.lance")
|
|
dataset = lance.write_dataset(table, path, max_rows_per_file=2)
|
|
assert dataset is not None
|
|
|
|
fragments = dataset.get_fragments()
|
|
ds = ray.data.read_lance(path, scanner_options={"fragments": fragments[:1]})
|
|
values = [[s["one"], s["two"]] for s in ds.take_all()]
|
|
assert values == [
|
|
[2, "b"],
|
|
[1, "a"],
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
|
|
def test_lance_read_many_files(data_path):
|
|
setup_data_path = _unwrap_protocol(data_path)
|
|
path = os.path.join(setup_data_path, "test.lance")
|
|
num_rows = 1024
|
|
data = pa.table({"id": pa.array(range(num_rows))})
|
|
lance.write_dataset(data, path, max_rows_per_file=1)
|
|
|
|
def test_lance():
|
|
ds = ray.data.read_lance(path)
|
|
return ds.count() == num_rows
|
|
|
|
wait_for_condition(test_lance, timeout=10)
|
|
|
|
|
|
@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
|
|
def test_lance_write(data_path):
|
|
schema = pa.schema([pa.field("id", pa.int64()), pa.field("str", pa.string())])
|
|
|
|
ray.data.range(10).map(
|
|
lambda x: {"id": x["id"], "str": f"str-{x['id']}"}
|
|
).write_lance(data_path, schema=schema)
|
|
|
|
ds = lance.dataset(data_path)
|
|
assert ds is not None
|
|
ds.count_rows() == 10
|
|
assert ds.schema.names == schema.names
|
|
# The schema is platform-dependent, because numpy uses int32 on Windows.
|
|
# So we observe the schema that is written and use that.
|
|
schema = ds.schema
|
|
|
|
tbl = ds.to_table()
|
|
assert sorted(tbl["id"].to_pylist()) == list(range(10))
|
|
assert set(tbl["str"].to_pylist()) == {f"str-{i}" for i in range(10)}
|
|
|
|
ray.data.range(10).map(
|
|
lambda x: {"id": x["id"] + 10, "str": f"str-{x['id'] + 10}"}
|
|
).write_lance(data_path, mode=SaveMode.APPEND)
|
|
|
|
ds = lance.dataset(data_path)
|
|
assert ds is not None
|
|
ds.count_rows() == 20
|
|
tbl = ds.to_table()
|
|
assert sorted(tbl["id"].to_pylist()) == list(range(20))
|
|
assert set(tbl["str"].to_pylist()) == {f"str-{i}" for i in range(20)}
|
|
|
|
ray.data.range(10).map(
|
|
lambda x: {"id": x["id"], "str": f"str-{x['id']}"}
|
|
).write_lance(data_path, schema=schema, mode=SaveMode.OVERWRITE)
|
|
|
|
ds = lance.dataset(data_path)
|
|
assert ds is not None
|
|
ds.count_rows() == 10
|
|
assert ds.schema == schema
|
|
|
|
|
|
@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
|
|
def test_lance_write_create_errors_if_exists(data_path):
|
|
table_path = os.path.join(data_path, "my_table")
|
|
ds = ray.data.range(10)
|
|
|
|
# First CREATE succeeds on an empty destination.
|
|
ds.write_lance(table_path, mode=SaveMode.CREATE)
|
|
assert lance.dataset(table_path).count_rows() == 10
|
|
|
|
# A second CREATE must error instead of silently overwriting.
|
|
with pytest.raises(ValueError, match="already exists"):
|
|
ray.data.range(5).write_lance(table_path, mode=SaveMode.CREATE)
|
|
|
|
# Existing data is untouched.
|
|
assert lance.dataset(table_path).count_rows() == 10
|
|
|
|
# CREATE is also the default mode, so it must guard too.
|
|
with pytest.raises(ValueError, match="already exists"):
|
|
ray.data.range(5).write_lance(table_path)
|
|
assert lance.dataset(table_path).count_rows() == 10
|
|
|
|
# OVERWRITE replaces the existing data.
|
|
ray.data.range(5).write_lance(table_path, mode=SaveMode.OVERWRITE)
|
|
assert lance.dataset(table_path).count_rows() == 5
|
|
|
|
|
|
@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
|
|
def test_lance_write_append_errors_if_missing(data_path):
|
|
table_path = os.path.join(data_path, "missing_table")
|
|
# APPEND surfaces Lance's own "not found" error. We don't pin the message,
|
|
# since it can change across Lance versions.
|
|
expected_errors: tuple[type[Exception], ...] = (
|
|
ValueError,
|
|
OSError,
|
|
FileNotFoundError,
|
|
)
|
|
with pytest.raises(expected_errors):
|
|
ray.data.range(5).write_lance(table_path, mode=SaveMode.APPEND)
|
|
assert not os.path.exists(table_path)
|
|
|
|
|
|
@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
|
|
def test_lance_write_min_rows_per_file(data_path):
|
|
schema = pa.schema([pa.field("id", pa.int64()), pa.field("str", pa.string())])
|
|
|
|
ray.data.range(10).map(
|
|
lambda x: {"id": x["id"], "str": f"str-{x['id']}"}
|
|
).write_lance(data_path, schema=schema, min_rows_per_file=100)
|
|
|
|
ds = lance.dataset(data_path)
|
|
assert ds is not None
|
|
assert ds.count_rows() == 10
|
|
assert ds.schema == schema
|
|
|
|
assert len(ds.get_fragments()) == 1
|
|
|
|
|
|
@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
|
|
def test_lance_write_max_rows_per_file(data_path):
|
|
schema = pa.schema([pa.field("id", pa.int64()), pa.field("str", pa.string())])
|
|
|
|
ray.data.range(10).map(
|
|
lambda x: {"id": x["id"], "str": f"str-{x['id']}"}
|
|
).write_lance(data_path, schema=schema, max_rows_per_file=1)
|
|
|
|
ds = lance.dataset(data_path)
|
|
assert ds is not None
|
|
assert ds.count_rows() == 10
|
|
assert ds.schema == schema
|
|
|
|
assert len(ds.get_fragments()) == 10
|
|
|
|
|
|
@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
|
|
def test_lance_read_with_version(data_path):
|
|
# Write an initial dataset (version 1)
|
|
df1 = pa.table({"one": [2, 1, 3, 4, 6, 5], "two": ["b", "a", "c", "e", "g", "f"]})
|
|
setup_data_path = _unwrap_protocol(data_path)
|
|
path = os.path.join(setup_data_path, "test_version.lance")
|
|
lance.write_dataset(df1, path)
|
|
|
|
# Merge new data to create a later version (latest)
|
|
ds_lance = lance.dataset(path)
|
|
assert ds_lance is not None
|
|
# Get the initial version
|
|
initial_version = ds_lance.version
|
|
|
|
df2 = pa.table(
|
|
{
|
|
"one": [1, 2, 3, 4, 5, 6],
|
|
"three": [4, 5, 8, 9, 12, 13],
|
|
"four": ["u", "v", "w", "x", "y", "z"],
|
|
}
|
|
)
|
|
ds_lance.merge(df2, "one")
|
|
|
|
# Default read should return the latest (merged) dataset.
|
|
ds_latest = ray.data.read_lance(path)
|
|
|
|
assert ds_latest.count() == 6
|
|
# Latest dataset should contain merged columns
|
|
assert "three" in ds_latest.schema().names
|
|
|
|
# Read the initial version and ensure it contains the original columns
|
|
ds_prev = ray.data.read_lance(path, version=initial_version)
|
|
assert ds_prev.count() == 6
|
|
assert ds_prev.schema().names == ["one", "two"]
|
|
|
|
values_prev = [[s["one"], s["two"]] for s in ds_prev.take_all()]
|
|
assert sorted(values_prev) == [
|
|
[1, "a"],
|
|
[2, "b"],
|
|
[3, "c"],
|
|
[4, "e"],
|
|
[5, "f"],
|
|
[6, "g"],
|
|
]
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_lance_write(monkeypatch):
|
|
captured = {}
|
|
|
|
class _FakeLanceDatasink:
|
|
def __init__(self, path, **kwargs):
|
|
captured["path"] = path
|
|
captured["kwargs"] = kwargs
|
|
|
|
def _fake_write_datasink(self, datasink, **kwargs):
|
|
captured["datasink"] = datasink
|
|
captured["write_kwargs"] = kwargs
|
|
|
|
monkeypatch.setattr(ray.data.dataset, "LanceDatasink", _FakeLanceDatasink)
|
|
monkeypatch.setattr(ray.data.Dataset, "write_datasink", _fake_write_datasink)
|
|
|
|
return captured, _FakeLanceDatasink
|
|
|
|
|
|
def test_write_lance_passes_namespace_args(mock_lance_write):
|
|
captured, fake_lance_datasink_cls = mock_lance_write
|
|
table_id = ["db", "table"]
|
|
namespace_impl = "dir"
|
|
namespace_properties = {"path": "/tmp/ns"}
|
|
|
|
ds = ray.data.range(1)
|
|
ds.write_lance(
|
|
"/tmp/lance-namespace-test",
|
|
table_id=table_id,
|
|
namespace_impl=namespace_impl,
|
|
namespace_properties=namespace_properties,
|
|
)
|
|
|
|
assert captured["path"] == "/tmp/lance-namespace-test"
|
|
assert captured["kwargs"]["table_id"] == table_id
|
|
assert captured["kwargs"]["namespace_impl"] == namespace_impl
|
|
assert captured["kwargs"]["namespace_properties"] == namespace_properties
|
|
assert isinstance(captured["datasink"], fake_lance_datasink_cls)
|
|
|
|
|
|
@pytest.mark.parametrize("mode", [SaveMode.APPEND, SaveMode.OVERWRITE])
|
|
def test_lance_namespace_write_rejects_non_create_mode(monkeypatch, mode):
|
|
class _FakeNamespace:
|
|
pass
|
|
|
|
monkeypatch.setattr(
|
|
"ray.data._internal.datasource.lance_datasink.get_or_create_namespace",
|
|
lambda namespace_impl, namespace_properties: _FakeNamespace(),
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="Namespace writes currently only support"):
|
|
LanceDatasink(
|
|
uri="/tmp/lance-namespace-test",
|
|
mode=mode,
|
|
table_id=["db", "table"],
|
|
namespace_impl="dir",
|
|
namespace_properties={"path": "/tmp/ns"},
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"max_attempts,expected_blocks_consumed_before_write",
|
|
[(1, 1), (2, 3)],
|
|
)
|
|
def test_write_fragment_only_materializes_stream_when_retrying(
|
|
monkeypatch, max_attempts, expected_blocks_consumed_before_write
|
|
):
|
|
import lance.fragment
|
|
|
|
consumed = {"count": 0}
|
|
blocks = [pa.table({"id": [i]}) for i in range(3)]
|
|
|
|
def block_stream():
|
|
for block in blocks:
|
|
consumed["count"] += 1
|
|
yield block
|
|
|
|
def fake_write_fragments(reader, uri, **kwargs):
|
|
assert consumed["count"] == expected_blocks_consumed_before_write
|
|
return []
|
|
|
|
monkeypatch.setattr(lance.fragment, "write_fragments", fake_write_fragments)
|
|
|
|
_write_fragment(
|
|
block_stream(),
|
|
"/tmp/lance-materialization-test",
|
|
retry_params={
|
|
"description": _WRITE_LANCE_FRAGMENTS_DESCRIPTION,
|
|
"match": [],
|
|
"max_attempts": max_attempts,
|
|
"max_backoff_s": 0,
|
|
},
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|