Add baseline read methods to TIMDEXDataset

Why these changes are being introduced:

A primary responsibility of the TIMDEXDataset class is to provide
performant and memory efficient reading of a dataset.  It is anticipated
that additional read methods may be required, for specific or niche
situations, but some simple baseline ones are needed at this time.

How this addresses that need:
* Adds methods for reading pyarrow batches, pandas dataframes, and python
dictionaries from a dataset.

Side effects of this change:
* Applications like timdex lambdas or TIM can now read records from dataset

Relevant ticket(s):
* https://mitlibraries.atlassian.net/browse/TIMX-417

Author: ghukill

tests/conftest.py

@@ -49,7 +49,23 @@ def fixed_local_dataset(tmp_path) -> TIMDEXDataset:
     method.
     """
     timdex_dataset = TIMDEXDataset(str(tmp_path / "fixed_local_dataset/"))
-    timdex_dataset.write(generate_sample_records(num_records=5_000, run_id="abc123"))
+    for source, run_id in [
+        ("alma", "abc123"),
+        ("dspace", "def456"),
+        ("aspace", "ghi789"),
+        ("libguides", "jkl123"),
+        ("gismit", "mno456"),
+    ]:
+        timdex_dataset.write(
+            generate_sample_records(
+                num_records=1_000,
+                timdex_record_id_prefix=source,
+                source=source,
+                run_date="2024-12-01",
+                run_id=run_id,
+            )
+        )
+    timdex_dataset.load()
     return timdex_dataset
 
 

tests/test_dataset_read.py

@@ -0,0 +1,91 @@
+# ruff: noqa: PLR2004, PD901
+
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+DATASET_COLUMNS_SET = {
+    "timdex_record_id",
+    "source_record",
+    "transformed_record",
+    "source",
+    "run_date",
+    "run_type",
+    "run_id",
+    "action",
+    "year",
+    "month",
+    "day",
+}
+
+
+def test_read_batches_yields_pyarrow_record_batches(fixed_local_dataset):
+    batches = fixed_local_dataset.read_batches_iter()
+    batch = next(batches)
+    assert isinstance(batch, pa.RecordBatch)
+
+
+def test_read_batches_all_columns_by_default(fixed_local_dataset):
+    batches = fixed_local_dataset.read_batches_iter()
+    batch = next(batches)
+    assert set(batch.column_names) == DATASET_COLUMNS_SET
+
+
+def test_read_batches_filter_columns(fixed_local_dataset):
+    columns_subset = ["source", "transformed_record"]
+    batches = fixed_local_dataset.read_batches_iter(columns=columns_subset)
+    batch = next(batches)
+    assert set(batch.column_names) == set(columns_subset)
+
+
+def test_read_batches_no_filters_gets_full_dataset(fixed_local_dataset):
+    batches = fixed_local_dataset.read_batches_iter()
+    table = pa.Table.from_batches(batches)
+    assert len(table) == fixed_local_dataset.row_count
+
+
+def test_read_batches_with_filters_gets_subset_of_dataset(fixed_local_dataset):
+    batches = fixed_local_dataset.read_batches_iter(
+        source="libguides",
+        run_date="2024-12-01",
+        run_type="daily",
+        action="index",
+    )
+
+    table = pa.Table.from_batches(batches)
+    assert len(table) == 1_000
+    assert len(table) < fixed_local_dataset.row_count
+
+    # assert loaded dataset is unchanged by filtering for a read method
+    assert fixed_local_dataset.row_count == 5_000
+
+
+def test_read_dataframe_batches_yields_dataframes(fixed_local_dataset):
+    df_iter = fixed_local_dataset.read_dataframes_iter()
+    df_batch = next(df_iter)
+    assert isinstance(df_batch, pd.DataFrame)
+    assert len(df_batch) == 1_000
+
+
+def test_read_dataframe_reads_all_dataset_rows_after_filtering(fixed_local_dataset):
+    df = fixed_local_dataset.read_dataframe()
+    assert isinstance(df, pd.DataFrame)
+    assert len(df) == fixed_local_dataset.row_count
+
+
+def test_read_dicts_yields_dictionary_for_each_dataset_record(fixed_local_dataset):
+    records = fixed_local_dataset.read_dicts_iter()
+    record = next(records)
+    assert isinstance(record, dict)
+    assert set(record.keys()) == DATASET_COLUMNS_SET
+
+
+def test_read_batches_filter_to_none_returns_empty_list(fixed_local_dataset):
+    batches = fixed_local_dataset.read_batches_iter(source="not-gonna-find-me")
+    assert list(batches) == []
+
+
+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)

timdex_dataset_api/dataset.py

@@ -10,6 +10,7 @@
 from typing import TYPE_CHECKING, TypedDict, Unpack
 
 import boto3
+import pandas as pd
 import pyarrow as pa
 import pyarrow.compute as pc
 import pyarrow.dataset as ds
@@ -388,3 +389,82 @@ def log_write_statistics(self, start_time: float) -> None:
             f"total rows: {total_rows}, "
             f"total size: {total_size}"
         )
+
+    def read_batches_iter(
+        self,
+        columns: list[str] | None = None,
+        batch_size: int = DEFAULT_BATCH_SIZE,
+        **filters: Unpack[DatasetFilters],
+    ) -> Iterator[pa.RecordBatch]:
+        """Yield pyarrow.RecordBatches from the dataset.
+
+        While batch_size will limit the max rows per batch, filtering may result in some
+        batches have less than this limit.
+
+        Args:
+            - columns: list[str], list of columns to return from the dataset
+            - batch_size: int, max number of rows to yield per batch
+            - filter_kwargs: pairs of column:value to filter the dataset
+        """
+        if not self.dataset:
+            raise DatasetNotLoadedError(
+                "Dataset is not loaded. Please call the `load` method first."
+            )
+        dataset = self._get_filtered_dataset(**filters)
+        for batch in dataset.to_batches(columns=columns, batch_size=batch_size):
+            if len(batch) > 0:
+                yield batch
+
+    def read_dataframes_iter(
+        self,
+        columns: list[str] | None = None,
+        batch_size: int = DEFAULT_BATCH_SIZE,
+        **filters: Unpack[DatasetFilters],
+    ) -> Iterator[pd.DataFrame]:
+        """Yield record batches as Pandas DataFrames from the dataset.
+
+        Args: see self.read_batches_iter()
+        """
+        for record_batch in self.read_batches_iter(
+            columns=columns, batch_size=batch_size, **filters
+        ):
+            yield record_batch.to_pandas()
+
+    def read_dataframe(
+        self,
+        columns: list[str] | None = None,
+        batch_size: int = DEFAULT_BATCH_SIZE,
+        **filters: Unpack[DatasetFilters],
+    ) -> pd.DataFrame | None:
+        """Yield record batches as Pandas DataFrames and concatenate to single dataframe.
+
+        WARNING: this will pull all records from currently filtered dataset into memory.
+
+        If no batches are found based on filtered dataset, None is returned.
+
+        Args: see self.read_batches_iter()
+        """
+        df_batches = [
+            record_batch.to_pandas()
+            for record_batch in self.read_batches_iter(
+                columns=columns, batch_size=batch_size, **filters
+            )
+        ]
+        if not df_batches:
+            return None
+        return pd.concat(df_batches)
+
+    def read_dicts_iter(
+        self,
+        columns: list[str] | None = None,
+        batch_size: int = DEFAULT_BATCH_SIZE,
+        **filters: Unpack[DatasetFilters],
+    ) -> Iterator[dict]:
+        """Yield individual record rows as dictionaries from the dataset.
+
+        Args: see self.read_batches_iter()
+        """
+        for record_batch in self.read_batches_iter(
+            columns=columns, batch_size=batch_size, **filters
+        ):
+            yield from record_batch.to_pylist()