Repository: MITLibraries/timdex-dataset-api
Pull request #144: TIMX 494 - yield deduped, most recent records
Author: login: ghukill
Created: 2025-05-21T18:27:48Z | Merged: 2025-05-23T19:37:05Z

### Purpose and background context

This PR marks a somewhat important milestone in this library: support for reading only "current" records from the parquet dataset in a performant and memory-efficient manner.

This builds on PR #143, which introduced the `TIMDEXRunManager` which provided explicit lists of parquet files based on run ETL metadata.  With that in hand, `TIMDEXDataset` was updated in two key ways.

First, the `.load()` method was updated to include a `current_records:bool` keyword argument, which when True, will utilize the `TIMDEXRunManager` to set an explicit list of parquet files for any future reading.

Second, when `current_records` are requested a private attribute is set `TIMDEXDataset._dedupe_on_read = True`.  This informs _all_ read methods to perform record deduping on read, which they can confidently and easily do with the dataset parquet files explicitly set in reverse chronological order.

The majority of the work was not getting this functional, but keeping the API minimal, intuitive, and non-surprising where possible.  For example, to read only current `dspace` records from the dataset, only the following is required:

```python
from timdex_dataset_api import TIMDEXDataset

timdex_dataset = TIMDEXDataset("s3://bucket/dataset")
timdex_dataset.load(current_records=True, source="dspace")

# from here, all read methods will produce only current records
```

The first anticipated use case is for TIM, where a new CLI command along the lines of `reindex-source` would utilize this new `.load()` + read technique to yield  _all and only current_ records for a source purely from dataset data.  No re-harvests would be required.  This is useful for development work, and could be a very useful option in disaster recovery situations.

### How can a reviewer manually see the effects of these changes?

1- Set **Dev Admin** AWS credentials

2- Set following env vars (if not already):
```shell
TDA_LOG_LEVEL=DEBUG
WARNING_ONLY_LOGGERS=asyncio,botocore,urllib3,s3transfer,boto3
```

3- Start Ipython
```shell
pipenv run shell
```

4- Load the dataset for `libguides` current records only
```python
from timdex_dataset_api import TIMDEXDataset
from timdex_dataset_api.config import configure_dev_logger

configure_dev_logger()

timdex_dataset = TIMDEXDataset("s3://ghukill-test/timdex-dataset/prod_2025_05_05")
timdex_dataset.load(current_records=True, source="libguides")
```
- Note some logging outputs that indicate the dataset is loaded and filtered

5- Given the small number of records for this source, load them all into a dataframe (takes ~5 seconds):
```python
df = timdex_dataset.read_dataframe()
```

6- Inspecting the dataframe to confirm that `timdex_record_id` is unique:
```python
df.timdex_record_id.is_unique
# Out[3]: True
```

7- Observe the ETL runs present in the records:
```python
df[['run_date','run_type','run_id']].value_counts()
```
- The count is effectively how many records from that run are "current" (i.e. no new versions after that)
- Note the single `run_type="full"` has the bulk of the records, and `run_type="daily"` runs are just incrementally updating

8- Lastly, optional and just for kicks, an more full fledged example showing looping through all current records for `dspace`:

```python
import time

from timdex_dataset_api import TIMDEXDataset
from timdex_dataset_api.config import configure_dev_logger

configure_dev_logger()

timdex_dataset = TIMDEXDataset("s3://ghukill-test/timdex-dataset/prod_2025_05_05")
timdex_dataset.load(current_records=True, source="dspace")

read_columns = ["timdex_record_id", "run_date", "action", "run_id"]
t0 = time.perf_counter()
count = 0
t1 = time.perf_counter()
for i, batch in enumerate(timdex_dataset.read_batches_iter(columns=read_columns)):
    df = batch.to_pandas()
    row = df.iloc[0]
    count += len(batch)
    print(
        f"batch {i}, rows {len(batch)}, total {count}, run_date {row.run_date}, "
        f"run_id {row.run_id}, elapsed {time.perf_counter() - t1}"
    )
    t1 = time.perf_counter()
print(f"total: {time.perf_counter()-t0}")
```
- Note the batches start with small numbers of records, just updates, until you finally get to the first `run_type="full"` where you get lots of records
- If you have resource monitoring open with something like `btop`, note that memory usage remains low

### Includes new or updated dependencies?
NO

### Changes expectations for external applications?
YES: applications like TIM may utilize new TDA functionality to read only current records

### What are the relevant tickets?
- https://mitlibraries.atlassian.net/browse/TIMX-493
- https://mitlibraries.atlassian.net/browse/TIMX-494

### Developer
- [X] All new ENV is documented in README
- [X] All new ENV has been added to staging and production environments
- [X] All related Jira tickets are linked in commit message(s)
- [X] Stakeholder approval has been confirmed (or is not needed)

### Code Reviewer(s)
- [ ] The commit message is clear and follows our guidelines (not just this PR message)
- [ ] There are appropriate tests covering any new functionality
- [ ] The provided documentation is sufficient for understanding any new functionality introduced
- [ ] Any manual tests have been performed **or** provided examples verified
- [ ] New dependencies are appropriate or there were no changes




Diff:
diff --git a/tests/conftest.py b/tests/conftest.py
index 4838a6a..2e5e9ed 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -138,6 +138,7 @@ def dataset_with_runs_location(tmp_path) -> str:
         num_records, source, run_date, run_type, action, run_id = params
         records = generate_sample_records(
             num_records,
+            timdex_record_id_prefix=source,
             source=source,
             run_date=run_date,
             run_type=run_type,
@@ -147,3 +148,8 @@ def dataset_with_runs_location(tmp_path) -> str:
         timdex_dataset.write(records)
 
     return location
+
+
+@pytest.fixture
+def local_dataset_with_runs(dataset_with_runs_location) -> TIMDEXDataset:
+    return TIMDEXDataset(dataset_with_runs_location)
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
index a0c4793..0b0cbf2 100644
--- a/tests/test_dataset.py
+++ b/tests/test_dataset.py
@@ -24,7 +24,7 @@
 def test_dataset_init_success(location, expected_file_system, expected_source):
     timdex_dataset = TIMDEXDataset(location=location)
     assert isinstance(timdex_dataset.filesystem, expected_file_system)
-    assert timdex_dataset.source == expected_source
+    assert timdex_dataset.paths == expected_source
 
 
 def test_dataset_init_env_vars_set_config(monkeypatch, local_dataset_location):
@@ -79,8 +79,7 @@ def test_dataset_load_s3_sets_filesystem_and_dataset_success(
     timdex_dataset = TIMDEXDataset(location="s3://bucket/path/to/dataset")
     result = timdex_dataset.load()
 
-    mock_get_s3_fs.assert_called_once()
-    mock_pyarrow_ds.assert_called_once_with(
+    mock_pyarrow_ds.assert_called_with(
         "bucket/path/to/dataset",
         schema=timdex_dataset.schema,
         format="parquet",
@@ -137,6 +136,26 @@ def test_dataset_load_with_multi_nonpartition_filters_success(fixed_local_datase
     assert fixed_local_dataset.row_count == 1
 
 
+def test_dataset_load_current_records_all_sources_success(dataset_with_runs_location):
+    timdex_dataset = TIMDEXDataset(dataset_with_runs_location)
+
+    # 16 total parquet files, with current_records=False we get them all
+    timdex_dataset.load(current_records=False)
+    assert len(timdex_dataset.dataset.files) == 16
+
+    # 16 total parquet files, with current_records=True we only get 12 for current runs
+    timdex_dataset.load(current_records=True)
+    assert len(timdex_dataset.dataset.files) == 12
+
+
+def test_dataset_load_current_records_one_source_success(dataset_with_runs_location):
+    timdex_dataset = TIMDEXDataset(dataset_with_runs_location)
+    timdex_dataset.load(current_records=True, source="alma")
+
+    # 7 total parquet files for source, only 6 related to current runs
+    assert len(timdex_dataset.dataset.files) == 6
+
+
 def test_dataset_get_filtered_dataset_with_single_nonpartition_success(
     fixed_local_dataset,
 ):
@@ -324,3 +343,57 @@ def test_dataset_local_dataset_row_count_missing_dataset_raise_error(local_datas
     td = TIMDEXDataset(location="path/to/nowhere")
     with pytest.raises(DatasetNotLoadedError):
         _ = td.row_count
+
+
+def test_dataset_all_records_not_current_and_not_deduped(local_dataset_with_runs):
+    local_dataset_with_runs.load()
+    all_records_df = local_dataset_with_runs.read_dataframe()
+
+    # assert counts reflect all records from dataset, no deduping
+    assert all_records_df.source.value_counts().to_dict() == {"alma": 254, "dspace": 194}
+
+    # assert run_date min/max dates align with min/max for all runs
+    assert all_records_df.run_date.min() == date(2024, 12, 1)
+    assert all_records_df.run_date.max() == date(2025, 2, 5)
+
+
+def test_dataset_all_current_records_deduped(local_dataset_with_runs):
+    local_dataset_with_runs.load(current_records=True)
+    all_records_df = local_dataset_with_runs.read_dataframe()
+
+    # assert both sources have accurate record counts for current records only
+    assert all_records_df.source.value_counts().to_dict() == {"dspace": 90, "alma": 100}
+
+    # assert only one "full" run, per source
+    assert len(all_records_df[all_records_df.run_type == "full"].run_id.unique()) == 2
+
+    # assert run_date min/max dates align with both sources min/max dates
+    assert all_records_df.run_date.min() == date(2025, 1, 1)  # both
+    assert all_records_df.run_date.max() == date(2025, 2, 5)  # dspace
+
+
+def test_dataset_source_current_records_deduped(local_dataset_with_runs):
+    local_dataset_with_runs.load(current_records=True, source="alma")
+    alma_records_df = local_dataset_with_runs.read_dataframe()
+
+    # assert only alma records present and correct count
+    assert alma_records_df.source.value_counts().to_dict() == {"alma": 100}
+
+    # assert only one "full" run
+    assert len(alma_records_df[alma_records_df.run_type == "full"].run_id.unique()) == 1
+
+    # assert run_date min/max dates are correct for single source
+    assert alma_records_df.run_date.min() == date(2025, 1, 1)
+    assert alma_records_df.run_date.max() == date(2025, 1, 5)
+
+
+def test_dataset_all_read_methods_get_deduplication(
+    local_dataset_with_runs,
+):
+    local_dataset_with_runs.load(current_records=True, source="alma")
+
+    full_df = local_dataset_with_runs.read_dataframe()
+    all_records = list(local_dataset_with_runs.read_dicts_iter())
+    transformed_records = list(local_dataset_with_runs.read_transformed_records_iter())
+
+    assert len(full_df) == len(all_records) == len(transformed_records)
diff --git a/tests/test_runs.py b/tests/test_runs.py
index 0867906..0372749 100644
--- a/tests/test_runs.py
+++ b/tests/test_runs.py
@@ -56,14 +56,27 @@ def test_timdex_run_manager_get_runs_df(timdex_run_manager):
     assert runs_df.source.value_counts().to_dict() == {"alma": 7, "dspace": 7}
 
 
+def test_timdex_run_manager_get_all_current_run_parquet_files_success(
+    timdex_run_manager,
+):
+    ordered_parquet_files = timdex_run_manager.get_current_parquet_files()
+
+    # assert 12 parquet files, despite being 14 total for ALL sources
+    # this represents the last full run and all daily since
+    assert len(ordered_parquet_files) == 12
+
+    # assert sorted reverse chronologically
+    assert "year=2025/month=01/day=01" in ordered_parquet_files[-1]
+
+
 def test_timdex_run_manager_get_source_current_run_parquet_files_success(
     timdex_run_manager,
 ):
-    ordered_parquet_files = timdex_run_manager.get_current_source_parquet_files("alma")
+    ordered_parquet_files = timdex_run_manager._get_current_source_parquet_files("alma")
 
-    # assert 6 parquet files, despite being 8 total for alma
+    # assert 6 parquet files, despite being 8 total for 'alma' source
     # this represents the last full run and all daily since
-    assert len(ordered_parquet_files)
+    assert len(ordered_parquet_files) == 6
 
     # assert sorted reverse chronologically
     assert "year=2025/month=01/day=05" in ordered_parquet_files[0]
diff --git a/timdex_dataset_api/__init__.py b/timdex_dataset_api/__init__.py
index 13cefd2..cb56a60 100644
--- a/timdex_dataset_api/__init__.py
+++ b/timdex_dataset_api/__init__.py
@@ -3,7 +3,7 @@
 from timdex_dataset_api.dataset import TIMDEXDataset
 from timdex_dataset_api.record import DatasetRecord
 
-__version__ = "1.0.0"
+__version__ = "2.0.0"
 
 __all__ = [
     "DatasetRecord",
diff --git a/timdex_dataset_api/dataset.py b/timdex_dataset_api/dataset.py
index 43f6ccd..f5710d8 100644
--- a/timdex_dataset_api/dataset.py
+++ b/timdex_dataset_api/dataset.py
@@ -20,6 +20,7 @@
 
 from timdex_dataset_api.config import configure_logger
 from timdex_dataset_api.exceptions import DatasetNotLoadedError
+from timdex_dataset_api.run import TIMDEXRunManager
 
 if TYPE_CHECKING:
     from timdex_dataset_api.record import DatasetRecord  # pragma: nocover
@@ -114,11 +115,12 @@ def __init__(
         self.location = location
         self.config = config or TIMDEXDatasetConfig()
 
-        self.filesystem, self.source = self.parse_location(self.location)
+        self.filesystem, self.paths = self.parse_location(self.location)
         self.dataset: ds.Dataset = None  # type: ignore[assignment]
         self.schema = TIMDEX_DATASET_SCHEMA
         self.partition_columns = TIMDEX_DATASET_PARTITION_COLUMNS
         self._written_files: list[ds.WrittenFile] = None  # type: ignore[assignment]
+        self._dedupe_on_read: bool = False
 
     @property
     def row_count(self) -> int:
@@ -129,6 +131,8 @@ def row_count(self) -> int:
 
     def load(
         self,
+        *,
+        current_records: bool = False,
         **filters: Unpack[DatasetFilters],
     ) -> None:
         """Lazy load a pyarrow.dataset.Dataset and set to self.dataset.
@@ -152,14 +156,24 @@ def load(
         """
         start_time = time.perf_counter()
 
-        # load dataset
-        self.dataset = ds.dataset(
-            self.source,
-            schema=self.schema,
-            format="parquet",
-            partitioning="hive",
-            filesystem=self.filesystem,
-        )
+        # reset paths from original location before load
+        _, self.paths = self.parse_location(self.location)
+
+        # perform initial load of full dataset
+        self._load_pyarrow_dataset()
+
+        # if current_records flag set, limit to parquet files associated with current runs
+        self._dedupe_on_read = current_records
+        if current_records:
+            timdex_run_manager = TIMDEXRunManager(timdex_dataset=self)
+
+            # update paths, limiting by source if set
+            self.paths = timdex_run_manager.get_current_parquet_files(
+                source=filters.get("source")
+            )
+
+            # reload pyarrow dataset
+            self._load_pyarrow_dataset()
 
         # filter dataset
         self.dataset = self._get_filtered_dataset(**filters)
@@ -169,6 +183,16 @@ def load(
             f"{round(time.perf_counter()-start_time, 2)}s"
         )
 
+    def _load_pyarrow_dataset(self) -> None:
+        """Load the pyarrow dataset per local filesystem and paths attributes."""
+        self.dataset = ds.dataset(
+            self.paths,
+            schema=self.schema,
+            format="parquet",
+            partitioning="hive",
+            filesystem=self.filesystem,
+        )
+
     def _get_filtered_dataset(
         self,
         **filters: Unpack[DatasetFilters],
@@ -345,7 +369,8 @@ def write(
         start_time = time.perf_counter()
         self._written_files = []
 
-        if isinstance(self.source, list):
+        dataset_filesystem, dataset_path = self.parse_location(self.location)
+        if isinstance(dataset_path, list):
             raise TypeError(
                 "Dataset location must be the root of a single dataset for writing"
             )
@@ -354,10 +379,10 @@ def write(
 
         ds.write_dataset(
             record_batches_iter,
-            base_dir=self.source,
+            base_dir=dataset_path,
             basename_template="%s-{i}.parquet" % (str(uuid.uuid4())),  # noqa: UP031
             existing_data_behavior="overwrite_or_ignore",
-            filesystem=self.filesystem,
+            filesystem=dataset_filesystem,
             file_visitor=lambda written_file: self._written_files.append(written_file),  # type: ignore[arg-type]
             format="parquet",
             max_open_files=500,
@@ -444,14 +469,55 @@ def read_batches_iter(
                 "Dataset is not loaded. Please call the `load` method first."
             )
         dataset = self._get_filtered_dataset(**filters)
-        for batch in dataset.to_batches(
+
+        batches = dataset.to_batches(
             columns=columns,
             batch_size=self.config.read_batch_size,
             batch_readahead=self.config.batch_read_ahead,
             fragment_readahead=self.config.fragment_read_ahead,
-        ):
+        )
+
+        if self._dedupe_on_read:
+            yield from self._yield_deduped_batches(batches)
+        else:
+            for batch in batches:
+                if len(batch) > 0:
+                    yield batch
+
+    def _yield_deduped_batches(
+        self, batches: Iterator[pa.RecordBatch]
+    ) -> Iterator[pa.RecordBatch]:
+        """Method to yield record deduped batches.
+
+        Extending the normal behavior of yielding batches untouched, this method keeps
+        track of seen timdex_record_id's, yielding them only once.  For this method to
+        yield the most current version of a record -- most common usage -- it is required
+        that the batches are pre-ordered so the most recent record version is encountered
+        first.
+        """
+        seen_records = set()
+        for batch in batches:
             if len(batch) > 0:
-                yield batch
+                # init list of batch indices for records unseen
+                unseen_batch_indices = []
+
+                # get list of timdex ids from batch
+                timdex_ids = batch.column("timdex_record_id").to_pylist()
+
+                # check each record id and track unseen ones
+                for i, record_id in enumerate(timdex_ids):
+                    if record_id not in seen_records:
+                        unseen_batch_indices.append(i)
+                        seen_records.add(record_id)
+
+                # if all records from batch were seen, continue
+                if not unseen_batch_indices:
+                    continue
+
+                # else, yield unseen records from batch
+                deduped_batch = batch.take(pa.array(unseen_batch_indices))  # type: ignore[arg-type]
+                if len(deduped_batch) > 0:
+                    yield deduped_batch
 
     def read_dataframes_iter(
         self,
@@ -513,13 +579,14 @@ def read_transformed_records_iter(
     ) -> 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.
+        If 'transformed_record' is None (common scenarios are action="skip"|"error"), the
+        yield statement will not be executed for the row.  Note that for action="delete" a
+        transformed record still may be yielded if present.
 
         Args: see self.read_batches_iter()
         """
         for record_dict in self.read_dicts_iter(
-            columns=["transformed_record"],
+            columns=["timdex_record_id", "transformed_record"],
             **filters,
         ):
             if transformed_record := record_dict["transformed_record"]:
diff --git a/timdex_dataset_api/run.py b/timdex_dataset_api/run.py
index 0e11536..8080295 100644
--- a/timdex_dataset_api/run.py
+++ b/timdex_dataset_api/run.py
@@ -85,7 +85,22 @@ def get_runs_metadata(self, *, refresh: bool = False) -> pd.DataFrame:
         )
         return grouped_runs_df
 
-    def get_current_source_parquet_files(self, source: str) -> list[str]:
+    def get_current_parquet_files(self, source: str | None = None) -> list[str]:
+        """Get reverse chronological list of parquet files associated with current runs.
+
+        Args:
+            source: if provided, limits parquet files to only that source
+        """
+        runs_df = self.get_runs_metadata()  # run metadata is cached for future calls
+        sources = [source] if source else list(runs_df.source.unique())
+
+        source_parquet_files = []
+        for _source in sources:
+            source_parquet_files.extend(self._get_current_source_parquet_files(_source))
+
+        return source_parquet_files
+
+    def _get_current_source_parquet_files(self, source: str) -> list[str]:
         """Get reverse chronological list of current parquet files for a source.
 
         Args:
@@ -166,8 +181,9 @@ def _parse_run_metadata_from_parquet_file(self, parquet_filepath: str) -> dict:
         """
         parquet_file = pq.ParquetFile(
             parquet_filepath,
-            filesystem=self.timdex_dataset.filesystem,  # type: ignore[union-attr]
+            filesystem=self.timdex_dataset.filesystem,
         )
+
         file_meta = parquet_file.metadata.to_dict()
         num_rows = file_meta["num_rows"]
         columns_meta = file_meta["row_groups"][0]["columns"]  # type: ignore[typeddict-item]


Reviews:
user: login: ehanson8
body: Looking good and works as expected, just a few questions!
state: COMMENTED
author_association: NONE
submitted_at: 2025-05-22T13:30:12Z
user: login: ghukill
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T14:32:03Z
user: login: ghukill
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T14:39:46Z
user: login: ghukill
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T14:41:20Z
user: login: ehanson8
state: COMMENTED
author_association: NONE
submitted_at: 2025-05-22T14:49:25Z
user: login: ehanson8
state: APPROVED
author_association: NONE
submitted_at: 2025-05-22T14:49:51Z
user: login: jonavellecuerdo
body: One small request that will likely be addressed in the follow-up PR!
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T19:29:45Z
user: login: ghukill
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T19:37:02Z

Review comments:
diff_hunk: @@ -137,6 +136,22 @@ def test_dataset_load_with_multi_nonpartition_filters_success(fixed_local_datase
     assert fixed_local_dataset.row_count == 1
 
 
+def test_dataset_load_current_records_all_sources_success(dataset_with_runs_location):
+    timdex_dataset = TIMDEXDataset(dataset_with_runs_location)
+    timdex_dataset.load(current_records=True)
+
+    # 14 total parquet files, only 12 related to current runs
+    assert len(timdex_dataset.dataset.files) == 12
path: tests/test_dataset.py
user: login: ehanson8
body: Optional: you could really illustrate functionality this by including as well:

```
timdex_dataset.load(current_records=False)
assert len(timdex_dataset.dataset.files) == 14
```
created_at: 2025-05-22T12:56:12Z
author_association: NONE
diff_hunk: @@ -152,14 +155,23 @@ def load(
         """
         start_time = time.perf_counter()
 
-        # load dataset
-        self.dataset = ds.dataset(
-            self.source,
-            schema=self.schema,
-            format="parquet",
-            partitioning="hive",
-            filesystem=self.filesystem,
-        )
+        # reset paths from original location before load
+        _, self.paths = self.parse_location(self.location)
+
+        # perform initial load of full dataset
+        self._load_pyarrow_dataset()
+
+        # if current_records flag set, limit to parquet files associated with current runs
+        if current_records:
+            timdex_run_manager = TIMDEXRunManager(timdex_dataset=self)
+
+            # if filters.source is set, further limit to only this source
+            source = filters.get("source")
+            if source:
+                self.paths = timdex_run_manager.get_current_source_parquet_files(source)
+            else:
+                self.paths = timdex_run_manager.get_current_parquet_files()
+            self._load_pyarrow_dataset()
 
         # filter dataset
         self.dataset = self._get_filtered_dataset(**filters)
path: timdex_dataset_api/dataset.py
user: login: ehanson8
body: Optional because this is likely opening a can of worms: it does feel odd that you're using the `source` filter separately in the `current_records` block when there is a `self._get_filtered_dataset(**filters)` afterwards. Is there a way to do the `self._get_filtered_dataset(**filters)` call prior to the `current_records` code?
created_at: 2025-05-22T13:12:48Z
author_association: NONE
diff_hunk: @@ -167,12 +167,12 @@ def load(
         if current_records:
             timdex_run_manager = TIMDEXRunManager(timdex_dataset=self)
 
-            # if filters.source is set, further limit to only this source
-            source = filters.get("source")
-            if source:
-                self.paths = timdex_run_manager.get_current_source_parquet_files(source)
-            else:
-                self.paths = timdex_run_manager.get_current_parquet_files()
+            # update paths, limiting by source if set
+            self.paths = timdex_run_manager.get_current_parquet_files(
+                source=filters.get("source")
+            )
+
+            # reload pyarrow dataset
path: timdex_dataset_api/dataset.py
user: login: ehanson8
body: Even though this block changed in this commit, still curious about your thoughts on my previous comment!
created_at: 2025-05-22T13:25:21Z
author_association: NONE
diff_hunk: @@ -137,6 +136,22 @@ def test_dataset_load_with_multi_nonpartition_filters_success(fixed_local_datase
     assert fixed_local_dataset.row_count == 1
 
 
+def test_dataset_load_current_records_all_sources_success(dataset_with_runs_location):
+    timdex_dataset = TIMDEXDataset(dataset_with_runs_location)
+    timdex_dataset.load(current_records=True)
+
+    # 14 total parquet files, only 12 related to current runs
+    assert len(timdex_dataset.dataset.files) == 12
path: tests/test_dataset.py
user: login: ghukill
body: Nice!  I like it.  _And_ when adding that, I realized my code comment was stale: it's actualyl 16 files.   Update coming that has this, and self-documents that way.  Good suggestion.
created_at: 2025-05-22T14:32:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -152,14 +155,23 @@ def load(
         """
         start_time = time.perf_counter()
 
-        # load dataset
-        self.dataset = ds.dataset(
-            self.source,
-            schema=self.schema,
-            format="parquet",
-            partitioning="hive",
-            filesystem=self.filesystem,
-        )
+        # reset paths from original location before load
+        _, self.paths = self.parse_location(self.location)
+
+        # perform initial load of full dataset
+        self._load_pyarrow_dataset()
+
+        # if current_records flag set, limit to parquet files associated with current runs
+        if current_records:
+            timdex_run_manager = TIMDEXRunManager(timdex_dataset=self)
+
+            # if filters.source is set, further limit to only this source
+            source = filters.get("source")
+            if source:
+                self.paths = timdex_run_manager.get_current_source_parquet_files(source)
+            else:
+                self.paths = timdex_run_manager.get_current_parquet_files()
+            self._load_pyarrow_dataset()
 
         # filter dataset
         self.dataset = self._get_filtered_dataset(**filters)
path: timdex_dataset_api/dataset.py
user: login: ghukill
body: Tis a can of worms!  But a good one.  When we think of grabbing records from the dataset, I kind of think of it in two major motions:

1. what parquet **files** do you want to limit to
2. of those parquet files, what **rows** do you want to limit to

The first time we use `filters["source"]` we are limiting what parquet files we get by passing that to `TIMDEXRunManager` to return a list of parquet files.  The second time we use it, it's limiting the rows we get back to only that source.  But as you're kind of alluding to, this has no effect: we have already gauranteed that all rows will be of that source by virtue of limiting the parquet files.

The approach here was leaning on  the `DatasetFilters` that can be passed around to load and read methods.  Given that `"source"` is already in there, it feels natural to use it.  The user may not be aware that, under the hood, their inclusion of `current_records` + `source` resulted in a pre-optimization of parquet file selection, and that the normal `dataset.filter(source=XYZ)` will have no effect.  But I think that's okay.  It's like a SQL optimization engine that will "allow" `WHERE` clauses that may have no effect if it can optimize around them silently for the user.

It's a _very_ good catch and comment, but hoping this reasoning explains why it's setup this way.

FWIW, we _could_ pop `"source"` from the dataset filters if we already used it to limit to the parquet files... but seeing as it's harmless, that felt unncessary.

This builds on something else subtle, but I suppose important: the `dataset.filter` operations are **lazy**, meaning they are only applied when actually looping through records.  So while we do perform `self.dataset = self._get_filtered_dataset(**filters)` that appears to filter `source` unncessarily, it's only applied when we perform reading and it's nearly instant and inconsequential at that time.
created_at: 2025-05-22T14:39:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -167,12 +167,12 @@ def load(
         if current_records:
             timdex_run_manager = TIMDEXRunManager(timdex_dataset=self)
 
-            # if filters.source is set, further limit to only this source
-            source = filters.get("source")
-            if source:
-                self.paths = timdex_run_manager.get_current_source_parquet_files(source)
-            else:
-                self.paths = timdex_run_manager.get_current_parquet_files()
+            # update paths, limiting by source if set
+            self.paths = timdex_run_manager.get_current_parquet_files(
+                source=filters.get("source")
+            )
+
+            # reload pyarrow dataset
path: timdex_dataset_api/dataset.py
user: login: ghukill
body: See response above!
created_at: 2025-05-22T14:41:20Z
author_association: CONTRIBUTOR
diff_hunk: @@ -152,14 +155,23 @@ def load(
         """
         start_time = time.perf_counter()
 
-        # load dataset
-        self.dataset = ds.dataset(
-            self.source,
-            schema=self.schema,
-            format="parquet",
-            partitioning="hive",
-            filesystem=self.filesystem,
-        )
+        # reset paths from original location before load
+        _, self.paths = self.parse_location(self.location)
+
+        # perform initial load of full dataset
+        self._load_pyarrow_dataset()
+
+        # if current_records flag set, limit to parquet files associated with current runs
+        if current_records:
+            timdex_run_manager = TIMDEXRunManager(timdex_dataset=self)
+
+            # if filters.source is set, further limit to only this source
+            source = filters.get("source")
+            if source:
+                self.paths = timdex_run_manager.get_current_source_parquet_files(source)
+            else:
+                self.paths = timdex_run_manager.get_current_parquet_files()
+            self._load_pyarrow_dataset()
 
         # filter dataset
         self.dataset = self._get_filtered_dataset(**filters)
path: timdex_dataset_api/dataset.py
user: login: ehanson8
body: That makes sense, thank you for the context!
created_at: 2025-05-22T14:49:25Z
author_association: NONE
diff_hunk: @@ -129,6 +131,8 @@ def row_count(self) -> int:
 
     def load(
         self,
+        *,
+        current_records: bool = False,
         **filters: Unpack[DatasetFilters],
     ) -> None:
path: timdex_dataset_api/dataset.py
user: login: jonavellecuerdo
body: Can you update the docstring for `TIMDEXDataset.load` to include a description for the new `current_records: bool` keyword arg?
created_at: 2025-05-23T15:01:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -129,6 +131,8 @@ def row_count(self) -> int:
 
     def load(
         self,
+        *,
+        current_records: bool = False,
         **filters: Unpack[DatasetFilters],
     ) -> None:
path: timdex_dataset_api/dataset.py
user: login: ghukill
body: It is coming in the followup PR!  Good catch.
created_at: 2025-05-23T19:37:02Z
author_association: CONTRIBUTOR

PR comments:
user: login: coveralls
created_at: 2025-05-21T18:29:54Z
body: ## Pull Request Test Coverage Report for [Build 15189612409](https://coveralls.io/builds/73800365)


### Details

* **42** of **43**   **(97.67%)**  changed or added relevant lines in **3** files are covered.
* No unchanged relevant lines lost coverage.
* Overall coverage increased (+**0.3%**) to **94.872%**

---

|  Changes Missing Coverage | Covered Lines | Changed/Added Lines | % |
| :-----|--------------|--------|---: |
| [timdex_dataset_api/dataset.py](https://coveralls.io/builds/73800365/source?filename=timdex_dataset_api%2Fdataset.py#L515) | 33 | 34 | 97.06%
<!-- | **Total:** | **42** | **43** | **97.67%** | -->


|  Totals | [![Coverage Status](https://coveralls.io/builds/73800365/badge)](https://coveralls.io/builds/73800365) |
| :-- | --: |
| Change from base [Build 15168616734](https://coveralls.io/builds/73782787): |  0.3% |
| Covered Lines: | 296 |
| Relevant Lines: | 312 |

---
##### 💛  - [Coveralls](https://coveralls.io)

author_association: NONE
user: login: ghukill
created_at: 2025-05-22T21:05:11Z
body: @jonavellecuerdo , @ehanson8 - I've stumbled on kind of an interesting wrinkle / edge case with this approach.  @ehanson8, in a way, it's related to your comment about `self.dataset = self._get_filtered_dataset(**filters)` after we've already filtered for the current parquet files.  It's relevant there, but also for any read method where do the same thing.

My preference would be to keep this PR as-is, which establishes some flags and methods that get us very close.  In fact, it would work as needed for TIM to refresh a source as long as TIM was careful with how it uses this new functionality.

But, I think there is room for `.load(current_records=True)` to be less surprising, even if additional filtering is performed after the initial load.  More on that in a follow-up PR, but wanted to note here in case it comes up during @jonavellecuerdo's review.

**The very short example of this edge case:**

  - there are 100 "current" records
  - the most recent "daily" run had 10 `action=delete` rows
  - therefore, you have 90 records with `action=index` and 10 records with `action=delete` yielded as "current" records

While those `action=delete` rows _are_ the current version of those records, we would _not_ see them in Opensearch, and this challenges our definition of "current" for those records.  Yes, it's the "current" version of the record from the dataset's POV, but it's not in Opensearch by virtue of that most recent version having `action=delete`.  

I would propose that when `load(current_records=True)` is used you:
- _always_ see only the most recent version of that record in the parquet dataset
- if filtering removes that particular version, the record is omitted entirely
author_association: CONTRIBUTOR
user: login: ehanson8
created_at: 2025-05-23T12:34:11Z
body: > I would propose that when `load(current_records=True)` is used you:
> 
> * _always_ see only the most recent version of that record in the parquet dataset
> * if filtering removes that particular version, the record is omitted entirely

I largely agree but did wonder if this could complicate troubleshooting if we can't find the ID of a particular record in the loaded records since it was deleted. However, I imagine `load(current_records=True)` will be primarily used for reindexing and not troubleshooting, so I think there is minimal risk if any risk at all
author_association: NONE
user: login: ghukill
created_at: 2025-05-23T13:13:46Z
body: > > I would propose that when `load(current_records=True)` is used you:
> > 
> > * _always_ see only the most recent version of that record in the parquet dataset
> > * if filtering removes that particular version, the record is omitted entirely
> 
> I largely agree but did wonder if this could complicate troubleshooting if we can't find the ID of a particular record in the loaded records since it was deleted. However, I imagine `load(current_records=True)` will be primarily used for reindexing and not troubleshooting, so I think there is minimal risk if any risk at all

Yeah, it's interesting.  I know that I probably overuse "ergonomics" but I think this is a really good example of it.  Functionality wise it's all pretty simple, but it's about matching public interfaces, internal arguments and naming conventions, and behavior with a user's expectations.  

My hope is that we could nudge any concept of "current records" towards 100% matching what we find in Opensearch.  So if `alma:1234` was most recently deleted from Opensearch as part of daily ETL run `run-99`, I see two concrete paths when you utilize `current_records=True`:

1. without any other filters, you encounter `alma:1234, action=delete, run_id=run-99` which is the form of ETL record that caused the delete
2. you apply filtering that removes `action=delete`, or `run_id=run-99`, or maybe `source=libguides`, you should **not** see this record at all

I think the implementation details for this have a _little_ complexity, but I think the more intuitive ergonomics of it will be worth it.
author_association: CONTRIBUTOR
user: login: ghukill
created_at: 2025-05-23T13:20:19Z
body: @jonavellecuerdo, @ehanson8  - building on [my comment above about an edge case](https://github.com/MITLibraries/timdex-dataset-api/pull/144#issuecomment-2902570292), if interested, here is a unit test I'm working on locally that gets at it.

My goal is to work through a solution for this test, which currently **fails**, and have that be the basis of a followup PR.

```python
def test_dataset_current_records_accurate_with_additional_filters(new_local_dataset):
    """Simulate ETL runs with the following flow:

        1. full run-1 establishes 5 records
            - current state = 5 to-index
        2. daily run-2 updates the first 5 and adds 5 new ones
            - current state = 10 to-index
        3. daily run-3 has 3 deletes
            a. current state = 7 to-index, 3 to-delete

    The interesting records here are alma:0|1|2 where the most recent version is
    action="delete".  If during a read method we filtered action="index", we would expect
    that we don't see alma:0|2 with action="delete".  However, we should *also* not see
    alma:0|1|2 with action="index" from run-2, as these are not "current" versions of the
    record.

    This test ensures that even if additional filtering is applied after
    load(current_records=True) is performed, we do not encounter non-current forms of a
    record.
    """
    run_params = [
        (5, "alma", "2025-01-01", "full", "index", "run-1"),
        (10, "alma", "2025-01-02", "daily", "index", "run-2"),
        (3, "alma", "2025-01-03", "daily", "delete", "run-3"),
    ]
    for params in run_params:
        num_records, source, run_date, run_type, action, run_id = params
        records = generate_sample_records(
            num_records,
            timdex_record_id_prefix=source,
            source=source,
            run_date=run_date,
            run_type=run_type,
            action=action,
            run_id=run_id,
        )
        new_local_dataset.write(records)

    # without any additional filtering, confirm current records are accurate
    new_local_dataset.load(current_records=True)
    assert new_local_dataset.read_dataframe().action.value_counts().to_dict() == {
        "index": 7,  # run-1 and run-2 established 10 records
        "delete": 3,  # run-3 had 3 deletes
    }

    # apply action="index" filtering on read, assert that alma:0|1|2 are not present
    new_local_dataset.load(current_records=True)
    current_to_index = new_local_dataset.read_dataframe(action="index")

    assert len(current_to_index) == 7
    for timdex_record_id in ["alma:0", "alma:1", "alma:2"]:
        assert timdex_record_id not in list(current_to_index.timdex_record_id)
```
author_association: CONTRIBUTOR
user: login: ehanson8
created_at: 2025-05-23T13:23:22Z
body: Makes sense to me!
author_association: NONE