mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-18 12:08:35 +00:00
fix(python): avoid deadlock for query-backed iterators
This commit is contained in:
@@ -5,6 +5,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import abstractmethod
|
||||
import asyncio
|
||||
from datetime import timedelta
|
||||
from pathlib import Path
|
||||
import sys
|
||||
@@ -1735,8 +1736,16 @@ class AsyncConnection(object):
|
||||
if fill_value is None:
|
||||
fill_value = 0.0
|
||||
|
||||
data, schema = sanitize_create_table(
|
||||
data, schema, metadata, on_bad_vectors, fill_value
|
||||
# Input preparation may advance a user-provided iterator. Keep that work
|
||||
# off the background event loop so an iterator can use the synchronous
|
||||
# LanceDB API without blocking the loop that API needs to make progress.
|
||||
data, schema = await asyncio.to_thread(
|
||||
sanitize_create_table,
|
||||
data,
|
||||
schema,
|
||||
metadata,
|
||||
on_bad_vectors,
|
||||
fill_value,
|
||||
)
|
||||
validate_schema(schema)
|
||||
|
||||
|
||||
@@ -5143,13 +5143,18 @@ class AsyncTable:
|
||||
if mode == "overwrite":
|
||||
# For overwrite, apply the same preprocessing as create_table
|
||||
# so vector columns are inferred as FixedSizeList.
|
||||
data, _ = sanitize_create_table(
|
||||
data, None, on_bad_vectors=on_bad_vectors, fill_value=fill_value
|
||||
data, _ = await asyncio.to_thread(
|
||||
sanitize_create_table,
|
||||
data,
|
||||
None,
|
||||
on_bad_vectors=on_bad_vectors,
|
||||
fill_value=fill_value,
|
||||
)
|
||||
elif on_bad_vectors != "error" or (
|
||||
schema.metadata is not None and b"embedding_functions" in schema.metadata
|
||||
):
|
||||
data = _sanitize_data(
|
||||
data = await asyncio.to_thread(
|
||||
_sanitize_data,
|
||||
data,
|
||||
schema,
|
||||
metadata=schema.metadata,
|
||||
@@ -5158,7 +5163,10 @@ class AsyncTable:
|
||||
allow_subschema=True,
|
||||
)
|
||||
_register_optional_converters()
|
||||
data = to_scannable(data)
|
||||
# Converting an iterator peeks at its first item. A synchronous query in
|
||||
# that iterator schedules work on LOOP, so peeking on LOOP's own thread
|
||||
# would deadlock waiting for itself.
|
||||
data = await asyncio.to_thread(to_scannable, data)
|
||||
progress, owns = _normalize_progress(progress)
|
||||
try:
|
||||
return await self._inner.add(
|
||||
|
||||
@@ -435,6 +435,30 @@ def test_add(mem_db: DBConnection):
|
||||
_add(table, schema)
|
||||
|
||||
|
||||
def test_add_from_iterator_that_queries_table(mem_db: DBConnection):
|
||||
source = mem_db.create_table("source", data=pa.table({"id": range(16)}))
|
||||
target = mem_db.create_table("target", schema=source.schema)
|
||||
|
||||
def batches():
|
||||
for _ in range(5):
|
||||
yield source.search().limit(10).to_arrow()
|
||||
|
||||
target.add(batches())
|
||||
|
||||
assert target.count_rows() == 50
|
||||
|
||||
|
||||
def test_create_table_from_iterator_that_queries_table(mem_db: DBConnection):
|
||||
source = mem_db.create_table("source", data=pa.table({"id": range(16)}))
|
||||
|
||||
def batches():
|
||||
yield source.search().limit(10).to_arrow()
|
||||
|
||||
target = mem_db.create_table("target", data=batches())
|
||||
|
||||
assert target.count_rows() == 10
|
||||
|
||||
|
||||
def test_add_write_parallelism(mem_db: DBConnection):
|
||||
schema = pa.schema([pa.field("id", pa.int64())])
|
||||
table = mem_db.create_table("test", schema=schema)
|
||||
|
||||
Reference in New Issue
Block a user