fix: shard batch iterator can reads partial batches (#1889)

Author: eddyxu

python/python/lance/_dataset/sharded_batch_iterator.py

@@ -141,7 +141,7 @@ def _gen_ranges():
                 total,
                 self._world_size * self._batch_size,
             ):
-                yield start, start + self._batch_size
+                yield start, min(start + self._batch_size, total)
 
         return self._ds._ds.take_scan(
             _gen_ranges(),

python/python/tests/test_dataset.py

@@ -1416,3 +1416,22 @@ def test_sharded_iterator_batches(tmp_path: Path):
             for j in range(i, i + BATCH_SIZE)
         ]
     )
+
+
+def test_sharded_iterator_non_full_batch(tmp_path: Path):
+    arr = pa.array(range(1186))
+    tbl = pa.table({"a": arr})
+
+    ds = lance.write_dataset(tbl, tmp_path)
+    shard_datast = ShardedBatchIterator(
+        ds,
+        1,
+        2,
+        columns=["a"],
+        batch_size=100,
+        granularity="batch",
+    )
+    batches = pa.concat_arrays([b["a"] for b in shard_datast])
+
+    # Can read partial batches
+    assert len(set(range(1100, 1186)) - set(batches.to_pylist())) == 0