diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index eeae8bf50..718979670 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -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) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index ae36bac7a..d6ef2ee74 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -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( diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 069527b21..df66e0177 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -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)