Refactor to configuration object

Why these changes are being introduced:

With the addition of two read configurations that would be passed
around beteween multiple methods, the dataset module tipped
to where a centralized configuration object would be helpful.

Additionally, we have learned that per-operation configurations are
rare, and much more likely to be set once during TIMDEXDataset
init, or even as env vars for the duration of the library import.

How this addresses that need:

Creates a dataclass TIMDEXDatasetConfig that is passed to
TIMDEXDataset on init.  This class provides a typed object, with sensible
defaults, that are shared throughout all read and write methods.

Side effects of this change:
* None

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

Author: ghukill

tests/test_dataset.py

@@ -1,4 +1,4 @@
-# ruff: noqa: S105, S106, SLF001
+# ruff: noqa: S105, S106, SLF001, PLR2004
 import os
 from datetime import date
 from unittest.mock import MagicMock, patch
@@ -7,7 +7,11 @@
 import pytest
 from pyarrow import fs
 
-from timdex_dataset_api.dataset import DatasetNotLoadedError, TIMDEXDataset
+from timdex_dataset_api.dataset import (
+    DatasetNotLoadedError,
+    TIMDEXDataset,
+    TIMDEXDatasetConfig,
+)
 
 
 @pytest.mark.parametrize(
@@ -23,6 +27,24 @@ def test_dataset_init_success(location, expected_file_system, expected_source):
     assert timdex_dataset.source == expected_source
 
 
+def test_dataset_init_env_vars_set_config(monkeypatch, local_dataset_location):
+    default_timdex_dataset = TIMDEXDataset(location=local_dataset_location)
+    default_read_batch_config = default_timdex_dataset.config.read_batch_size
+    assert default_read_batch_config == 1_000
+
+    monkeypatch.setenv("TDA_READ_BATCH_SIZE", "100_000")
+    env_var_timdex_dataset = TIMDEXDataset(location=local_dataset_location)
+    env_var_read_batch_config = env_var_timdex_dataset.config.read_batch_size
+    assert env_var_read_batch_config == 100_000
+
+
+def test_dataset_init_custom_config_object(monkeypatch, local_dataset_location):
+    config = TIMDEXDatasetConfig()
+    config.max_rows_per_file = 42
+    timdex_dataset = TIMDEXDataset(location=local_dataset_location, config=config)
+    assert timdex_dataset.config.max_rows_per_file == 42
+
+
 @patch("timdex_dataset_api.dataset.fs.LocalFileSystem")
 @patch("timdex_dataset_api.dataset.ds.dataset")
 def test_dataset_load_local_sets_filesystem_and_dataset_success(
@@ -73,28 +95,28 @@ def test_dataset_load_without_filters_success(fixed_local_dataset):
     fixed_local_dataset.load()
 
     assert os.path.exists(fixed_local_dataset.location)
-    assert fixed_local_dataset.row_count == 5_000  # noqa: PLR2004
+    assert fixed_local_dataset.row_count == 5_000
 
 
 def test_dataset_load_with_run_date_str_filters_success(fixed_local_dataset):
     fixed_local_dataset.load(run_date="2024-12-01")
 
     assert os.path.exists(fixed_local_dataset.location)
-    assert fixed_local_dataset.row_count == 5_000  # noqa: PLR2004
+    assert fixed_local_dataset.row_count == 5_000
 
 
 def test_dataset_load_with_run_date_obj_filters_success(fixed_local_dataset):
     fixed_local_dataset.load(run_date=date(2024, 12, 1))
 
     assert os.path.exists(fixed_local_dataset.location)
-    assert fixed_local_dataset.row_count == 5_000  # noqa: PLR2004
+    assert fixed_local_dataset.row_count == 5_000
 
 
 def test_dataset_load_with_ymd_filters_success(fixed_local_dataset):
     fixed_local_dataset.load(year="2024", month="12", day="01")
 
     assert os.path.exists(fixed_local_dataset.location)
-    assert fixed_local_dataset.row_count == 5_000  # noqa: PLR2004
+    assert fixed_local_dataset.row_count == 5_000
 
 
 def test_dataset_load_with_single_nonpartition_filters_success(fixed_local_dataset):
@@ -158,7 +180,7 @@ def test_dataset_get_filtered_dataset_with_or_nonpartition_filters_success(
         timdex_record_id=["alma:0", "alma:1"]
     )
     filtered_local_df = filtered_local_dataset.to_table().to_pandas()
-    assert len(filtered_local_df) == 2  # noqa: PLR2004
+    assert len(filtered_local_df) == 2
     assert filtered_local_df["timdex_record_id"].tolist() == ["alma:0", "alma:1"]
 
 

tests/test_write.py

@@ -8,7 +8,6 @@
 
 from tests.utils import generate_sample_records
 from timdex_dataset_api.dataset import (
-    MAX_ROWS_PER_FILE,
     TIMDEX_DATASET_SCHEMA,
     TIMDEXDataset,
 )
@@ -28,28 +27,29 @@ def test_dataset_write_records_to_new_local_dataset(
 def test_dataset_write_default_max_rows_per_file(new_local_dataset, sample_records_iter):
     """Default is 100k rows per file, therefore writing 200,033 records should result in
     3 files (x2 @ 100k rows, x1 @ 33 rows)."""
+    default_max_rows_per_file = new_local_dataset.config.max_rows_per_file
     total_records = 200_033
 
     new_local_dataset.write(sample_records_iter(total_records))
     new_local_dataset.load()
 
     assert new_local_dataset.row_count == total_records
     assert len(new_local_dataset.dataset.files) == math.ceil(
-        total_records / MAX_ROWS_PER_FILE
+        total_records / default_max_rows_per_file
     )
 
 
 def test_dataset_write_record_batches_uses_batch_size(
     new_local_dataset, sample_records_iter
 ):
     total_records = 101
-    batch_size = 50
+    new_local_dataset.config.write_batch_size = 50
     batches = list(
-        new_local_dataset.create_record_batches(
-            sample_records_iter(total_records), batch_size=batch_size
-        )
+        new_local_dataset.create_record_batches(sample_records_iter(total_records))
+    )
+    assert len(batches) == math.ceil(
+        total_records / new_local_dataset.config.write_batch_size
     )
-    assert len(batches) == math.ceil(total_records / batch_size)
 
 
 def test_dataset_write_to_multiple_locations_raise_error(sample_records_iter):

timdex_dataset_api/dataset.py

@@ -3,9 +3,11 @@
 import itertools
 import json
 import operator
+import os
 import time
 import uuid
 from collections.abc import Iterator
+from dataclasses import dataclass, field
 from datetime import UTC, date, datetime
 from functools import reduce
 from typing import TYPE_CHECKING, TypedDict, Unpack
@@ -61,27 +63,57 @@ class DatasetFilters(TypedDict, total=False):
     day: str | None
 
 
-DEFAULT_BATCH_SIZE = 1_000
-MAX_ROWS_PER_GROUP = DEFAULT_BATCH_SIZE
-MAX_ROWS_PER_FILE = 100_000
-DEFAULT_BATCH_READ_AHEAD = 0
-DEFAULT_FRAGMENT_READ_AHEAD = 0
+@dataclass
+class TIMDEXDatasetConfig:
+    """Configurations for dataset operations.
 
+    - read_batch_size: row size of batches read, affecting memory consumption
+    - write_batch_size: row size of batches written, directly affecting row group size in
+        final parquet files
+    - max_rows_per_group: max number of rows per row group in a parquet file
+    - max_rows_per_file: max number of rows in a single parquet file
+    - batch_read_ahead: number of batches to optimistically read ahead when batch reading
+        from a dataset; pyarrow default is 16
+    - fragment_read_ahead: number of fragments to optimistically read ahead when batch
+        reaching from a dataset; pyarrow default is 4
+    """
 
-def strict_date_parse(date_string: str) -> date:
-    return datetime.strptime(date_string, "%Y-%m-%d").astimezone(UTC).date()
+    read_batch_size: int = field(
+        default_factory=lambda: int(os.getenv("TDA_READ_BATCH_SIZE", "1_000"))
+    )
+    write_batch_size: int = field(
+        default_factory=lambda: int(os.getenv("TDA_WRITE_BATCH_SIZE", "1_000"))
+    )
+    max_rows_per_group: int = field(
+        default_factory=lambda: int(os.getenv("TDA_MAX_ROWS_PER_GROUP", "1_000"))
+    )
+    max_rows_per_file: int = field(
+        default_factory=lambda: int(os.getenv("TDA_MAX_ROWS_PER_FILE", "100_000"))
+    )
+    batch_read_ahead: int = field(
+        default_factory=lambda: int(os.getenv("TDA_BATCH_READ_AHEAD", "0"))
+    )
+    fragment_read_ahead: int = field(
+        default_factory=lambda: int(os.getenv("TDA_FRAGMENT_READ_AHEAD", "0"))
+    )
 
 
 class TIMDEXDataset:
 
-    def __init__(self, location: str | list[str]):
+    def __init__(
+        self,
+        location: str | list[str],
+        config: TIMDEXDatasetConfig | None = None,
+    ):
         """Initialize TIMDEXDataset object.
 
         Args:
             location (str | list[str]): Local filesystem path or an S3 URI to
                 a parquet dataset. For partitioned datasets, set to the base directory.
         """
         self.location = location
+        self.config = config or TIMDEXDatasetConfig()
+
         self.filesystem, self.source = self.parse_location(self.location)
         self.dataset: ds.Dataset = None  # type: ignore[assignment]
         self.schema = TIMDEX_DATASET_SCHEMA
@@ -171,7 +203,7 @@ def _get_filtered_dataset(
 
         # create filter expressions for element-wise equality comparisons
         expressions = []
-        for field, value in filters.items():
+        for field, value in filters.items():  # noqa: F402
             if isinstance(value, list):
                 expressions.append(ds.field(field).isin(value))
             else:
@@ -207,7 +239,7 @@ def _parse_date_filters(self, run_date: str | date | None) -> DatasetFilters:
             DatasetFilters[dict]: values for run_date, year, month, and day
         """
         if isinstance(run_date, str):
-            run_date_obj = strict_date_parse(run_date)
+            run_date_obj = datetime.strptime(run_date, "%Y-%m-%d").astimezone(UTC).date()
         elif isinstance(run_date, date):
             run_date_obj = run_date
         else:
@@ -286,7 +318,6 @@ def write(
         self,
         records_iter: Iterator["DatasetRecord"],
         *,
-        batch_size: int = DEFAULT_BATCH_SIZE,
         use_threads: bool = True,
     ) -> list[ds.WrittenFile]:
         """Write records to the TIMDEX parquet dataset.
@@ -309,8 +340,6 @@ def write(
 
         Args:
             - records_iter: Iterator of DatasetRecord instances
-            - batch_size: size for batches to yield and write, directly affecting row
-                group size in final parquet files
             - use_threads: boolean if threads should be used for writing
         """
         start_time = time.perf_counter()
@@ -321,10 +350,7 @@ def write(
                 "Dataset location must be the root of a single dataset for writing"
             )
 
-        record_batches_iter = self.create_record_batches(
-            records_iter,
-            batch_size=batch_size,
-        )
+        record_batches_iter = self.create_record_batches(records_iter)
 
         ds.write_dataset(
             record_batches_iter,
@@ -335,8 +361,8 @@ def write(
             file_visitor=lambda written_file: self._written_files.append(written_file),  # type: ignore[arg-type]
             format="parquet",
             max_open_files=500,
-            max_rows_per_file=MAX_ROWS_PER_FILE,
-            max_rows_per_group=MAX_ROWS_PER_GROUP,
+            max_rows_per_file=self.config.max_rows_per_file,
+            max_rows_per_group=self.config.max_rows_per_group,
             partitioning=self.partition_columns,
             partitioning_flavor="hive",
             schema=self.schema,
@@ -349,8 +375,6 @@ def write(
     def create_record_batches(
         self,
         records_iter: Iterator["DatasetRecord"],
-        *,
-        batch_size: int = DEFAULT_BATCH_SIZE,
     ) -> Iterator[pa.RecordBatch]:
         """Yield pyarrow.RecordBatches for writing.
 
@@ -361,10 +385,10 @@ def create_record_batches(
 
         Args:
             - records_iter: Iterator of DatasetRecord instances
-            - batch_size: size for batches to yield and write, directly affecting row
-                group size in final parquet files
         """
-        for i, record_batch in enumerate(itertools.batched(records_iter, batch_size)):
+        for i, record_batch in enumerate(
+            itertools.batched(records_iter, self.config.write_batch_size)
+        ):
             batch = pa.RecordBatch.from_pylist(
                 [record.to_dict() for record in record_batch]
             )
@@ -395,9 +419,6 @@ def log_write_statistics(self, start_time: float) -> None:
     def read_batches_iter(
         self,
         columns: list[str] | None = None,
-        batch_size: int = DEFAULT_BATCH_SIZE,
-        batch_read_ahead: int = DEFAULT_BATCH_READ_AHEAD,
-        fragment_read_ahead: int = DEFAULT_FRAGMENT_READ_AHEAD,
         **filters: Unpack[DatasetFilters],
     ) -> Iterator[pa.RecordBatch]:
         """Yield pyarrow.RecordBatches from the dataset.
@@ -416,7 +437,7 @@ def read_batches_iter(
                 number will increase RAM usage but could also improve IO utilization.
                 Pyarrow default is 4, but this library defaults to 0 to prioritize memory
                 footprint.
-            - filter_kwargs: pairs of column:value to filter the dataset
+            - filters: pairs of column:value to filter the dataset
         """
         if not self.dataset:
             raise DatasetNotLoadedError(
@@ -425,19 +446,16 @@ def read_batches_iter(
         dataset = self._get_filtered_dataset(**filters)
         for batch in dataset.to_batches(
             columns=columns,
-            batch_size=batch_size,
-            batch_readahead=batch_read_ahead,
-            fragment_readahead=fragment_read_ahead,
+            batch_size=self.config.read_batch_size,
+            batch_readahead=self.config.batch_read_ahead,
+            fragment_readahead=self.config.fragment_read_ahead,
         ):
             if len(batch) > 0:
                 yield batch
 
     def read_dataframes_iter(
         self,
         columns: list[str] | None = None,
-        batch_size: int = DEFAULT_BATCH_SIZE,
-        batch_read_ahead: int = DEFAULT_BATCH_READ_AHEAD,
-        fragment_read_ahead: int = DEFAULT_FRAGMENT_READ_AHEAD,
         **filters: Unpack[DatasetFilters],
     ) -> Iterator[pd.DataFrame]:
         """Yield record batches as Pandas DataFrames from the dataset.
@@ -446,19 +464,13 @@ def read_dataframes_iter(
         """
         for record_batch in self.read_batches_iter(
             columns=columns,
-            batch_size=batch_size,
-            batch_read_ahead=batch_read_ahead,
-            fragment_read_ahead=fragment_read_ahead,
             **filters,
         ):
             yield record_batch.to_pandas()
 
     def read_dataframe(
         self,
         columns: list[str] | None = None,
-        batch_size: int = DEFAULT_BATCH_SIZE,
-        batch_read_ahead: int = DEFAULT_BATCH_READ_AHEAD,
-        fragment_read_ahead: int = DEFAULT_FRAGMENT_READ_AHEAD,
         **filters: Unpack[DatasetFilters],
     ) -> pd.DataFrame | None:
         """Yield record batches as Pandas DataFrames and concatenate to single dataframe.
@@ -473,9 +485,6 @@ def read_dataframe(
             record_batch.to_pandas()
             for record_batch in self.read_batches_iter(
                 columns=columns,
-                batch_size=batch_size,
-                batch_read_ahead=batch_read_ahead,
-                fragment_read_ahead=fragment_read_ahead,
                 **filters,
             )
         ]
@@ -486,9 +495,6 @@ def read_dataframe(
     def read_dicts_iter(
         self,
         columns: list[str] | None = None,
-        batch_size: int = DEFAULT_BATCH_SIZE,
-        batch_read_ahead: int = DEFAULT_BATCH_READ_AHEAD,
-        fragment_read_ahead: int = DEFAULT_FRAGMENT_READ_AHEAD,
         **filters: Unpack[DatasetFilters],
     ) -> Iterator[dict]:
         """Yield individual record rows as dictionaries from the dataset.
@@ -497,18 +503,12 @@ def read_dicts_iter(
         """
         for record_batch in self.read_batches_iter(
             columns=columns,
-            batch_size=batch_size,
-            batch_read_ahead=batch_read_ahead,
-            fragment_read_ahead=fragment_read_ahead,
             **filters,
         ):
             yield from record_batch.to_pylist()
 
     def read_transformed_records_iter(
         self,
-        batch_size: int = DEFAULT_BATCH_SIZE,
-        batch_read_ahead: int = DEFAULT_BATCH_READ_AHEAD,
-        fragment_read_ahead: int = DEFAULT_FRAGMENT_READ_AHEAD,
         **filters: Unpack[DatasetFilters],
     ) -> Iterator[dict]:
         """Yield individual transformed records as dictionaries from the dataset.
@@ -520,9 +520,6 @@ def read_transformed_records_iter(
         """
         for record_dict in self.read_dicts_iter(
             columns=["transformed_record"],
-            batch_size=batch_size,
-            batch_read_ahead=batch_read_ahead,
-            fragment_read_ahead=fragment_read_ahead,
             **filters,
         ):
             if transformed_record := record_dict["transformed_record"]: