fix: deliver the last batch from the sharded sampler (#1993)

python/python/lance/sampler.py

@@ -404,6 +404,12 @@ def _shard_scan(
                 rows_accumulated = big_batch.num_rows
                 if big_batch.num_rows > 0:
                     accumulated_batches.append(big_batch)
+        # deliver any batches left over, they will be <= batch
+        # size but that is ok because we are done
+        last_batch = (
+            pa.Table.from_batches(accumulated_batches).combine_chunks().to_batches()[0]
+        )
+        yield last_batch
 
     def _sample_filtered(
         self,

python/python/tests/torch_tests/test_data.py

@@ -225,6 +225,9 @@ def test_sample_batches_with_filter(tmp_path: Path):
     # - - - 26 -
     # ...
     all_ids = list(chain.from_iterable([batch.cpu().numpy() for batch in ds]))
+    # Half of the data is filtered out, divided amongst 5 workers s
+    # each should see 1/10th of the data
+    assert len(all_ids) == 1000
     assert all_ids == [6 + (10 * i) for i in range(len(all_ids))]
 
     # Now test with random order