lightly-studio 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

lightly_studio/core/dataset_query/boolean_expression.py

@@ -0,0 +1,67 @@
+"""Classes for boolean expressions in dataset queries."""
+
+from __future__ import annotations
+
+from sqlalchemy import ColumnElement, and_, false, not_, or_, true
+
+from lightly_studio.core.dataset_query.match_expression import MatchExpression
+
+
+class AND(MatchExpression):
+    """Logical AND operation between other MatchExpression objects."""
+
+    def __init__(self, *terms: MatchExpression) -> None:
+        """Initialize AND Expression with multiple MatchExpression terms.
+
+        Args:
+            terms: The MatchExpression instances to combine with AND. They can also be nested.
+        """
+        self.terms = terms
+
+    def get(self) -> ColumnElement[bool]:
+        """Combine expressions of all terms using AND.
+
+        Returns:
+            The combined SQLAlchemy expression.
+        """
+        return and_(true(), *(term.get() for term in self.terms))
+
+
+class OR(MatchExpression):
+    """Logical OR operation between other MatchExpression objects."""
+
+    def __init__(self, *terms: MatchExpression) -> None:
+        """Initialize OR Expression with multiple MatchExpression terms.
+
+        Args:
+            terms: The MatchExpression instances to combine with OR. They can also be nested.
+        """
+        self.terms = terms
+
+    def get(self) -> ColumnElement[bool]:
+        """Combine expressions of all terms using OR.
+
+        Returns:
+            The combined SQLAlchemy expression.
+        """
+        return or_(false(), *(term.get() for term in self.terms))
+
+
+class NOT(MatchExpression):
+    """Logical NOT operation for a MatchExpression object."""
+
+    def __init__(self, term: MatchExpression) -> None:
+        """Initialize NOT Expression with a MatchExpression term.
+
+        Args:
+            term: The MatchExpression to be negated. It can also be nested.
+        """
+        self.term = term
+
+    def get(self) -> ColumnElement[bool]:
+        """Negate the expression of the term using NOT.
+
+        Returns:
+            The combined SQLAlchemy expression.
+        """
+        return not_(self.term.get())

lightly_studio/core/dataset_query/dataset_query.py

@@ -0,0 +1,216 @@
+"""Matching functionality for filtering database samples based on field conditions."""
+
+from __future__ import annotations
+
+from typing import Iterator
+
+from sqlmodel import Session, select
+
+from lightly_studio.core.dataset_query.match_expression import MatchExpression
+from lightly_studio.core.dataset_query.order_by import OrderByExpression, OrderByField
+from lightly_studio.core.dataset_query.sample_field import SampleField
+from lightly_studio.core.sample import Sample
+from lightly_studio.export.export_dataset import DatasetExport
+from lightly_studio.models.dataset import DatasetTable
+from lightly_studio.models.sample import SampleTable
+from lightly_studio.resolvers import tag_resolver
+from lightly_studio.selection.select import Selection
+
+_SliceType = slice  # to avoid shadowing built-in slice in type annotations
+
+
+class DatasetQuery:
+    """Class for executing querying on a dataset."""
+
+    def __init__(self, dataset: DatasetTable, session: Session) -> None:
+        """Initialize with dataset and database session.
+
+        Args:
+            dataset: The dataset to query.
+            session: Database session for executing queries.
+        """
+        self.dataset = dataset
+        self.session = session
+        self.match_expression: MatchExpression | None = None
+        self.order_by_expressions: list[OrderByExpression] | None = None
+        self._slice: _SliceType | None = None
+
+    def match(self, match_expression: MatchExpression) -> DatasetQuery:
+        """Store a field condition for filtering.
+
+        Args:
+            match_expression: Defines the filter.
+
+        Returns:
+            Self for method chaining.
+
+        Raises:
+            ValueError: If match() has already been called on this instance.
+        """
+        if self.match_expression is not None:
+            raise ValueError("match() can only be called once per DatasetQuery instance")
+
+        self.match_expression = match_expression
+        return self
+
+    def order_by(self, *order_by: OrderByExpression) -> DatasetQuery:
+        """Store ordering expressions.
+
+        Args:
+            order_by: One or more ordering expressions. They are applied in order.
+                E.g. first ordering by sample width and then by sample file_name will
+                only order the samples with the same sample width by file_name.
+
+        Returns:
+            Self for method chaining.
+
+        Raises:
+            ValueError: If order_by() has already been called on this instance.
+        """
+        if self.order_by_expressions:
+            raise ValueError("order_by() can only be called once per DatasetQuery instance")
+
+        self.order_by_expressions = list(order_by)
+        return self
+
+    def slice(self, offset: int = 0, limit: int | None = None) -> DatasetQuery:
+        """Apply offset and limit to results.
+
+        Args:
+            offset: Number of items to skip from beginning (default: 0).
+            limit: Maximum number of items to return (None = no limit).
+
+        Returns:
+            Self for method chaining.
+
+        Raises:
+            ValueError: If slice() has already been called on this instance.
+        """
+        if self._slice is not None:
+            raise ValueError("slice() can only be called once per DatasetQuery instance")
+
+        # Convert to slice object for internal consistency
+        stop = None if limit is None else offset + limit
+        self._slice = _SliceType(offset, stop)
+        return self
+
+    def __getitem__(self, key: _SliceType) -> DatasetQuery:
+        """Enable bracket notation for slicing.
+
+        Args:
+            key: A slice object (e.g., [10:20], [:50], [100:]).
+
+        Returns:
+            Self with slice applied.
+
+        Raises:
+            TypeError: If key is not a slice object.
+            ValueError: If slice contains unsupported features or conflicts with existing slice.
+        """
+        if not isinstance(key, _SliceType):
+            raise TypeError(
+                "DatasetQuery only supports slice notation, not integer indexing. "
+                "Use execute() to get results as a list for element access."
+            )
+
+        # Validate unsupported features
+        if key.step is not None:
+            raise ValueError("Strides are not supported. Use simple slices like [start:stop].")
+
+        if (key.start is not None and key.start < 0) or (key.stop is not None and key.stop < 0):
+            raise ValueError("Negative indices are not supported. Use positive indices only.")
+
+        # Check for conflicts with existing slice
+        if self._slice is not None:
+            raise ValueError("Cannot use bracket notation after slice() has been called.")
+
+        # Set slice and return self
+        self._slice = key
+        return self
+
+    def __iter__(self) -> Iterator[Sample]:
+        """Iterate over the query results.
+
+        Returns:
+            Iterator of Sample objects from the database.
+        """
+        # Build query
+        query = select(SampleTable).where(SampleTable.dataset_id == self.dataset.dataset_id)
+
+        # Apply filter if present
+        if self.match_expression:
+            query = query.where(self.match_expression.get())
+
+        # Apply ordering
+        if self.order_by_expressions:
+            for order_by in self.order_by_expressions:
+                query = order_by.apply(query)
+        else:
+            # Order by SampleField.created_at by default.
+            default_order_by = OrderByField(SampleField.created_at)
+            query = default_order_by.apply(query)
+
+        # Apply slicing if present
+        if self._slice is not None:
+            start = self._slice.start or 0
+            query = query.offset(start)
+            if self._slice.stop is not None:
+                limit = max(self._slice.stop - start, 0)
+                query = query.limit(limit)
+
+        # Execute query and yield results
+        for sample_table in self.session.exec(query):
+            yield Sample(inner=sample_table)
+
+    def to_list(self) -> list[Sample]:
+        """Execute the query and return the results as a list.
+
+        Returns:
+            List of Sample objects from the database.
+        """
+        return list(self)
+
+    def add_tag(self, tag_name: str) -> None:
+        """Add a tag to all samples returned by this query.
+
+        First, creates the tag if it doesn't exist. Then applies the tag to all samples
+        that match the current query filters. Samples already having that tag are unchanged,
+        as the database prevents duplicates.
+
+        Args:
+            tag_name: Name of the tag to add to matching samples.
+        """
+        # Get or create the tag
+        tag = tag_resolver.get_or_create_sample_tag_by_name(
+            session=self.session, dataset_id=self.dataset.dataset_id, tag_name=tag_name
+        )
+
+        # Execute query to get matching samples
+        samples = self.to_list()
+        sample_ids = [sample.sample_id for sample in samples]
+
+        # Use resolver to bulk assign tag (handles validation and edge cases)
+        tag_resolver.add_sample_ids_to_tag_id(
+            session=self.session, tag_id=tag.tag_id, sample_ids=sample_ids
+        )
+
+    def selection(self) -> Selection:
+        """Selection interface for this query.
+
+        The returned Selection snapshots the current query results immediately.
+        Mutating the query after calling this method will therefore not affect
+        the samples used by that Selection instance.
+
+        Returns:
+            Selection interface operating on the current query result snapshot.
+        """
+        input_sample_ids = (sample.sample_id for sample in self)
+        return Selection(
+            dataset_id=self.dataset.dataset_id,
+            session=self.session,
+            input_sample_ids=input_sample_ids,
+        )
+
+    def export(self) -> DatasetExport:
+        """Return a DatasetExport instance which can export the dataset in various formats."""
+        return DatasetExport(session=self.session, samples=self)

lightly_studio/core/dataset_query/field.py

@@ -0,0 +1,113 @@
+"""Base field classes for building dataset queries."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from datetime import datetime
+from typing import Any, Generic, TypeVar, Union
+
+from sqlalchemy.orm import Mapped
+
+from lightly_studio.core.dataset_query.field_expression import (
+    OrdinalFieldExpression,
+    StringFieldExpression,
+)
+
+T = TypeVar("T")
+
+
+class Field(ABC):
+    """Abstract base class for all field types in dataset queries."""
+
+    @abstractmethod
+    def get_sqlmodel_field(self) -> Mapped[Any]:
+        """Get the database column or property that this field represents.
+
+        Returns:
+            The database column or property for queries.
+        """
+
+
+class OrdinalField(Field, Generic[T]):
+    """Generic field for ordinal values that support comparison operations.
+
+    Ordinal values have a natural ordering and support all comparison operators:
+    >, <, >=, <=, ==, !=
+    """
+
+    def __init__(self, column: Mapped[T]) -> None:
+        """Initialize the ordinal field with a database column.
+
+        Args:
+            column: The database column this field represents.
+        """
+        self._column = column
+
+    def get_sqlmodel_field(self) -> Mapped[T]:
+        """Get the ordinal database column or property.
+
+        Returns:
+            The ordinal column for database queries.
+        """
+        return self._column
+
+    def __gt__(self, other: T) -> OrdinalFieldExpression[T]:
+        """Create a greater-than expression."""
+        return OrdinalFieldExpression(field=self, operator=">", value=other)
+
+    def __lt__(self, other: T) -> OrdinalFieldExpression[T]:
+        """Create a less-than expression."""
+        return OrdinalFieldExpression(field=self, operator="<", value=other)
+
+    def __ge__(self, other: T) -> OrdinalFieldExpression[T]:
+        """Create a greater-than-or-equal expression."""
+        return OrdinalFieldExpression(field=self, operator=">=", value=other)
+
+    def __le__(self, other: T) -> OrdinalFieldExpression[T]:
+        """Create a less-than-or-equal expression."""
+        return OrdinalFieldExpression(field=self, operator="<=", value=other)
+
+    def __eq__(self, other: T) -> OrdinalFieldExpression[T]:  # type: ignore[override]
+        """Create an equality expression."""
+        return OrdinalFieldExpression(field=self, operator="==", value=other)
+
+    def __ne__(self, other: T) -> OrdinalFieldExpression[T]:  # type: ignore[override]
+        """Create a not-equal expression."""
+        return OrdinalFieldExpression(field=self, operator="!=", value=other)
+
+
+NumericalField = OrdinalField[Union[float, int]]
+DatetimeField = OrdinalField[datetime]
+
+
+class StringField(Field):
+    """Field for string values that supports equality operations.
+
+    Optional refactor when needed: Split into
+    - StringField(ABC) with the comparison operators.
+    - StringColumnField(StringField) for the __init__ and get_sqlmodel_field implementation.
+    """
+
+    def __init__(self, column: Mapped[str]) -> None:
+        """Initialize the string field with a database column.
+
+        Args:
+            column: The database column this field represents.
+        """
+        self._column = column
+
+    def get_sqlmodel_field(self) -> Mapped[str]:
+        """Get the string database column or property.
+
+        Returns:
+            The string column for database queries.
+        """
+        return self._column
+
+    def __eq__(self, other: str) -> StringFieldExpression:  # type: ignore[override]
+        """Create an equality expression."""
+        return StringFieldExpression(field=self, operator="==", value=other)
+
+    def __ne__(self, other: str) -> StringFieldExpression:  # type: ignore[override]
+        """Create a not-equal expression."""
+        return StringFieldExpression(field=self, operator="!=", value=other)

lightly_studio/core/dataset_query/field_expression.py

@@ -0,0 +1,79 @@
+"""Field expressions for building specific query conditions."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime
+from typing import TYPE_CHECKING, Callable, Generic, Literal, TypeVar, Union
+
+from sqlalchemy import ColumnElement
+from sqlalchemy.orm import Mapped
+
+from lightly_studio.core.dataset_query.match_expression import MatchExpression
+
+if TYPE_CHECKING:
+    from lightly_studio.core.dataset_query.field import (
+        OrdinalField,
+        StringField,
+    )
+
+
+T = TypeVar("T")
+
+"""Conditions themselves, in the format <field> <operator> <value>
+
+Example:
+SampleField.file_name == "img1.jpg",
+becomes StringField(file_name) == "img1.jpg",
+becomes StringFieldExpression(field=StringField(file_name), operator="==", value="img1.jpg")
+becomes SQLQuery.where(...)
+"""
+StringOperator = Literal["==", "!="]
+OrdinalOperator = Literal[">", "<", "==", ">=", "<=", "!="]
+
+
+@dataclass
+class OrdinalFieldExpression(MatchExpression, Generic[T]):
+    """Generic expression for ordinal field comparisons."""
+
+    field: OrdinalField[T]
+    operator: OrdinalOperator
+    value: T
+
+    def get(self) -> ColumnElement[bool]:
+        """Return the SQLAlchemy expression for this ordinal field expression."""
+        table_property = self.field.get_sqlmodel_field()
+        operations: dict[
+            OrdinalOperator,
+            Callable[[Mapped[T], T], ColumnElement[bool]],
+        ] = {
+            "<": lambda tp, v: tp < v,
+            "<=": lambda tp, v: tp <= v,
+            ">": lambda tp, v: tp > v,
+            ">=": lambda tp, v: tp >= v,
+            "==": lambda tp, v: tp == v,
+            "!=": lambda tp, v: tp != v,
+        }
+        return operations[self.operator](table_property, self.value)
+
+
+NumericalFieldExpression = OrdinalFieldExpression[Union[float, int]]
+DatetimeFieldExpression = OrdinalFieldExpression[datetime]
+
+
+@dataclass
+class StringFieldExpression(MatchExpression):
+    """Expression for string field comparisons."""
+
+    field: StringField
+    operator: StringOperator
+    value: str
+
+    def get(self) -> ColumnElement[bool]:
+        """Return the SQLAlchemy expression for this string field expression."""
+        table_property = self.field.get_sqlmodel_field()
+        operations: dict[StringOperator, Callable[[Mapped[str], str], ColumnElement[bool]]] = {
+            "==": lambda tp, v: tp == v,
+            "!=": lambda tp, v: tp != v,
+        }
+        return operations[self.operator](table_property, self.value)

lightly_studio/core/dataset_query/match_expression.py

@@ -0,0 +1,23 @@
+"""Base classes for match expressions in dataset queries."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+from sqlalchemy import ColumnElement
+
+
+class MatchExpression(ABC):
+    """Base class for all match expressions that can be applied to database queries.
+
+    This class provides the foundation for implementing complex query expressions
+    that can be combined using AND/OR operations in the future.
+    """
+
+    @abstractmethod
+    def get(self) -> ColumnElement[bool]:
+        """Get the SQLAlchemy expression for this match expression.
+
+        Returns:
+            The combined SQLAlchemy expression.
+        """

lightly_studio/core/dataset_query/order_by.py

@@ -0,0 +1,79 @@
+"""Classes for order by expressions in dataset queries."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+from sqlmodel.sql.expression import SelectOfScalar
+from typing_extensions import Self
+
+from lightly_studio.core.dataset_query.field import Field
+from lightly_studio.models.sample import SampleTable
+
+
+class OrderByExpression(ABC):
+    """Base class for all order by expressions that can be applied to database queries."""
+
+    def __init__(self, *, ascending: bool = True) -> None:
+        """Initialize the order by expression.
+
+        Args:
+            ascending: Whether to order in ascending (True) or descending (False) order.
+        """
+        self.ascending = ascending
+
+    @abstractmethod
+    def apply(self, query: SelectOfScalar[SampleTable]) -> SelectOfScalar[SampleTable]:
+        """Apply this ordering to a SQLModel Select query.
+
+        Args:
+            query: The SQLModel Select query to modify.
+
+        Returns:
+            The modified query after ordering
+        """
+
+    def asc(self) -> Self:
+        """Set the ordering to ascending.
+
+        Returns:
+            Self for method chaining.
+        """
+        self.ascending = True
+        return self
+
+    def desc(self) -> Self:
+        """Set the ordering to descending.
+
+        Returns:
+            Self for method chaining.
+        """
+        self.ascending = False
+        return self
+
+
+class OrderByField(OrderByExpression):
+    """Order by a specific field, either ascending or descending.
+
+    Args:
+        field: The field to order by.
+        ascending: Whether to order in ascending (True) or descending (False) order.
+    """
+
+    def __init__(self, field: Field) -> None:
+        """Initialize with field and order direction."""
+        super().__init__()
+        self.field = field
+
+    def apply(self, query: SelectOfScalar[SampleTable]) -> SelectOfScalar[SampleTable]:
+        """Apply this ordering to a SQLModel Select query.
+
+        Args:
+            query: The SQLModel Select query to modify.
+
+        Returns:
+            The modified query after ordering
+        """
+        if self.ascending:
+            return query.order_by(self.field.get_sqlmodel_field().asc())
+        return query.order_by(self.field.get_sqlmodel_field().desc())

lightly_studio/core/dataset_query/sample_field.py

@@ -0,0 +1,28 @@
+"""Fields for querying sample properties in the dataset query system."""
+
+from __future__ import annotations
+
+from sqlmodel import col
+
+from lightly_studio.core.dataset_query.field import (
+    DatetimeField,
+    NumericalField,
+    StringField,
+)
+from lightly_studio.core.dataset_query.tags_expression import TagsAccessor
+from lightly_studio.models.sample import SampleTable
+
+
+class SampleField:
+    """Providing access to predefined sample fields for queries.
+
+    This class provides static instances of different field types that can be
+    used to build database queries on sample properties.
+    """
+
+    file_name = StringField(col(SampleTable.file_name))
+    width = NumericalField(col(SampleTable.width))
+    height = NumericalField(col(SampleTable.height))
+    file_path_abs = StringField(col(SampleTable.file_path_abs))
+    created_at = DatetimeField(col(SampleTable.created_at))
+    tags = TagsAccessor()

lightly_studio/core/dataset_query/tags_expression.py

@@ -0,0 +1,46 @@
+"""Tag field classes for building dataset queries on sample tags."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from sqlalchemy import ColumnElement
+from sqlmodel import col
+
+from lightly_studio.core.dataset_query.match_expression import MatchExpression
+from lightly_studio.models.sample import SampleTable
+from lightly_studio.models.tag import TagTable
+
+
+class TagsAccessor:
+    """Provides access to tag operations for query building.
+
+    This class enables checking tag membership using the contains method:
+    SampleField.tags.contains("tag_name") returns a TagsContainsExpression.
+    """
+
+    def contains(self, tag_name: str) -> TagsContainsExpression:
+        """Check if a tag name is in the sample's tags.
+
+        Args:
+            tag_name: The name of the tag to check for.
+
+        Returns:
+            A TagsContainsExpression for building queries.
+        """
+        return TagsContainsExpression(tag_name=tag_name)
+
+
+@dataclass
+class TagsContainsExpression(MatchExpression):
+    """Expression for checking if a sample contains a specific tag."""
+
+    tag_name: str
+
+    def get(self) -> ColumnElement[bool]:
+        """Get the tag contains expression.
+
+        Returns:
+            The SQLAlchemy expression for this field expression.
+        """
+        return SampleTable.tags.any(col(TagTable.name) == self.tag_name)