version

Author: jtwinrelevanceai

ai_transform/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.31.3"
+__version__ = "0.32.0"
 
 from ai_transform.timer import Timer
 

ai_transform/engine/abstract_engine.py

@@ -106,8 +106,10 @@ def __init__(
             filters = []
         assert isinstance(filters, list), "Filters must be applied as a list of Dictionaries"
 
-        if not refresh:
-            filters += self._get_refresh_filter(select_fields, dataset)
+        self._refresh = refresh
+        self._after_id = after_id
+
+        filters += self._get_refresh_filter()
         filters += self._get_workflow_filter()
 
         self._filters = filters
@@ -117,9 +119,6 @@ def __init__(
         else:
             self._size = dataset.len(filters=filters) if self._limit_documents is None else self._limit_documents
 
-        self._refresh = refresh
-        self._after_id = after_id
-
         self._successful_documents = 0
         self._success_ratio = None
 
@@ -205,36 +204,36 @@ def _operate(self, mini_batch):
             self._successful_documents += len(mini_batch)
             return transformed_batch
 
-    def _get_refresh_filter(self, select_fields: List[str], dataset: Dataset):
+    def _get_refresh_filter(self):
         # initialize the refresh filter container
-        refresh_filters = {"filter_type": "or", "condition_value": []}
+        input_field_filters = {"filter_type": "or", "condition_value": []}
 
         # initialize where the filters are going
-        input_field_filters = []
         output_field_filters = {"filter_type": "or", "condition_value": []}
 
-        # We want documents where all select_fields exists
+        # We want documents where any of the select_fields exists
         # as these are needed for operator ...
-        for field in select_fields:
-            input_field_filters += dataset[field].exists()
-
-        # ... and where any of its output_fields dont exist
-        for operator in self.operators:
-            if operator.output_fields is not None:
-                for output_field in operator.output_fields:
-                    output_field_filters["condition_value"] += dataset[output_field].not_exists()
-
         # We construct this as:
         #
-        #   input_field1 and input_field2 and (not output_field1 or not output_field2)
+        #   (input_field1 or input_field2) and (not output_field1 or not output_field2)
         #
         # This use case here is for two input fields and two output fields
         # tho this extends to arbitrarily many.
-        refresh_filters["condition_value"] = input_field_filters
-        refresh_filters["condition_value"] += [output_field_filters]
+        for field in self._select_fields:
+            input_field_filters["condition_value"] += self.dataset[field].exists()
+
+        # ... and where any of its output_fields dont exist
+        if not self._refresh:
+            for operator in self.operators:
+                if operator.output_fields is not None:
+                    for output_field in operator.output_fields:
+                        output_field_filters["condition_value"] += self.dataset[output_field].not_exists()
+
+            return [input_field_filters, output_field_filters]
 
-        # Wrap in list at end
-        return [refresh_filters]
+        else:
+            # Wrap in list at end
+            return [input_field_filters]
 
     def _get_workflow_filter(self, field: str = "_id"):
         # Get the required workflow filter as an environment variable

tests/conftest.py

@@ -69,14 +69,34 @@ def partial_dataset(test_client: Client) -> Dataset:
     dataset_id = f"_sample_dataset_{salt}"
     dataset = test_client.Dataset(dataset_id, expire=True)
     documents = mock_documents(1000)
+    fields = ["sample_1_label", "sample_2_label", "sample_3_label"]
     for document in documents:
-        for field in random.choices(document.keys(), k=min(len(document), 5)):
+        for field in random.sample(fields, k=random.randint(1, 3)):
             document.pop(field)
     dataset.insert_documents(documents)
     yield dataset
     test_client.delete_dataset(dataset_id)
 
 
+@pytest.fixture(scope="class")
+def partial_dataset_with_outputs(test_client: Client) -> Dataset:
+    salt = "".join(random.choices(string.ascii_lowercase, k=10))
+    dataset_id = f"_sample_dataset_{salt}"
+    dataset = test_client.Dataset(dataset_id, expire=True)
+    documents = mock_documents(1000)
+    fields = ["sample_1_label", "sample_2_label", "sample_3_label"]
+    for document in documents:
+        for field in random.sample(fields, k=random.randint(1, 3)):
+            document.pop(field)
+    for document in documents:
+        for field in fields:
+            if document.get(field) and random.random() < 0.5:
+                document[field + "_output"] = document[field] + "_output"
+    dataset.insert_documents(documents)
+    yield dataset
+    test_client.delete_dataset(dataset_id)
+
+
 @pytest.fixture(scope="class")
 def mixed_dataset(test_client: Client) -> Dataset:
     salt = "".join(random.choices(string.ascii_lowercase, k=10))
@@ -164,6 +184,26 @@ def transform(self, documents: DocumentList) -> DocumentList:
     return ExampleOperator()
 
 
+@pytest.fixture(scope="function")
+def test_partial_operator() -> AbstractOperator:
+    class PartialOperator(AbstractOperator):
+        def __init__(self, fields):
+            super().__init__(input_fields=fields, output_fields=[field + "_output" for field in fields])
+
+        def transform(self, documents: DocumentList) -> DocumentList:
+            """
+            Main transform function
+            """
+            for input_field, output_field in zip(self.input_fields, self.output_fields):
+                for document in documents:
+                    if document.get(input_field):
+                        document[output_field] = document[input_field] + "_output"
+
+            return documents
+
+    return PartialOperator
+
+
 @pytest.fixture(scope="function")
 def test_paid_operator() -> AbstractOperator:
     class ExampleOperator(AbstractOperator):

tests/core/test_engine/test_stable_engine.py

@@ -1,3 +1,7 @@
+import uuid
+
+from typing import Type
+
 from ai_transform.dataset.dataset import Dataset
 from ai_transform.engine.stable_engine import StableEngine
 from ai_transform.engine.small_batch_stable_engine import SmallBatchStableEngine
@@ -6,21 +10,55 @@
 from ai_transform.workflow.abstract_workflow import Workflow
 
 
+def _random_id():
+    return str(uuid.uuid4())
+
+
 class TestStableEngine:
     def test_stable_engine(self, full_dataset: Dataset, test_operator: AbstractOperator):
         engine = StableEngine(full_dataset, test_operator, worker_number=0)
-        workflow = Workflow(name="workflow_test123", engine=engine, job_id="test_job123")
+        workflow = Workflow(name=_random_id(), engine=engine, job_id=_random_id())
         workflow.run()
         assert engine.success_ratio == 1
 
     def test_small_batch_stable_engine(self, full_dataset: Dataset, test_operator: AbstractOperator):
         engine = SmallBatchStableEngine(full_dataset, test_operator)
-        workflow = Workflow(name="workflow_test123", engine=engine, job_id="test_job123")
+        workflow = Workflow(name=_random_id(), engine=engine, job_id=_random_id())
         workflow.run()
         assert engine.success_ratio == 1
 
-    def test_stable_engine_filters(self, partial_dataset: Dataset, test_operator: AbstractOperator):
-        engine = StableEngine(partial_dataset, test_operator, select_fields=["sample_1_label"])
-        workflow = Workflow(name="workflow_test123", engine=engine, job_id="test_job123")
+
+class TestStableEngineFilters:
+    _SELECTED_FIELDS = ["sample_1_label", "sample_2_label", "sample_3_label"]
+
+    def test_stable_engine_filters1(self, partial_dataset: Dataset, test_partial_operator: Type[AbstractOperator]):
+        prev_health = partial_dataset.health()
+        operator = test_partial_operator(self._SELECTED_FIELDS)
+
+        engine = StableEngine(partial_dataset, operator, select_fields=self._SELECTED_FIELDS)
+        workflow = Workflow(name=_random_id(), engine=engine, job_id=_random_id())
         workflow.run()
+
+        post_health = partial_dataset.health()
+        for input_field, output_field in zip(operator.input_fields, operator.output_fields):
+            assert prev_health[input_field]["exists"] == post_health[output_field]["exists"]
+
+        assert engine.success_ratio == 1
+
+    def test_stable_engine_filters2(
+        self, partial_dataset_with_outputs: Dataset, test_partial_operator: Type[AbstractOperator]
+    ):
+        prev_health = partial_dataset_with_outputs.health()
+        operator = test_partial_operator(self._SELECTED_FIELDS)
+
+        engine = StableEngine(
+            partial_dataset_with_outputs, operator, select_fields=self._SELECTED_FIELDS, refresh=False
+        )
+        workflow = Workflow(name=_random_id(), engine=engine, job_id=_random_id())
+        workflow.run()
+
+        post_health = partial_dataset_with_outputs.health()
+        for input_field, output_field in zip(operator.input_fields, operator.output_fields):
+            assert prev_health[input_field]["exists"] == post_health[output_field]["exists"]
+
         assert engine.success_ratio == 1