Fix filtering by run metadata (#3344)

* Fix filtering artifact versions by metadata

* More fixes

* Json encode after extracting filter operator

* More fixes and tests

* Linting

* Linting

Author: schustmi

src/zenml/client.py

@@ -3821,7 +3821,7 @@ def list_pipeline_runs(
         templatable: Optional[bool] = None,
         tag: Optional[str] = None,
         user: Optional[Union[UUID, str]] = None,
-        run_metadata: Optional[Dict[str, str]] = None,
+        run_metadata: Optional[Dict[str, Any]] = None,
         pipeline: Optional[Union[UUID, str]] = None,
         code_repository: Optional[Union[UUID, str]] = None,
         model: Optional[Union[UUID, str]] = None,
@@ -3974,6 +3974,7 @@ def list_run_steps(
         user: Optional[Union[UUID, str]] = None,
         model_version_id: Optional[Union[str, UUID]] = None,
         model: Optional[Union[UUID, str]] = None,
+        run_metadata: Optional[Dict[str, Any]] = None,
         hydrate: bool = False,
     ) -> Page[StepRunResponse]:
         """List all pipelines.
@@ -4000,6 +4001,7 @@ def list_run_steps(
             cache_key: The cache key of the step run to filter by.
             code_hash: The code hash of the step run to filter by.
             status: The name of the run to filter by.
+            run_metadata: Filter by run metadata.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
 
@@ -4028,6 +4030,7 @@ def list_run_steps(
             user=user,
             model_version_id=model_version_id,
             model=model,
+            run_metadata=run_metadata,
         )
         step_run_filter_model.set_scope_workspace(self.active_workspace.id)
         return self.zen_store.list_run_steps(
@@ -4254,7 +4257,7 @@ def list_artifact_versions(
         user: Optional[Union[UUID, str]] = None,
         model: Optional[Union[UUID, str]] = None,
         pipeline_run: Optional[Union[UUID, str]] = None,
-        run_metadata: Optional[Dict[str, str]] = None,
+        run_metadata: Optional[Dict[str, Any]] = None,
         tag: Optional[str] = None,
         hydrate: bool = False,
     ) -> Page[ArtifactVersionResponse]:
@@ -4320,6 +4323,7 @@ def list_artifact_versions(
             user=user,
             model=model,
             pipeline_run=pipeline_run,
+            run_metadata=run_metadata,
         )
         artifact_version_filter_model.set_scope_workspace(
             self.active_workspace.id

src/zenml/models/v2/base/filter.py

@@ -171,6 +171,8 @@ def generate_query_conditions_from_column(self, column: Any) -> Any:
 class StrFilter(Filter):
     """Filter for all string fields."""
 
+    json_encode_value: bool = False
+
     ALLOWED_OPS: ClassVar[List[str]] = [
         GenericFilterOps.EQUALS,
         GenericFilterOps.NOT_EQUALS,
@@ -211,16 +213,6 @@ def generate_query_conditions_from_column(self, column: Any) -> Any:
         Raises:
             ValueError: the comparison of the column to a numeric value fails.
         """
-        if self.operation == GenericFilterOps.CONTAINS:
-            return column.like(f"%{self.value}%")
-        if self.operation == GenericFilterOps.STARTSWITH:
-            return column.startswith(f"{self.value}")
-        if self.operation == GenericFilterOps.ENDSWITH:
-            return column.endswith(f"{self.value}")
-        if self.operation == GenericFilterOps.NOT_EQUALS:
-            return column != self.value
-        if self.operation == GenericFilterOps.ONEOF:
-            return column.in_(self.value)
         if self.operation in {
             GenericFilterOps.GT,
             GenericFilterOps.LT,
@@ -254,7 +246,33 @@ def generate_query_conditions_from_column(self, column: Any) -> Any:
                     f"value '{self.value}' (must be numeric): {e}"
                 )
 
-        return column == self.value
+        if self.operation == GenericFilterOps.ONEOF:
+            assert isinstance(self.value, list)
+            # Convert the list of values to a list of json strings
+            json_list = (
+                [json.dumps(v) for v in self.value]
+                if self.json_encode_value
+                else self.value
+            )
+            return column.in_(json_list)
+
+        # Don't convert the value to a json string if the operation is contains
+        # because the quotes around strings will mess with the comparison
+        if self.operation == GenericFilterOps.CONTAINS:
+            return column.like(f"%{self.value}%")
+
+        json_value = (
+            json.dumps(self.value) if self.json_encode_value else self.value
+        )
+
+        if self.operation == GenericFilterOps.STARTSWITH:
+            return column.startswith(f"{json_value}")
+        if self.operation == GenericFilterOps.ENDSWITH:
+            return column.endswith(f"{json_value}")
+        if self.operation == GenericFilterOps.NOT_EQUALS:
+            return column != json_value
+
+        return column == json_value
 
 
 class UUIDFilter(StrFilter):
@@ -733,13 +751,15 @@ def generate_custom_query_conditions_for_column(
         value: Any,
         table: Type[SQLModel],
         column: str,
+        json_encode_value: bool = False,
     ) -> "ColumnElement[bool]":
         """Generate custom filter conditions for a column of a table.
 
         Args:
             value: The filter value.
             table: The table which contains the column.
             column: The column name.
+            json_encode_value: Whether to json encode the value.
 
         Returns:
             The query conditions.
@@ -748,6 +768,9 @@ def generate_custom_query_conditions_for_column(
         filter_ = FilterGenerator(table).define_filter(
             column=column, value=value, operator=operator
         )
+        if isinstance(filter_, StrFilter):
+            filter_.json_encode_value = json_encode_value
+
         return filter_.generate_query_conditions(table=table)
 
     @property

src/zenml/models/v2/core/artifact_version.py

@@ -559,7 +559,7 @@ class ArtifactVersionFilter(WorkspaceScopedFilter, TaggableFilter):
         description="Name/ID of a pipeline run that is associated with this "
         "artifact version.",
     )
-    run_metadata: Optional[Dict[str, str]] = Field(
+    run_metadata: Optional[Dict[str, Any]] = Field(
         default=None,
         description="The run_metadata to filter the artifact versions by.",
     )
@@ -683,13 +683,19 @@ def get_custom_filters(
                     RunMetadataResourceSchema.resource_id
                     == ArtifactVersionSchema.id,
                     RunMetadataResourceSchema.resource_type
-                    == MetadataResourceTypes.ARTIFACT_VERSION,
+                    == MetadataResourceTypes.ARTIFACT_VERSION.value,
                     RunMetadataResourceSchema.run_metadata_id
                     == RunMetadataSchema.id,
+                    self.generate_custom_query_conditions_for_column(
+                        value=key,
+                        table=RunMetadataSchema,
+                        column="key",
+                    ),
                     self.generate_custom_query_conditions_for_column(
                         value=value,
                         table=RunMetadataSchema,
                         column="value",
+                        json_encode_value=True,
                     ),
                 )
                 custom_filters.append(additional_filter)

src/zenml/models/v2/core/pipeline_run.py

@@ -688,7 +688,7 @@ class PipelineRunFilter(WorkspaceScopedFilter, TaggableFilter):
         union_mode="left_to_right",
     )
     unlisted: Optional[bool] = None
-    run_metadata: Optional[Dict[str, str]] = Field(
+    run_metadata: Optional[Dict[str, Any]] = Field(
         default=None,
         description="The run_metadata to filter the pipeline runs by.",
     )
@@ -915,13 +915,19 @@ def get_custom_filters(
                     RunMetadataResourceSchema.resource_id
                     == PipelineRunSchema.id,
                     RunMetadataResourceSchema.resource_type
-                    == MetadataResourceTypes.PIPELINE_RUN,
+                    == MetadataResourceTypes.PIPELINE_RUN.value,
                     RunMetadataResourceSchema.run_metadata_id
                     == RunMetadataSchema.id,
+                    self.generate_custom_query_conditions_for_column(
+                        value=key,
+                        table=RunMetadataSchema,
+                        column="key",
+                    ),
                     self.generate_custom_query_conditions_for_column(
                         value=value,
                         table=RunMetadataSchema,
                         column="value",
+                        json_encode_value=True,
                     ),
                 )
                 custom_filters.append(additional_filter)

src/zenml/models/v2/core/step_run.py

@@ -16,6 +16,7 @@
 from datetime import datetime
 from typing import (
     TYPE_CHECKING,
+    Any,
     ClassVar,
     Dict,
     List,
@@ -574,7 +575,7 @@ class StepRunFilter(WorkspaceScopedFilter):
         default=None,
         description="Name/ID of the model associated with the step run.",
     )
-    run_metadata: Optional[Dict[str, str]] = Field(
+    run_metadata: Optional[Dict[str, Any]] = Field(
         default=None,
         description="The run_metadata to filter the step runs by.",
     )
@@ -619,13 +620,19 @@ def get_custom_filters(
                 additional_filter = and_(
                     RunMetadataResourceSchema.resource_id == StepRunSchema.id,
                     RunMetadataResourceSchema.resource_type
-                    == MetadataResourceTypes.STEP_RUN,
+                    == MetadataResourceTypes.STEP_RUN.value,
                     RunMetadataResourceSchema.run_metadata_id
                     == RunMetadataSchema.id,
+                    self.generate_custom_query_conditions_for_column(
+                        value=key,
+                        table=RunMetadataSchema,
+                        column="key",
+                    ),
                     self.generate_custom_query_conditions_for_column(
                         value=value,
                         table=RunMetadataSchema,
                         column="value",
+                        json_encode_value=True,
                     ),
                 )
                 custom_filters.append(additional_filter)

tests/integration/functional/zen_stores/test_zen_store.py

@@ -2924,7 +2924,7 @@ def pipeline_to_log_metadata(metadata):
     step_to_log_metadata(metadata)
 
 
-def test_pipeline_run_filters_with_oneof_and_run_metadata(clean_client):
+def test_pipeline_run_filters_with_oneof(clean_client):
     store = clean_client.zen_store
 
     metadata_values = [3, 25, 100, "random_string", True]
@@ -2961,6 +2961,49 @@ def test_pipeline_run_filters_with_oneof_and_run_metadata(clean_client):
         PipelineRunFilter(name="oneof:random_value")
 
 
+def test_run_metadata_filtering(clean_client):
+    store = clean_client.zen_store
+
+    metadata_values = [3, 25, 100, "random_string", True]
+
+    for v in metadata_values:
+        pipeline_to_log_metadata(v)
+
+    # Test run metadata filtering with string value
+    runs_filter = StepRunFilter(run_metadata={"blupus": "random_string"})
+    runs = store.list_run_steps(step_run_filter_model=runs_filter)
+    assert len(runs.items) == 1
+
+    # Test run metadata filtering with boolean value
+    runs_filter = StepRunFilter(run_metadata={"blupus": True})
+    runs = store.list_run_steps(step_run_filter_model=runs_filter)
+    assert len(runs.items) == 1
+
+    # Test run metadata filtering with int value
+    runs_filter = StepRunFilter(run_metadata={"blupus": 3})
+    runs = store.list_run_steps(step_run_filter_model=runs_filter)
+    assert len(runs.items) == 1
+
+    # Test run metadata filtering for a non-existent key
+    runs_filter = StepRunFilter(run_metadata={"non-existent": 3})
+    runs = store.list_run_steps(step_run_filter_model=runs_filter)
+    assert len(runs) == 0
+
+    # Test run metadata filtering with non-existent value
+    runs_filter = StepRunFilter(run_metadata={"blupus": "non-existent"})
+    runs = store.list_run_steps(step_run_filter_model=runs_filter)
+    assert len(runs.items) == 0
+
+    # Test run metadata filtering with operator value
+    runs_filter = StepRunFilter(run_metadata={"blupus": "lt:30"})
+    runs = store.list_run_steps(step_run_filter_model=runs_filter)
+    assert len(runs) == 2  # The run with 3 and 25
+
+    for r in runs:
+        assert isinstance(r.run_metadata["blupus"], int)
+        assert r.run_metadata["blupus"] < 30
+
+
 # .--------------------.
 # | Pipeline run steps |
 # '--------------------'