mirror of
https://github.com/langchain-ai/delta-rs.git
synced 2026-07-16 10:14:27 -04:00
Robert Pack
84 lines
2.3 KiB
Python
84 lines
2.3 KiB
Python
#!/usr/bin/env python3
|
|
#
|
|
# This filue contains all the tests of the deltalake python package in a
|
|
# multithreaded environment
|
|
|
|
import pathlib
|
|
import threading
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from typing import TYPE_CHECKING
|
|
|
|
import pytest
|
|
|
|
from deltalake import DeltaTable, write_deltalake
|
|
from deltalake.exceptions import CommitFailedError
|
|
|
|
if TYPE_CHECKING:
|
|
import pyarrow as pa
|
|
|
|
|
|
@pytest.mark.pyarrow
|
|
def test_concurrency(existing_table: DeltaTable, sample_data_pyarrow: "pa.Table"):
|
|
exception = None
|
|
|
|
def comp():
|
|
nonlocal exception
|
|
dt = DeltaTable(existing_table.table_uri)
|
|
for _ in range(5):
|
|
# We should always be able to get a consistent table state
|
|
data = DeltaTable(dt.table_uri).to_pyarrow_table()
|
|
# If two overwrites delete the same file and then add their own
|
|
# concurrently, then this will fail.
|
|
assert data.num_rows == sample_data_pyarrow.num_rows
|
|
try:
|
|
write_deltalake(dt.table_uri, sample_data_pyarrow, mode="overwrite")
|
|
except Exception as e:
|
|
exception = e
|
|
|
|
n_threads = 2
|
|
threads = [threading.Thread(target=comp) for _ in range(n_threads)]
|
|
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
assert isinstance(exception, CommitFailedError)
|
|
assert "a concurrent transaction deleted data this operation read" in str(exception)
|
|
|
|
|
|
@pytest.mark.polars
|
|
def test_multithreaded_write_using_table(tmp_path: pathlib.Path):
|
|
import polars as pl
|
|
|
|
table = pl.DataFrame({"a": [1, 2, 3]})
|
|
write_deltalake(tmp_path, table, mode="overwrite")
|
|
|
|
dt = DeltaTable(tmp_path)
|
|
|
|
with ThreadPoolExecutor() as exe:
|
|
list(
|
|
exe.map(
|
|
lambda i: write_deltalake(dt, pl.DataFrame({"a": [i]}), mode="append"),
|
|
range(5),
|
|
)
|
|
)
|
|
|
|
|
|
@pytest.mark.polars
|
|
def test_multithreaded_write_using_path(tmp_path: pathlib.Path):
|
|
import polars as pl
|
|
|
|
table = pl.DataFrame({"a": [1, 2, 3]})
|
|
write_deltalake(tmp_path, table, mode="overwrite")
|
|
|
|
with ThreadPoolExecutor() as exe:
|
|
list(
|
|
exe.map(
|
|
lambda _: write_deltalake(
|
|
tmp_path, pl.DataFrame({"a": [1, 2, 3]}), mode="append"
|
|
),
|
|
range(5),
|
|
)
|
|
)
|