Add dataset writing functionality

Why these changes are being introduced:

A primary use case for the library will be Transmogrifier writing new
records to the parquet dataset.  This library is intended to make
that work simple for Transmogrifier, where all it needs to do is
yield DatasetRecords (imported from this lib) to the write method.

How this addresses that need:
* Adds new entrypoint write() method
* Includes helper methods to batch records yielded to write method
* Adds DatasetRecord class that is designed to encapsulat each
record (row) that will get written to the dataset.

Side effects of this change:
* Library supports writing to local or remote dataset

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

Author: ghukill

pyproject.toml

@@ -86,6 +86,7 @@ ignore = [
     "D103",
     "D104",
     "D415",
+    "D417",
     "EM102",
     "G004",
     "PLR0912",

tests/conftest.py

@@ -1,7 +1,11 @@
 """tests/conftest.py"""
 
+# ruff: noqa: D205, D209
+
+
 import pytest
 
+from tests.utils import generate_sample_records
 from timdex_dataset_api import TIMDEXDataset
 
 
@@ -22,3 +26,37 @@ def local_dataset_location():
 @pytest.fixture
 def local_dataset(local_dataset_location):
     return TIMDEXDataset.load(local_dataset_location)
+
+
+@pytest.fixture
+def new_temp_dataset(tmp_path) -> TIMDEXDataset:
+    location = str(tmp_path / "new_dataset")
+    return TIMDEXDataset(location=location)
+
+
+@pytest.fixture
+def small_records_iter():
+    """Simulates an iterator of X number of valid DatasetRecord instances."""
+
+    def _records_iter(num_records):
+        return generate_sample_records(num_records)
+
+    return _records_iter
+
+
+@pytest.fixture
+def small_records_iter_without_partitions():
+    """Simulates an iterator of X number of DatasetRecord instances WITHOUT partition
+    values included."""
+
+    def _records_iter(num_records):
+        return generate_sample_records(
+            num_records,
+            source=None,
+            run_date=None,
+            run_type=None,
+            action=None,
+            run_id=None,
+        )
+
+    return _records_iter

tests/test_dataset.py

@@ -1,5 +1,4 @@
 # ruff: noqa: S105, S106, SLF001
-
 from unittest.mock import MagicMock, patch
 
 import pyarrow as pa
@@ -33,7 +32,7 @@
     ],
 )
 @patch("timdex_dataset_api.dataset.TIMDEXDataset.get_s3_filesystem")
-def test_parse_location_single_local_directory(
+def test_parse_location_success_scenarios(
     get_s3_filesystem,
     location,
     expected_filesystem,
@@ -45,6 +44,28 @@ def test_parse_location_single_local_directory(
     assert source == expected_source
 
 
+@pytest.mark.parametrize(
+    ("location", "expected_exception"),
+    [
+        # None is invalid location type
+        (None, TypeError),
+        # mixed local and S3 locations
+        (
+            [
+                "/local/path/to/dataset/records.parquet",
+                "s3://path/to/dataset/records.parquet",
+            ],
+            ValueError,
+        ),
+    ],
+)
+@patch("timdex_dataset_api.dataset.TIMDEXDataset.get_s3_filesystem")
+def test_parse_location_error_scenarios(get_s3_filesystem, location, expected_exception):
+    get_s3_filesystem.return_value = fs.S3FileSystem()
+    with pytest.raises(expected_exception):
+        _ = TIMDEXDataset.parse_location(location)
+
+
 def test_get_s3_filesystem_success(mocker):
     mocked_s3_filesystem = mocker.spy(fs, "S3FileSystem")
     s3_filesystem = TIMDEXDataset.get_s3_filesystem()

tests/test_dataset_write.py

@@ -0,0 +1,181 @@
+# ruff: noqa: S105, S106, SLF001, PLR2004, PD901, D209, D205
+
+import datetime
+import math
+import os
+
+import pyarrow.dataset as ds
+import pytest
+
+from timdex_dataset_api.dataset import (
+    MAX_ROWS_PER_FILE,
+    TIMDEX_DATASET_SCHEMA,
+    DatasetNotLoadedError,
+    TIMDEXDataset,
+)
+from timdex_dataset_api.record import DatasetRecord
+
+
+def test_dataset_record_serialization():
+    dataset_record = DatasetRecord(
+        timdex_record_id="alma:123",
+        source_record=b"<record><title>Hello World.</title></record>",
+        transformed_record=b"""{"title":["Hello World."]}""",
+    )
+    assert dataset_record.to_dict() == {
+        "timdex_record_id": "alma:123",
+        "source_record": b"<record><title>Hello World.</title></record>",
+        "transformed_record": b"""{"title":["Hello World."]}""",
+        "source": None,
+        "run_date": None,
+        "run_type": None,
+        "action": None,
+        "run_id": None,
+    }
+
+
+def test_dataset_record_serialization_with_partition_values_provided():
+    dataset_record = DatasetRecord(
+        timdex_record_id="alma:123",
+        source_record=b"<record><title>Hello World.</title></record>",
+        transformed_record=b"""{"title":["Hello World."]}""",
+    )
+    partition_values = {
+        "source": "alma",
+        "run_date": "2024-12-01",
+        "run_type": "daily",
+        "action": "index",
+        "run_id": "000-111-aaa-bbb",
+    }
+    assert dataset_record.to_dict(partition_values=partition_values) == {
+        "timdex_record_id": "alma:123",
+        "source_record": b"<record><title>Hello World.</title></record>",
+        "transformed_record": b"""{"title":["Hello World."]}""",
+        "source": "alma",
+        "run_date": "2024-12-01",
+        "run_type": "daily",
+        "action": "index",
+        "run_id": "000-111-aaa-bbb",
+    }
+
+
+def test_dataset_write_records_to_new_dataset(new_temp_dataset, small_records_iter):
+    files_written = new_temp_dataset.write(small_records_iter(10_000))
+    assert len(files_written) == 1
+    assert os.path.exists(new_temp_dataset.location)
+
+    # load newly created dataset as new TIMDEXDataset instance
+    dataset = TIMDEXDataset.load(new_temp_dataset.location)
+    assert dataset.row_count == 10_000
+
+
+def test_dataset_reload_after_write(new_temp_dataset, small_records_iter):
+    files_written = new_temp_dataset.write(small_records_iter(10_000))
+    assert len(files_written) == 1
+    assert os.path.exists(new_temp_dataset.location)
+
+    # attempt row count before reload
+    with pytest.raises(DatasetNotLoadedError):
+        _ = new_temp_dataset.row_count
+
+    # attempt row count after reload
+    new_temp_dataset.reload()
+    assert new_temp_dataset.row_count == 10_000
+
+
+def test_dataset_write_default_max_rows_per_file(new_temp_dataset, small_records_iter):
+    """Default is 100k rows per file, therefore writing 200,033 records should result in
+    3 files (x2 @ 100k rows, x1 @ 33 rows)."""
+    total_records = 200_033
+
+    new_temp_dataset.write(small_records_iter(total_records))
+    new_temp_dataset.reload()
+
+    assert new_temp_dataset.row_count == total_records
+    assert len(new_temp_dataset.dataset.files) == math.ceil(
+        total_records / MAX_ROWS_PER_FILE
+    )
+
+
+def test_dataset_write_record_batches_uses_batch_size(
+    new_temp_dataset, small_records_iter
+):
+    total_records = 101
+    batch_size = 50
+    batches = list(
+        new_temp_dataset.get_dataset_record_batches(
+            small_records_iter(total_records), batch_size=batch_size
+        )
+    )
+    assert len(batches) == math.ceil(total_records / batch_size)
+
+
+def test_dataset_write_to_multiple_locations_raise_error(small_records_iter):
+    timdex_dataset = TIMDEXDataset(
+        location=["/path/to/records-1.parquet", "/path/to/records-2.parquet"]
+    )
+    with pytest.raises(
+        TypeError,
+        match="Dataset location must be the root of a single dataset for writing",
+    ):
+        timdex_dataset.write(small_records_iter(10))
+
+
+def test_dataset_write_mixin_partition_values_used(
+    new_temp_dataset, small_records_iter_without_partitions
+):
+    partition_values = {
+        "source": "alma",
+        "run_date": "2024-12-01",
+        "run_type": "daily",
+        "action": "index",
+        "run_id": "000-111-aaa-bbb",
+    }
+    _written_files = new_temp_dataset.write(
+        small_records_iter_without_partitions(10),
+        partition_values=partition_values,
+    )
+    new_temp_dataset.reload()
+
+    # load as pandas dataframe and assert column values
+    df = new_temp_dataset.dataset.to_table().to_pandas()
+    row = df.iloc[0]
+    assert row.source == partition_values["source"]
+    assert row.run_date == datetime.date(2024, 12, 1)
+    assert row.run_type == partition_values["run_type"]
+    assert row.action == partition_values["action"]
+    assert row.action == partition_values["action"]
+
+
+def test_dataset_write_schema_partitions_correctly_ordered(
+    new_temp_dataset, small_records_iter
+):
+    written_files = new_temp_dataset.write(
+        small_records_iter(10),
+        partition_values={
+            "source": "alma",
+            "run_date": "2024-12-01",
+            "run_type": "daily",
+            "action": "index",
+            "run_id": "000-111-aaa-bbb",
+        },
+    )
+    file = written_files[0]
+    assert (
+        "/source=alma/run_date=2024-12-01/run_type=daily"
+        "/action=index/run_id=000-111-aaa-bbb" in file.path
+    )
+
+
+def test_dataset_write_schema_applied_to_dataset(new_temp_dataset, small_records_iter):
+    new_temp_dataset.write(small_records_iter(10))
+
+    # manually load dataset to confirm schema without TIMDEXDataset projecting schema
+    # during load
+    dataset = ds.dataset(
+        new_temp_dataset.location,
+        format="parquet",
+        partitioning="hive",
+    )
+
+    assert set(dataset.schema.names) == set(TIMDEX_DATASET_SCHEMA.names)

tests/utils.py

@@ -0,0 +1,60 @@
+"""tests/utils.py"""
+
+# ruff: noqa: S311
+
+import random
+import uuid
+from collections.abc import Iterator
+
+from timdex_dataset_api import DatasetRecord
+
+
+def generate_sample_records(
+    num_records: int,
+    timdex_record_id_prefix: str = "alma",
+    source: str | None = "alma",
+    run_date: str | None = "2024-12-01",
+    run_type: str | None = "daily",
+    action: str | None = "index",
+    run_id: str | None = None,
+) -> Iterator[DatasetRecord]:
+    """Generate sample DatasetRecords."""
+    if not run_id:
+        run_id = str(uuid.uuid4())
+
+    for x in range(num_records):
+        yield DatasetRecord(
+            timdex_record_id=f"{timdex_record_id_prefix}:{x}",
+            source_record=b"<record><title>Hello World.</title></record>",
+            transformed_record=b"""{"title":["Hello World."]}""",
+            source=source,
+            run_date=run_date,
+            run_type=run_type,
+            action=action,
+            run_id=run_id,
+        )
+
+
+def generate_sample_records_with_simulated_partitions(
+    num_records: int, num_run_ids: int = 4
+) -> Iterator[DatasetRecord]:
+    """Generate sample DatasetRecords, with simulated sampling of partitions."""
+    sources = ["alma", "dspsace", "aspace", "libguides", "gismit", "gisogm"]
+    run_dates = ["2024-01-01", "2024-06-15", "2024-12-31"]
+    run_types = ["full", "daily"]
+    actions = ["index", "delete"]
+    run_ids = [str(uuid.uuid4()) for x in range(num_run_ids)]
+
+    records_remaining = num_records
+    while records_remaining > 0:
+        batch_size = random.randint(1, min(100, records_remaining))
+        yield from generate_sample_records(
+            num_records=batch_size,
+            timdex_record_id_prefix=random.choice(sources),
+            source=random.choice(sources),
+            run_date=random.choice(run_dates),
+            run_type=random.choice(run_types),
+            action=random.choice(actions),
+            run_id=random.choice(run_ids),
+        )
+        records_remaining -= batch_size

timdex_dataset_api/__init__.py

@@ -1,9 +1,11 @@
 """timdex_dataset_api/__init__.py"""
 
 from timdex_dataset_api.dataset import TIMDEXDataset
+from timdex_dataset_api.record import DatasetRecord
 
 __version__ = "0.1.0"
 
 __all__ = [
+    "DatasetRecord",
     "TIMDEXDataset",
 ]