Merge pull request #58 from MITLibraries/TIMX-453-read-transformed-records-from-dataset

Add read method for yielding transformed records to TIMDEXDataset

Author: jonavellecuerdo

tests/test_dataset_read.py

@@ -89,3 +89,10 @@ def test_read_dicts_filter_to_none_stopiteration_immediately(fixed_local_dataset
     batches = fixed_local_dataset.read_dicts_iter(source="not-gonna-find-me")
     with pytest.raises(StopIteration):
         next(batches)
+
+
+def test_read_transformed_records_yields_parsed_dictionary(fixed_local_dataset):
+    batches = fixed_local_dataset.read_transformed_records_iter()
+    transformed_record = next(batches)
+    assert isinstance(transformed_record, dict)
+    assert transformed_record == {"title": ["Hello World."]}

timdex_dataset_api/dataset.py

@@ -1,6 +1,7 @@
 """timdex_dataset_api/dataset.py"""
 
 import itertools
+import json
 import operator
 import time
 import uuid
@@ -468,3 +469,19 @@ def read_dicts_iter(
             columns=columns, batch_size=batch_size, **filters
         ):
             yield from record_batch.to_pylist()
+
+    def read_transformed_records_iter(
+        self,
+        batch_size: int = DEFAULT_BATCH_SIZE,
+        **filters: Unpack[DatasetFilters],
+    ) -> Iterator[dict]:
+        """Yield individual transformed records as dictionaries from the dataset.
+
+        If 'transformed_record' is None (i.e., action="skip"|"error"), the yield
+        statement will not be executed for the row.
+        """
+        for record_dict in self.read_dicts_iter(
+            columns=["transformed_record"], batch_size=batch_size, **filters
+        ):
+            if transformed_record := record_dict["transformed_record"]:
+                yield json.loads(transformed_record)