fix(python): avoid deadlock for query-backed iterators

This commit is contained in:
Gatefixer
2026-08-05 23:33:37 +00:00
parent 7357d63e87
commit f382c8548f
3 changed files with 47 additions and 6 deletions
+11 -2
View File
@@ -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)
+12 -4
View File
@@ -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(
+24
View File
@@ -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)