diff --git a/python/python/lancedb/scannable.py b/python/python/lancedb/scannable.py index baf8e0c2b..9a6f71a5d 100644 --- a/python/python/lancedb/scannable.py +++ b/python/python/lancedb/scannable.py @@ -276,6 +276,14 @@ def _from_iterable(data: Iterator) -> Scannable: _registered_modules: set[str] = set() +def _hf_batches(data): + if data._indices is None: + yield from data.data.to_batches() + else: + for table in data.with_format("arrow").iter(batch_size=1000): + yield from table.to_batches() + + def _register_optional_converters(): """Register converters for optional dependencies that are already imported.""" @@ -323,7 +331,15 @@ def _register_optional_converters(): def _from_hf_dataset(data: HFDataset) -> Scannable: table = data.data.table # Access underlying Arrow table return Scannable( - schema=table.schema, num_rows=len(data), reader=table.to_reader + schema=table.schema, + num_rows=len(data), + reader=( + table.to_reader + if data._indices is None + else lambda: pa.RecordBatchReader.from_batches( + table.schema, _hf_batches(data) + ) + ), ) @to_scannable.register(HFDatasetDict) @@ -335,7 +351,7 @@ def _register_optional_converters(): def gen(): for split_name, dataset in data.items(): - for batch in dataset.data.to_batches(): + for batch in _hf_batches(dataset): split_arr = pa.array( [split_name] * len(batch), type=pa.string() ) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 9b96a1b2a..f9cfac6e1 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -28,7 +28,7 @@ from typing import ( ) from urllib.parse import urlparse -from lancedb.scannable import _register_optional_converters, to_scannable +from lancedb.scannable import _hf_batches, _register_optional_converters, to_scannable from . import __version__ from ._blob import ( @@ -400,7 +400,7 @@ def _into_pyarrow_reader( if _check_for_hugging_face(data): if isinstance(data, datasets.Dataset): schema = data.features.arrow_schema - return pa.RecordBatchReader.from_batches(schema, data.data.to_batches()) + return pa.RecordBatchReader.from_batches(schema, _hf_batches(data)) elif isinstance(data, datasets.dataset_dict.DatasetDict): schema = _schema_from_hf(data, None) if "split" not in schema.names: @@ -1195,7 +1195,7 @@ def _to_batches_with_split(data): with an extra `split` column """ for key, dataset in data.items(): - for batch in dataset.data.to_batches(): + for batch in _hf_batches(dataset): table = pa.Table.from_batches([batch]) if "split" not in table.column_names: table = table.append_column( diff --git a/python/python/tests/test_huggingface.py b/python/python/tests/test_huggingface.py index 9b597bbc9..91421b79c 100644 --- a/python/python/tests/test_huggingface.py +++ b/python/python/tests/test_huggingface.py @@ -129,3 +129,106 @@ def test_generator(tmp_path: Path): assert len(tbl) == 2 assert tbl.schema == ds.features.arrow_schema + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("operation", ["create", "add"]) +@pytest.mark.parametrize( + "selection", + ["normal", "contiguous", "selected", "duplicates", "filtered", "shuffled", "empty"], +) +async def test_write_selected_hf_dataset(tmp_path, asynchronous, operation, selection): + source = datasets.Dataset.from_dict({"id": list(range(5)), "text": list("abcde")}) + data = { + "normal": source, + "contiguous": source.select(range(1, 4)), + "selected": source.select([4, 1]), + "duplicates": source.select([2, 2, 0]), + "filtered": source.filter(lambda row: row["id"] % 2 == 1), + "shuffled": source.shuffle(seed=42), + "empty": source.select([]), + }[selection] + expected = data.to_list() + if asynchronous: + db = await lancedb.connect_async(tmp_path) + if operation == "create": + table = await db.create_table("selected", data) + else: + table = await db.create_table("selected", schema=data.features.arrow_schema) + await table.add(data) + actual = await table.to_arrow() + else: + db = lancedb.connect(tmp_path) + if operation == "create": + table = db.create_table("selected", data) + else: + table = db.create_table("selected", schema=data.features.arrow_schema) + table.add(data) + actual = table.to_arrow() + assert actual.to_pylist() == expected + assert actual.schema.equals(data.features.arrow_schema, check_metadata=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("operation", ["create", "add"]) +async def test_write_selected_hf_splits(tmp_path, asynchronous, operation): + source = datasets.Dataset.from_dict({"id": list(range(5)), "text": list("abcde")}) + data = datasets.DatasetDict( + { + "train": source.select([4, 1]), + "test": source.select([3, 0, 2]), + "empty": source.select([]), + } + ) + expected = [ + dict(row, split=split) + for split, dataset in data.items() + for row in dataset.to_list() + ] + schema = source.features.arrow_schema.append(pa.field("split", pa.string())) + if asynchronous: + db = await lancedb.connect_async(tmp_path) + if operation == "create": + table = await db.create_table("splits", data) + else: + table = await db.create_table("splits", schema=schema) + await table.add(data) + actual = await table.to_arrow() + else: + db = lancedb.connect(tmp_path) + if operation == "create": + table = db.create_table("splits", data) + else: + table = db.create_table("splits", schema=schema) + table.add(data) + actual = table.to_arrow() + assert actual.to_pylist() == expected + assert actual.schema.equals(schema, check_metadata=True) + + +@pytest.mark.parametrize("format_type", [None, "numpy", "pandas", "arrow"]) +@pytest.mark.parametrize("empty", [False, True]) +def test_selected_hf_reader_is_rescannable(format_type, empty): + from lancedb.scannable import _register_optional_converters, to_scannable + + source = datasets.Dataset.from_dict( + {"id": list(range(2105)), "text": [str(i) for i in range(2105)]} + ) + indices = [] if empty else list(reversed(range(2105))) + [1, 1] + data = source.select(indices).with_format(format_type, columns=["id"]) + original_format = data.format + expected = source.select(indices).to_list() + _register_optional_converters() + scannable = to_scannable(data) + assert scannable.rescannable + assert scannable.num_rows == len(expected) + for _ in range(2): + batches = list(scannable.reader()) + if not empty: + assert all(batch.num_rows <= 1000 for batch in batches) + actual = pa.Table.from_batches(batches, schema=scannable.schema) + assert actual.to_pylist() == expected + assert actual.schema.equals(data.features.arrow_schema, check_metadata=True) + assert data.format == original_format