Reflect samples argument in SpectrumDataset lengths (#49)

* Fix length of dataset

* Bump changelog

* Add condition when samples is 0

CHANGELOG.md

@@ -6,6 +6,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [v0.4.3]
+
+### Changed
+- Length of the `SpectrumDataset` and `AnnotatedSpectrumDataset` now reflect the `samples` parameter of the `lance.pytorch.LanceDataset` parent class.
+
 ## [v0.4.2]
 
 ### Changed

depthcharge/data/spectrum_datasets.py

@@ -174,7 +174,11 @@ def __getitem__(self, idx: int) -> dict[str, Any]:
 
     def __len__(self) -> int:
         """The number of batches in the lance dataset."""
-        return math.ceil(self._dataset.count_rows() / self.batch_size)
+        num = self._dataset.count_rows()
+        if self.samples:
+            num = min(self.samples, num)
+
+        return math.ceil(num / self.batch_size)
 
     def __del__(self) -> None:
         """Cleanup the temporary directory."""

tests/unit_tests/test_data/test_loaders.py

@@ -1,6 +1,7 @@
 """Test PyTorch DataLoaders."""
 
 import pyarrow as pa
+import pytest
 import torch
 from torch.utils.data import DataLoader
 
@@ -25,6 +26,21 @@ def test_spectrum_loader(mgf_small, tmp_path):
     assert isinstance(batch["mz_array"], torch.Tensor)
 
 
+@pytest.mark.parametrize(["samples", "batches"], [(1, 1), (50, 2), (4, 2)])
+def test_spectrum_loader_samples(mgf_small, tmp_path, samples, batches):
+    """Test sampling."""
+    dset = SpectrumDataset(
+        mgf_small,
+        batch_size=1,
+        path=tmp_path / "test",
+        samples=samples,
+    )
+
+    loaded = list(DataLoader(dset))
+    assert len(dset) == batches
+    assert len(loaded) == batches
+
+
 def test_streaming_spectrum_loader(mgf_small, tmp_path):
     """Test streaming spectra."""
     streamer = StreamingSpectrumDataset(mgf_small, batch_size=2)