fix(python): preserve logical rows from Hugging Face datasets (#4461)

Hugging Face `select`, `filter`, and `shuffle` can retain a physical
Arrow table while storing logical row selection/order in an indices
table. Ingesting the physical table then silently writes excluded rows
or the wrong order and can lose duplicate selections.

Read indexed datasets through bounded Arrow-format iteration, and keep
the existing direct batch path for datasets without indices. Apply the
same conversion to sync/async table creation and additions, including
DatasetDict splits. The conversion does not mutate the user's formatting
settings.

Regression coverage includes reordered/duplicate selections, filtering,
shuffling, empty selections, projected numpy/pandas/Arrow formatting,
DatasetDicts, rescanning, and bounded batches.

Validation on base d864e7198f with its
byte-identical native/config:
- Python 3.11 full CI test selection: 1564 passed, 36 skipped, 68
deselected; 162.51s pytest /171.68s wall.
- Python 3.13 full CI test selection: 1521 passed, 37 skipped, 68
deselected; 160.32s pytest /163.74s wall. Optional HF tests were skipped
in this environment.
- Python 3.13 complete doctest command: 71 passed, 6 skipped; 8.42s
pytest /12.14s wall.
- Python 3.11 HF regression module from a candidate wheel: 43 passed;
the same regression module on the unchanged baseline gives 24 failed /19
passed.
- Ruff check/format, repository pre-commit hooks, pyright with the
source import path, and `git diff --check` passed.

This is a Python-only change. Native binaries were reused from the
verified baseline; no new native build or external-service integration
tests were run. The separate schema-metadata preservation change in
#4451 addresses a different conversion defect.

The target branch currently includes newer native changes than the
verified base. Latest-main native compatibility, Linux CI, and
remote-service tests remain for CI/maintainer validation. Prepared with
AI assistance (gpt-6.1-sol).

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()
                         )

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(

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