Add subset filtering for dataset items (#5004)

Author: A-Artemis

application/backend/app/api/routers/datasets.py

@@ -17,7 +17,7 @@
     SetDatasetItemAnnotations,
 )
 from app.core.models import Pagination
-from app.models import DatasetItemAnnotationStatus
+from app.models import DatasetItemAnnotationStatus, DatasetItemSubset
 from app.schemas import ProjectView
 from app.services import DatasetService, ResourceNotFoundError
 from app.services.dataset_service import AnnotationValidationError, InvalidImageError, SubsetAlreadyAssignedError
@@ -106,14 +106,20 @@ def list_dataset_items(  # noqa: PLR0913
     end_date: Annotated[datetime | None, Query()] = None,
     annotation_status: Annotated[DatasetItemAnnotationStatus | None, Query()] = None,
     labels: Annotated[list[UUID] | None, Query()] = None,
+    subset: Annotated[DatasetItemSubset | None, Query()] = None,
 ) -> DatasetItemsWithPagination:
     """List the available dataset items and their metadata. This endpoint supports pagination."""
     if start_date is not None and end_date is not None and start_date > end_date:
         raise HTTPException(
             status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail="Start date must be before end date."
         )
     total = dataset_service.count_dataset_items(
-        project=project, start_date=start_date, end_date=end_date, annotation_status=annotation_status, label_ids=labels
+        project=project,
+        start_date=start_date,
+        end_date=end_date,
+        annotation_status=annotation_status,
+        label_ids=labels,
+        subset=subset,
     )
     dataset_items = dataset_service.list_dataset_items(
         project=project,
@@ -123,6 +129,7 @@ def list_dataset_items(  # noqa: PLR0913
         end_date=end_date,
         annotation_status=annotation_status,
         label_ids=labels,
+        subset=subset,
     )
     return DatasetItemsWithPagination(
         items=[DatasetItemView.model_validate(dataset_item, from_attributes=True) for dataset_item in dataset_items],

application/backend/app/repositories/dataset_item_repo.py

@@ -50,6 +50,13 @@ def _apply_annotation_status_filter(stmt: Select, annotation_status: str | None
             stmt = stmt.where(DatasetItemDB.annotation_data.is_not(None), DatasetItemDB.user_reviewed.is_(False))
         return stmt
 
+    @staticmethod
+    def _apply_subset_filter(stmt: Select, subset: str | None = None) -> Select:
+        """Apply subset filter to a select statement."""
+        if subset is not None:
+            stmt = stmt.where(DatasetItemDB.subset == subset)
+        return stmt
+
     def save(self, dataset_item_db: DatasetItemDB) -> DatasetItemDB:
         dataset_item_db.updated_at = datetime.now(UTC)
         self.db.add(dataset_item_db)
@@ -62,6 +69,7 @@ def count(
         end_date: datetime | None = None,
         annotation_status: str | None = None,
         label_ids: list[str] | None = None,
+        subset: str | None = None,
     ) -> int:
         # When the query involves a JOIN (e.g. when filtering by labels), count distinct items to avoid duplicates
         if label_ids:
@@ -71,6 +79,7 @@ def count(
         stmt = select(select_fn).select_from(DatasetItemDB).where(DatasetItemDB.project_id == self.project_id)
         stmt = self._apply_date_filters(stmt, start_date, end_date)
         stmt = self._apply_annotation_status_filter(stmt, annotation_status)
+        stmt = self._apply_subset_filter(stmt, subset)
         if label_ids:
             stmt = stmt.join(DatasetItemLabelDB).where(DatasetItemLabelDB.label_id.in_(label_ids))
         return self.db.scalar(stmt) or 0
@@ -83,10 +92,12 @@ def list_items(
         end_date: datetime | None = None,
         annotation_status: str | None = None,
         label_ids: list[str] | None = None,
+        subset: str | None = None,
     ) -> list[DatasetItemDB]:
         stmt = self._base_select()
         stmt = self._apply_date_filters(stmt, start_date, end_date)
         stmt = self._apply_annotation_status_filter(stmt, annotation_status)
+        stmt = self._apply_subset_filter(stmt, subset)
         if label_ids:
             stmt = stmt.join(DatasetItemLabelDB).where(DatasetItemLabelDB.label_id.in_(label_ids)).distinct()
         stmt = stmt.order_by(DatasetItemDB.created_at.desc()).offset(offset).limit(limit)

application/backend/app/services/dataset_service.py

@@ -158,15 +158,20 @@ def count_dataset_items(
         end_date: datetime | None = None,
         annotation_status: str | None = None,
         label_ids: list[UUID] | None = None,
+        subset: str | None = None,
     ) -> int:
         """Get number of available dataset items (within date range if specified)"""
         repo = DatasetItemRepository(project_id=str(project.id), db=self._db_session)
         label_ids_str = [str(label_id) for label_id in label_ids] if label_ids else None
         return repo.count(
-            start_date=start_date, end_date=end_date, annotation_status=annotation_status, label_ids=label_ids_str
+            start_date=start_date,
+            end_date=end_date,
+            annotation_status=annotation_status,
+            label_ids=label_ids_str,
+            subset=subset,
         )
 
-    def list_dataset_items(
+    def list_dataset_items(  # noqa: PLR0913
         self,
         project: ProjectView,
         limit: int = 20,
@@ -175,6 +180,7 @@ def list_dataset_items(
         end_date: datetime | None = None,
         annotation_status: str | None = None,
         label_ids: list[UUID] | None = None,
+        subset: str | None = None,
     ) -> list[DatasetItem]:
         """Get information about available dataset items"""
         repo = DatasetItemRepository(project_id=str(project.id), db=self._db_session)
@@ -188,6 +194,7 @@ def list_dataset_items(
                 end_date=end_date,
                 annotation_status=annotation_status,
                 label_ids=label_ids_str,
+                subset=subset,
             )
         ]
 

application/backend/tests/integration/services/test_dataset_service.py

@@ -331,6 +331,122 @@ def fxt_project_with_labeled_dataset_items(
     return project, db_dataset_items
 
 
+@pytest.fixture
+def fxt_project_with_subset_items(fxt_project_with_pipeline, db_session) -> tuple[ProjectView, list[DatasetItemDB]]:
+    """Fixture with dataset items covering all subset types."""
+    project, _ = fxt_project_with_pipeline
+
+    # Unassigned items
+    unassigned_items = [
+        DatasetItemDB(
+            name="unassigned1",
+            format="jpg",
+            size=1024,
+            width=1024,
+            height=768,
+            subset=DatasetItemSubset.UNASSIGNED,
+            user_reviewed=False,
+            project_id=str(project.id),
+            created_at=datetime.fromisoformat("2025-02-01T00:00:00Z"),
+        ),
+        DatasetItemDB(
+            name="unassigned2",
+            format="jpg",
+            size=1024,
+            width=1024,
+            height=768,
+            subset=DatasetItemSubset.UNASSIGNED,
+            user_reviewed=False,
+            project_id=str(project.id),
+            created_at=datetime.fromisoformat("2025-02-02T00:00:00Z"),
+        ),
+    ]
+
+    # Training items
+    training_items = [
+        DatasetItemDB(
+            name="training1",
+            format="jpg",
+            size=1024,
+            width=1024,
+            height=768,
+            subset=DatasetItemSubset.TRAINING,
+            user_reviewed=False,
+            project_id=str(project.id),
+            created_at=datetime.fromisoformat("2025-02-03T00:00:00Z"),
+        ),
+        DatasetItemDB(
+            name="training2",
+            format="jpg",
+            size=1024,
+            width=1024,
+            height=768,
+            subset=DatasetItemSubset.TRAINING,
+            user_reviewed=False,
+            project_id=str(project.id),
+            created_at=datetime.fromisoformat("2025-02-04T00:00:00Z"),
+        ),
+        DatasetItemDB(
+            name="training3",
+            format="jpg",
+            size=1024,
+            width=1024,
+            height=768,
+            subset=DatasetItemSubset.TRAINING,
+            user_reviewed=False,
+            project_id=str(project.id),
+            created_at=datetime.fromisoformat("2025-02-05T00:00:00Z"),
+        ),
+    ]
+
+    # Validation items
+    validation_items = [
+        DatasetItemDB(
+            name="validation1",
+            format="jpg",
+            size=1024,
+            width=1024,
+            height=768,
+            subset=DatasetItemSubset.VALIDATION,
+            user_reviewed=False,
+            project_id=str(project.id),
+            created_at=datetime.fromisoformat("2025-02-06T00:00:00Z"),
+        ),
+        DatasetItemDB(
+            name="validation2",
+            format="jpg",
+            size=1024,
+            width=1024,
+            height=768,
+            subset=DatasetItemSubset.VALIDATION,
+            user_reviewed=False,
+            project_id=str(project.id),
+            created_at=datetime.fromisoformat("2025-02-07T00:00:00Z"),
+        ),
+    ]
+
+    # Testing items
+    testing_items = [
+        DatasetItemDB(
+            name="testing1",
+            format="jpg",
+            size=1024,
+            width=1024,
+            height=768,
+            subset=DatasetItemSubset.TESTING,
+            user_reviewed=False,
+            project_id=str(project.id),
+            created_at=datetime.fromisoformat("2025-02-08T00:00:00Z"),
+        ),
+    ]
+
+    db_dataset_items = [*unassigned_items, *training_items, *validation_items, *testing_items]
+    db_session.add_all(db_dataset_items)
+    db_session.flush()
+
+    return project, db_dataset_items
+
+
 class TestDatasetServiceIntegration:
     """Integration tests for DatasetService."""
 
@@ -1069,3 +1185,155 @@ def test_list_dataset_items_no_label_filter(
         assert len(dataset_items) == 4
         item_names = {item.name for item in dataset_items}
         assert item_names == {"item_no_labels", "item_label_0", "item_label_1", "item_both_labels"}
+
+    @pytest.mark.parametrize(
+        "subset, expected_count",
+        [
+            (None, 8),  # All items
+            ("unassigned", 2),  # 2 unassigned items
+            ("training", 3),  # 3 training items
+            ("validation", 2),  # 2 validation items
+            ("testing", 1),  # 1 testing item
+        ],
+    )
+    def test_count_dataset_items_with_subset(
+        self,
+        fxt_dataset_service: DatasetService,
+        fxt_project_with_subset_items: tuple[ProjectView, list[DatasetItemDB]],
+        subset: str | None,
+        expected_count: int,
+    ) -> None:
+        """Test counting dataset items with subset filter."""
+        project, db_dataset_items = fxt_project_with_subset_items
+
+        count = fxt_dataset_service.count_dataset_items(project=project, subset=subset)
+
+        assert count == expected_count
+
+    @pytest.mark.parametrize(
+        "subset, expected_names",
+        [
+            (
+                None,
+                [
+                    "unassigned1",
+                    "unassigned2",
+                    "training1",
+                    "training2",
+                    "training3",
+                    "validation1",
+                    "validation2",
+                    "testing1",
+                ],
+            ),
+            ("unassigned", ["unassigned1", "unassigned2"]),
+            ("training", ["training1", "training2", "training3"]),
+            ("validation", ["validation1", "validation2"]),
+            ("testing", ["testing1"]),
+        ],
+    )
+    def test_list_dataset_items_with_subset(
+        self,
+        fxt_dataset_service: DatasetService,
+        fxt_project_with_subset_items: tuple[ProjectView, list[DatasetItemDB]],
+        subset: str | None,
+        expected_names: list[str],
+    ) -> None:
+        """Test listing dataset items with subset filter."""
+        project, db_dataset_items = fxt_project_with_subset_items
+
+        dataset_items = fxt_dataset_service.list_dataset_items(
+            project=project,
+            limit=20,
+            offset=0,
+            subset=subset,
+        )
+
+        assert len(dataset_items) == len(expected_names)
+        actual_names = sorted([item.name for item in dataset_items])
+        assert actual_names == sorted(expected_names)
+
+    @pytest.mark.parametrize(
+        "subset, limit, offset, expected_count",
+        [
+            ("unassigned", 1, 0, 1),  # First page of unassigned
+            ("unassigned", 1, 1, 1),  # Second page of unassigned
+            ("unassigned", 1, 2, 0),  # Beyond available unassigned items
+            ("training", 2, 0, 2),  # First page of training
+            ("training", 2, 2, 1),  # Second page of training (only 1 left)
+            ("validation", 10, 0, 2),  # All validation items
+            ("testing", 10, 0, 1),  # All testing items
+        ],
+    )
+    def test_list_dataset_items_with_subset_pagination(
+        self,
+        fxt_dataset_service: DatasetService,
+        fxt_project_with_subset_items: tuple[ProjectView, list[DatasetItemDB]],
+        subset: str | None,
+        limit: int,
+        offset: int,
+        expected_count: int,
+    ) -> None:
+        """Test listing dataset items with subset filter and pagination."""
+        project, db_dataset_items = fxt_project_with_subset_items
+
+        dataset_items = fxt_dataset_service.list_dataset_items(
+            project=project,
+            limit=limit,
+            offset=offset,
+            subset=subset,
+        )
+
+        assert len(dataset_items) == expected_count
+
+    def test_subset_filter_verifies_data_correctness(
+        self,
+        fxt_dataset_service: DatasetService,
+        fxt_project_with_subset_items: tuple[ProjectView, list[DatasetItemDB]],
+    ) -> None:
+        """Test that subset filter returns items with correct subset values."""
+        project, db_dataset_items = fxt_project_with_subset_items
+
+        # Unassigned items should have subset=unassigned
+        unassigned_items = fxt_dataset_service.list_dataset_items(
+            project=project,
+            limit=20,
+            offset=0,
+            subset="unassigned",
+        )
+        assert len(unassigned_items) == 2
+        for item in unassigned_items:
+            assert item.subset == DatasetItemSubset.UNASSIGNED
+
+        # Training items should have subset=training
+        training_items = fxt_dataset_service.list_dataset_items(
+            project=project,
+            limit=20,
+            offset=0,
+            subset="training",
+        )
+        assert len(training_items) == 3
+        for item in training_items:
+            assert item.subset == DatasetItemSubset.TRAINING
+
+        # Validation items should have subset=validation
+        validation_items = fxt_dataset_service.list_dataset_items(
+            project=project,
+            limit=20,
+            offset=0,
+            subset="validation",
+        )
+        assert len(validation_items) == 2
+        for item in validation_items:
+            assert item.subset == DatasetItemSubset.VALIDATION
+
+        # Testing items should have subset=testing
+        testing_items = fxt_dataset_service.list_dataset_items(
+            project=project,
+            limit=20,
+            offset=0,
+            subset="testing",
+        )
+        assert len(testing_items) == 1
+        for item in testing_items:
+            assert item.subset == DatasetItemSubset.TESTING