rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2025-12-27 23:12:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(python): merge_insert in async Python (#1707)

Browse Source 
Fixes #1401
This commit is contained in:
Will Jones
2024-10-01 10:06:52 -07:00
committed by GitHub
parent 33b402c861
commit 2c4b07eb17
4 changed files with 139 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

81
python/python/tests/test_table.py
View File

@@ -636,11 +636,13 @@ def test_merge_insert(db):
new_data = pa.table({"a": [2, 4], "b": ["x", "z"]})
# replace-range
table.merge_insert(
"a"
).when_matched_update_all().when_not_matched_insert_all().when_not_matched_by_source_delete(
"a > 2"
).execute(new_data)
(
table.merge_insert("a")
.when_matched_update_all()
.when_not_matched_insert_all()
.when_not_matched_by_source_delete("a > 2")
.execute(new_data)
)
expected = pa.table({"a": [1, 2, 4], "b": ["a", "x", "z"]})
assert table.to_arrow().sort_by("a") == expected
@@ -658,6 +660,75 @@ def test_merge_insert(db):
assert table.to_arrow().sort_by("a") == expected
@pytest.mark.asyncio
async def test_merge_insert_async(db_async: AsyncConnection):
data = pa.table({"a": [1, 2, 3], "b": ["a", "b", "c"]})
table = await db_async.create_table("some_table", data=data)
assert await table.count_rows() == 3
version = await table.version()
new_data = pa.table({"a": [2, 3, 4], "b": ["x", "y", "z"]})
# upsert
await (
table.merge_insert("a")
.when_matched_update_all()
.when_not_matched_insert_all()
.execute(new_data)
)
expected = pa.table({"a": [1, 2, 3, 4], "b": ["a", "x", "y", "z"]})
assert (await table.to_arrow()).sort_by("a") == expected
await table.checkout(version)
await table.restore()
# conditional update
await (
table.merge_insert("a")
.when_matched_update_all(where="target.b = 'b'")
.execute(new_data)
)
expected = pa.table({"a": [1, 2, 3], "b": ["a", "x", "c"]})
assert (await table.to_arrow()).sort_by("a") == expected
await table.checkout(version)
await table.restore()
# insert-if-not-exists
await table.merge_insert("a").when_not_matched_insert_all().execute(new_data)
expected = pa.table({"a": [1, 2, 3, 4], "b": ["a", "b", "c", "z"]})
assert (await table.to_arrow()).sort_by("a") == expected
await table.checkout(version)
await table.restore()
# replace-range
new_data = pa.table({"a": [2, 4], "b": ["x", "z"]})
await (
table.merge_insert("a")
.when_matched_update_all()
.when_not_matched_insert_all()
.when_not_matched_by_source_delete("a > 2")
.execute(new_data)
)
expected = pa.table({"a": [1, 2, 4], "b": ["a", "x", "z"]})
assert (await table.to_arrow()).sort_by("a") == expected
await table.checkout(version)
await table.restore()
# replace-range no condition
await (
table.merge_insert("a")
.when_matched_update_all()
.when_not_matched_insert_all()
.when_not_matched_by_source_delete()
.execute(new_data)
)
expected = pa.table({"a": [2, 4], "b": ["x", "z"]})
assert (await table.to_arrow()).sort_by("a") == expected
def test_create_with_embedding_function(db):
class MyTable(LanceModel):
text: str