From 6f7525425cfff07694e294af3c8117797b6475d8 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Thu, 7 Aug 2025 16:45:40 -0400 Subject: [PATCH] Load pyarrow dataset on TIMDEXDataset init Why these changes are being introduced: As the TIMDEXDatasetMetadata becomes more integrated, there is less need to be explicit about how we load the pyarrow dataset. Formerly, the method .load() needed to be called manually and supported options like 'current_records' or 'include_parquet_files'. This also reflected a time when 'TIMDEXDataset.load()' suggested that "loading" was the pyarrow dataset only. With the introduction of metadata, it is also better to be specific we are loading a pyarrow dataset which is only one of many assets associated with a TIMDEXDataset instance. How this addresses that need: Renames .load() to .load_pyarrow_dataset() to be explicit about what is happening. We no longer store the pyarrow dataset filesystem or paths on self, as they are only used briefly during this dataset load. We can get them anytime via .dataset. Really most important, we limit the root 'location' that we init a TIMDEXDataset instance to be a string only, the root of the dataset. Now that we don't allow a list of strings at that level, we can trust the nature of self.location to be a string, and the root of the TIMDEX dataset. Side effects of this change: * TIMDEXDataset and TIMDEXDatasetMetadata can only be initialized with a string, which is the root of the TIMDEX dataset. From there, both know where their assets can be found. * You cannot "pre-filter" the pyarrow dataset when loading, which had confusing overlap with the read methods; the read methods themselves may change somewhat dramatically now that we have metadata to use. Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-533 --- README.md | 6 - pyproject.toml | 4 +- tests/conftest.py | 7 - tests/test_dataset.py | 431 ++++++------------------------- tests/test_read.py | 151 ++++++++++- tests/test_write.py | 34 +-- timdex_dataset_api/dataset.py | 307 +++++++++------------- timdex_dataset_api/exceptions.py | 4 - 8 files changed, 346 insertions(+), 598 deletions(-) diff --git a/README.md b/README.md index 7220e3e..f44bd1e 100644 --- a/README.md +++ b/README.md @@ -110,12 +110,6 @@ timdex_dataset = TIMDEXDataset("s3://my-bucket/path/to/dataset") # or, local dataset (e.g. testing or development) timdex_dataset = TIMDEXDataset("/path/to/dataset") - -# load the dataset, which discovers all parquet files -timdex_dataset.load() - -# or, load the dataset but ensure that only current records are ever yielded -timdex_dataset.load(current_records=True) ``` All read methods for `TIMDEXDataset` allow for the same group of filters which are defined in `timdex_dataset_api.dataset.DatasetFilters`. Examples are shown below. diff --git a/pyproject.toml b/pyproject.toml index 45c168d..d6caccc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,7 +54,7 @@ line-length = 90 [tool.mypy] disallow_untyped_calls = true disallow_untyped_defs = true -exclude = ["tests/", "output/"] +exclude = ["tests/", "output/", "migrations/"] [[tool.mypy.overrides]] module = [] @@ -95,6 +95,8 @@ ignore = [ "PLR0915", "S321", "S608", + "TD002", + "TD003", "TRY003" ] diff --git a/tests/conftest.py b/tests/conftest.py index 18f4a4f..8265145 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -82,7 +82,6 @@ def timdex_dataset(tmp_path, timdex_dataset_config) -> TIMDEXDataset: ), write_append_deltas=False, ) - dataset.load() return dataset @@ -110,8 +109,6 @@ def timdex_dataset_multi_source(tmp_path) -> TIMDEXDataset: ), write_append_deltas=False, ) - - dataset.load() return dataset @@ -165,8 +162,6 @@ def timdex_dataset_with_runs(tmp_path, timdex_dataset_config_small) -> TIMDEXDat ), write_append_deltas=False, ) - - dataset.load() return dataset @@ -202,8 +197,6 @@ def timdex_dataset_same_day_runs(tmp_path) -> TIMDEXDataset: ), write_append_deltas=False, ) - - dataset.load() return dataset diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 3d3ec1f..d93d7ce 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -11,44 +11,14 @@ from pyarrow import fs from timdex_dataset_api.dataset import ( - DatasetNotLoadedError, TIMDEXDataset, TIMDEXDatasetConfig, ) -@pytest.mark.parametrize( - ("location_param", "expected_file_system", "expected_source_param"), - [ - ( - "path/to/dataset", - fs.LocalFileSystem, - "path/to/dataset/data/records", - ), - ( - "s3://timdex/path/to/dataset", - fs.S3FileSystem, - "timdex/path/to/dataset/data/records", - ), - ], -) -def test_dataset_init_success( - location_param, - expected_file_system, - expected_source_param, - s3_bucket_mocked, - tmp_path, -): - location = location_param - expected_source = expected_source_param - - if not location.startswith("s3://"): - location = str(tmp_path / location) - expected_source = str(tmp_path / expected_source) - - timdex_dataset = TIMDEXDataset(location=location) - assert isinstance(timdex_dataset.filesystem, expected_file_system) - assert timdex_dataset.paths == expected_source +def test_dataset_init_success(tmp_path): + timdex_dataset = TIMDEXDataset(str(tmp_path / "path/to/dataset")) + assert isinstance(timdex_dataset.dataset.filesystem, fs.LocalFileSystem) def test_dataset_init_env_vars_set_config(monkeypatch, tmp_path): @@ -73,130 +43,109 @@ def test_dataset_init_custom_config_object(monkeypatch, tmp_path): @patch("timdex_dataset_api.dataset.fs.LocalFileSystem") @patch("timdex_dataset_api.dataset.ds.dataset") -def test_dataset_load_local_sets_filesystem_and_dataset_success( +def test_load_pyarrow_dataset_default_uses_data_records_root( mock_pyarrow_ds, mock_local_fs, tmp_path ): + """Ensure load_pyarrow_dataset() without args calls pyarrow.dataset with the + dataset's data_records_root path as the source and the proper filesystem.""" mock_local_fs.return_value = MagicMock() mock_pyarrow_ds.return_value = MagicMock() - location = str(Path(tmp_path) / "local/path/to/dataset") + location = str(Path(tmp_path) / "local/path/to/default_dataset") timdex_dataset = TIMDEXDataset(location=location) - result = timdex_dataset.load() + # call the explicit loader to exercise the code path + dataset_obj = timdex_dataset.load_pyarrow_dataset() - mock_pyarrow_ds.assert_called_once_with( + mock_pyarrow_ds.assert_called_with( f"{location}/data/records", schema=timdex_dataset.schema, format="parquet", partitioning="hive", filesystem=mock_local_fs.return_value, ) - + assert dataset_obj == mock_pyarrow_ds.return_value assert timdex_dataset.dataset == mock_pyarrow_ds.return_value - assert result is None -@patch("timdex_dataset_api.dataset.TIMDEXDataset.get_s3_filesystem") +@patch("timdex_dataset_api.dataset.fs.LocalFileSystem") @patch("timdex_dataset_api.dataset.ds.dataset") -def test_dataset_load_s3_sets_filesystem_and_dataset_success( - mock_pyarrow_ds, mock_get_s3_fs, s3_bucket_mocked +def test_load_pyarrow_dataset_with_parquet_files_list( + mock_pyarrow_ds, mock_local_fs, tmp_path ): - mock_get_s3_fs.return_value = MagicMock() + """Ensure load_pyarrow_dataset(parquet_files=...) passes the explicit list + of parquet files as the source to pyarrow.dataset.""" + mock_local_fs.return_value = MagicMock() mock_pyarrow_ds.return_value = MagicMock() - timdex_dataset = TIMDEXDataset(location="s3://timdex/path/to/dataset") - result = timdex_dataset.load() + location = str(Path(tmp_path) / "local/path/to/dataset_with_files") + + timdex_dataset = TIMDEXDataset(location=location) + + parquet_files = [ + f"{timdex_dataset.data_records_root}/source=alma/run_date=2024-12-01/part-0.parquet", + f"{timdex_dataset.data_records_root}/source=alma/run_date=2024-12-01/part-1.parquet", + ] + + dataset_obj = timdex_dataset.load_pyarrow_dataset(parquet_files=parquet_files) mock_pyarrow_ds.assert_called_with( - "timdex/path/to/dataset/data/records", + parquet_files, schema=timdex_dataset.schema, format="parquet", partitioning="hive", - filesystem=mock_get_s3_fs.return_value, + filesystem=mock_local_fs.return_value, ) + assert dataset_obj == mock_pyarrow_ds.return_value assert timdex_dataset.dataset == mock_pyarrow_ds.return_value - assert result is None - - -def test_dataset_load_without_filters_success(timdex_dataset_multi_source): - timdex_dataset_multi_source.load() - - assert os.path.exists(timdex_dataset_multi_source.location) - assert timdex_dataset_multi_source.row_count == 5_000 - - -def test_dataset_load_with_run_date_str_filters_success(timdex_dataset_multi_source): - timdex_dataset_multi_source.load(run_date="2024-12-01") - - assert os.path.exists(timdex_dataset_multi_source.location) - assert timdex_dataset_multi_source.row_count == 5_000 - - -def test_dataset_load_with_run_date_obj_filters_success(timdex_dataset_multi_source): - timdex_dataset_multi_source.load(run_date=date(2024, 12, 1)) - assert os.path.exists(timdex_dataset_multi_source.location) - assert timdex_dataset_multi_source.row_count == 5_000 - -def test_dataset_load_with_ymd_filters_success(timdex_dataset_multi_source): - timdex_dataset_multi_source.load(year="2024", month="12", day="01") - - assert os.path.exists(timdex_dataset_multi_source.location) - assert timdex_dataset_multi_source.row_count == 5_000 - - -def test_dataset_load_with_single_nonpartition_filters_success( - timdex_dataset_multi_source, +@patch("timdex_dataset_api.dataset.fs.LocalFileSystem") +@patch("timdex_dataset_api.dataset.ds.dataset") +def test_dataset_load_local_sets_filesystem_and_dataset_success( + mock_pyarrow_ds, mock_local_fs, tmp_path ): - timdex_dataset_multi_source.load(timdex_record_id="alma:0") + mock_local_fs.return_value = MagicMock() + mock_pyarrow_ds.return_value = MagicMock() - assert timdex_dataset_multi_source.row_count == 1 + location = str(Path(tmp_path) / "local/path/to/dataset") + timdex_dataset = TIMDEXDataset(location=location) -def test_dataset_load_with_multi_nonpartition_filters_success( - timdex_dataset_multi_source, -): - timdex_dataset_multi_source.load( - timdex_record_id="alma:0", - source="alma", - run_type="daily", - run_id="abc123", - action="index", + mock_pyarrow_ds.assert_called_once_with( + f"{location}/data/records", + schema=timdex_dataset.schema, + format="parquet", + partitioning="hive", + filesystem=mock_local_fs.return_value, ) - assert timdex_dataset_multi_source.row_count == 1 + assert timdex_dataset.dataset == mock_pyarrow_ds.return_value -@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") -def test_dataset_load_current_records_all_sources_success( - timdex_timdex_dataset_with_runs, +@patch("timdex_dataset_api.dataset.TIMDEXDataset.get_s3_filesystem") +@patch("timdex_dataset_api.dataset.ds.dataset") +def test_dataset_load_s3_sets_filesystem_and_dataset_success( + mock_pyarrow_ds, mock_get_s3_fs, s3_bucket_mocked ): - timdex_dataset = TIMDEXDataset(timdex_timdex_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 - + mock_get_s3_fs.return_value = MagicMock() + mock_pyarrow_ds.return_value = MagicMock() -@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") -def test_dataset_load_current_records_one_source_success(timdex_timdex_dataset_with_runs): - timdex_dataset = TIMDEXDataset(timdex_timdex_dataset_with_runs.location) - timdex_dataset.load(current_records=True, source="alma") + timdex_dataset = TIMDEXDataset(location="s3://timdex/path/to/dataset") - # 7 total parquet files for source, only 6 related to current runs - assert len(timdex_dataset.dataset.files) == 6 + mock_pyarrow_ds.assert_called_with( + "timdex/path/to/dataset/data/records", + schema=timdex_dataset.schema, + format="parquet", + partitioning="hive", + filesystem=mock_get_s3_fs.return_value, + ) + assert timdex_dataset.dataset == mock_pyarrow_ds.return_value def test_dataset_get_filtered_dataset_with_single_nonpartition_success( timdex_dataset_multi_source, ): - timdex_dataset_multi_source.load() # initial load dataset, no filters passed - filtered_timdex_dataset = timdex_dataset_multi_source._get_filtered_dataset( run_id="abc123", ) @@ -211,8 +160,6 @@ def test_dataset_get_filtered_dataset_with_single_nonpartition_success( def test_dataset_get_filtered_dataset_with_multi_nonpartition_filters_success( timdex_dataset_multi_source, ): - timdex_dataset_multi_source.load() # initial load dataset, no filters passed - filtered_timdex_dataset = timdex_dataset_multi_source._get_filtered_dataset( timdex_record_id="alma:0", source="alma", @@ -229,8 +176,6 @@ def test_dataset_get_filtered_dataset_with_multi_nonpartition_filters_success( def test_dataset_get_filtered_dataset_with_or_nonpartition_filters_success( timdex_dataset_multi_source, ): - timdex_dataset_multi_source.load() - filtered_timdex_dataset = timdex_dataset_multi_source._get_filtered_dataset( timdex_record_id=["alma:0", "alma:1"] ) @@ -242,8 +187,6 @@ def test_dataset_get_filtered_dataset_with_or_nonpartition_filters_success( def test_dataset_get_filtered_dataset_with_run_date_str_successs( timdex_dataset_multi_source, ): - timdex_dataset_multi_source.load() # initial load dataset, no filters passed - filtered_timdex_dataset = timdex_dataset_multi_source._get_filtered_dataset( run_date="2024-12-01" ) @@ -253,15 +196,16 @@ def test_dataset_get_filtered_dataset_with_run_date_str_successs( # timdex_dataset_multi_source consists of single 'run_date' value # therefore, filtered_timdex_dataset includes all records - assert filtered_timdex_dataset.count_rows() == timdex_dataset_multi_source.row_count + assert ( + filtered_timdex_dataset.count_rows() + == timdex_dataset_multi_source.dataset.count_rows() + ) assert empty_timdex_dataset.count_rows() == 0 def test_dataset_get_filtered_dataset_with_run_date_obj_success( timdex_dataset_multi_source, ): - timdex_dataset_multi_source.load() # initial load dataset, no filters passed - filtered_timdex_dataset = timdex_dataset_multi_source._get_filtered_dataset( run_date=date(2024, 12, 1) ) @@ -271,13 +215,14 @@ def test_dataset_get_filtered_dataset_with_run_date_obj_success( # timdex_dataset_multi_source consists of single 'run_date' value # therefore, filtered_timdex_dataset includes all records - assert filtered_timdex_dataset.count_rows() == timdex_dataset_multi_source.row_count + assert ( + filtered_timdex_dataset.count_rows() + == timdex_dataset_multi_source.dataset.count_rows() + ) assert empty_timdex_dataset.count_rows() == 0 def test_dataset_get_filtered_dataset_with_ymd_success(timdex_dataset_multi_source): - timdex_dataset_multi_source.load() # initial load dataset, no filters passed - filtered_timdex_dataset = timdex_dataset_multi_source._get_filtered_dataset( year="2024" ) @@ -285,15 +230,16 @@ def test_dataset_get_filtered_dataset_with_ymd_success(timdex_dataset_multi_sour # timdex_dataset_multi_source consists of single 'run_date' value # therefore, filtered_timdex_dataset includes all records - assert filtered_timdex_dataset.count_rows() == timdex_dataset_multi_source.row_count + assert ( + filtered_timdex_dataset.count_rows() + == timdex_dataset_multi_source.dataset.count_rows() + ) assert empty_timdex_dataset.count_rows() == 0 def test_dataset_get_filtered_dataset_with_run_date_invalid_raise_error( timdex_dataset_multi_source, ): - timdex_dataset_multi_source.load() # initial load dataset, no filters passed - with pytest.raises( TypeError, match=( @@ -317,107 +263,15 @@ def test_dataset_get_s3_filesystem_success(mocker): assert isinstance(s3_filesystem, pa._s3fs.S3FileSystem) -@pytest.mark.parametrize( - ("location_param", "expected_filesystem", "expected_source_param"), - [ - ("path/to/dataset", fs.LocalFileSystem, "path/to/dataset"), - ( - ["path/to/records1.parquet", "path/to/records2.parquet"], - fs.LocalFileSystem, - ["path/to/records1.parquet", "path/to/records2.parquet"], - ), - ("s3://bucket/path/to/dataset", fs.S3FileSystem, "bucket/path/to/dataset"), - ( - [ - "s3://bucket/path/to/dataset/records1.parquet", - "s3://bucket/path/to/dataset/records2.parquet", - ], - fs.S3FileSystem, - [ - "bucket/path/to/dataset/records1.parquet", - "bucket/path/to/dataset/records2.parquet", - ], - ), - ], -) -@patch("timdex_dataset_api.dataset.TIMDEXDataset.get_s3_filesystem") -def test_dataset_parse_location_success( - get_s3_filesystem, - location_param, - expected_filesystem, - expected_source_param, - tmp_path, -): - get_s3_filesystem.return_value = fs.S3FileSystem() - - location = location_param - expected_source = expected_source_param - - if isinstance(location, str) and not location.startswith("s3://"): - location = str(tmp_path / location) - expected_source = str(tmp_path / expected_source) - elif isinstance(location, list) and not location[0].startswith("s3://"): - location = [str(tmp_path / path) for path in location] - expected_source = [str(tmp_path / path) for path in expected_source] - - filesystem, source = TIMDEXDataset.parse_location(location) - assert isinstance(filesystem, expected_filesystem) - assert source == expected_source - - -@pytest.mark.parametrize( - ("location_param", "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_dataset_parse_location_error( - get_s3_filesystem, location_param, expected_exception, tmp_path -): - get_s3_filesystem.return_value = fs.S3FileSystem() - - location = location_param - if isinstance(location, list) and not all( - path.startswith("s3://") for path in location - ): - # Update the local path with tmp_path - location = [ - str(tmp_path / path) if not path.startswith("s3://") else path - for path in location - ] - - with pytest.raises(expected_exception): - _ = TIMDEXDataset.parse_location(location) - - def test_dataset_timdex_dataset_validate_success(timdex_dataset): assert timdex_dataset.dataset.to_table().validate() is None # where None is valid def test_dataset_timdex_dataset_row_count_success(timdex_dataset): - assert timdex_dataset.dataset.count_rows() == timdex_dataset.row_count - - -def test_dataset_timdex_dataset_row_count_missing_dataset_raise_error( - timdex_dataset, tmp_path -): - td = TIMDEXDataset(location=str(tmp_path / "path/to/nowhere")) - with pytest.raises(DatasetNotLoadedError): - _ = td.row_count + assert timdex_dataset.dataset.count_rows() == timdex_dataset.dataset.count_rows() def test_dataset_all_records_not_current_and_not_deduped(timdex_dataset_with_runs): - timdex_dataset_with_runs.load() all_records_df = timdex_dataset_with_runs.read_dataframe() # assert counts reflect all records from dataset, no deduping @@ -428,145 +282,6 @@ def test_dataset_all_records_not_current_and_not_deduped(timdex_dataset_with_run assert all_records_df.run_date.max() == date(2025, 2, 5) -@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") -def test_dataset_all_current_records_deduped(timdex_dataset_with_runs): - timdex_dataset_with_runs.load(current_records=True) - all_records_df = timdex_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 - - -@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") -def test_dataset_source_current_records_deduped(timdex_dataset_with_runs): - timdex_dataset_with_runs.load(current_records=True, source="alma") - alma_records_df = timdex_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) - - -@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") -def test_dataset_all_read_methods_get_deduplication( - timdex_dataset_with_runs, -): - timdex_dataset_with_runs.load(current_records=True, source="alma") - - full_df = timdex_dataset_with_runs.read_dataframe() - all_records = list(timdex_dataset_with_runs.read_dicts_iter()) - transformed_records = list(timdex_dataset_with_runs.read_transformed_records_iter()) - - assert len(full_df) == len(all_records) == len(transformed_records) - - -@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") -def test_dataset_current_records_no_additional_filtering_accurate_records_yielded( - timdex_dataset_with_runs, -): - timdex_dataset_with_runs.load(current_records=True, source="alma") - df = timdex_dataset_with_runs.read_dataframe() - assert df.action.value_counts().to_dict() == {"index": 99, "delete": 1} - - -@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") -def test_dataset_current_records_action_filtering_accurate_records_yielded( - timdex_dataset_with_runs, -): - timdex_dataset_with_runs.load(current_records=True, source="alma") - df = timdex_dataset_with_runs.read_dataframe(action="index") - assert df.action.value_counts().to_dict() == {"index": 99} - - -@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") -def test_dataset_current_records_index_filtering_accurate_records_yielded( - timdex_dataset_with_runs, -): - """This is a somewhat complex test, but demonstrates that only 'current' records - are yielded when .load(current_records=True) is applied. - - Given these runs from the fixture: - [ - ... - (25, "alma", "2025-01-03", "daily", "index", "run-5"), <---- filtered to - (10, "alma", "2025-01-04", "daily", "delete", "run-6"), <---- influences current - ... - ] - - Though we are filtering to run-5, which has 25 total records to-index, we see only 15 - records yielded. Why? This is because while we have filtered to only yield from - run-5, run-6 had 10 deletes which made records alma:0|9 no longer "current" in run-5. - As we yielded records reverse chronologically, the deletes from run-6 (alma:0-alma:9) - "influenced" what records we would see as we continue backwards in time. - """ - # with current_records=False, we get all 25 records from run-5 - timdex_dataset_with_runs.load(current_records=False, source="alma") - df = timdex_dataset_with_runs.read_dataframe(run_id="run-5") - assert len(df) == 25 - - # with current_records=True, we only get 15 records from run-5 - # because newer run-6 influenced what records are current for older run-5 - timdex_dataset_with_runs.load(current_records=True, source="alma") - df = timdex_dataset_with_runs.read_dataframe(run_id="run-5") - assert len(df) == 15 - assert list(df.timdex_record_id) == [ - "alma:10", - "alma:11", - "alma:12", - "alma:13", - "alma:14", - "alma:15", - "alma:16", - "alma:17", - "alma:18", - "alma:19", - "alma:20", - "alma:21", - "alma:22", - "alma:23", - "alma:24", - ] - - -@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") -def test_dataset_load_current_records_gets_correct_same_day_full_run( - timdex_dataset_same_day_runs, -): - """Two full runs were performed on the same day, but 'run-2' was performed most - recently. current_records=True should discover the more recent of the two 'run-2', - not 'run-1'.""" - timdex_dataset_same_day_runs.load(current_records=True, run_type="full") - df = timdex_dataset_same_day_runs.read_dataframe() - - assert list(df.run_id.unique()) == ["run-2"] - - -@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") -def test_dataset_load_current_records_gets_correct_same_day_daily_runs_ordering( - timdex_dataset_same_day_runs, -): - """Two runs were performed on 2025-01-02, but the most recent records should be from - run 'run-5' which are action='delete', not 'run-4' with action='index'.""" - timdex_dataset_same_day_runs.load(current_records=True, run_type="daily") - first_record = next(timdex_dataset_same_day_runs.read_dicts_iter()) - - assert first_record["run_id"] == "run-5" - assert first_record["action"] == "delete" - - def test_dataset_records_data_structure_is_idempotent(timdex_dataset_with_runs): assert os.path.exists(timdex_dataset_with_runs.data_records_root) start_file_count = glob.glob(f"{timdex_dataset_with_runs.data_records_root}/**/*") diff --git a/tests/test_read.py b/tests/test_read.py index 33f5197..0072aad 100644 --- a/tests/test_read.py +++ b/tests/test_read.py @@ -1,4 +1,6 @@ -# ruff: noqa: PLR2004 +# ruff: noqa: D205, D209, PLR2004 + +from datetime import date import pandas as pd import pyarrow as pa @@ -31,7 +33,7 @@ def test_read_batches_filter_columns(timdex_dataset_multi_source): def test_read_batches_no_filters_gets_full_dataset(timdex_dataset_multi_source): batches = timdex_dataset_multi_source.read_batches_iter() table = pa.Table.from_batches(batches) - assert len(table) == timdex_dataset_multi_source.row_count + assert len(table) == timdex_dataset_multi_source.dataset.count_rows() def test_read_batches_with_filters_gets_subset_of_dataset(timdex_dataset_multi_source): @@ -44,10 +46,10 @@ def test_read_batches_with_filters_gets_subset_of_dataset(timdex_dataset_multi_s table = pa.Table.from_batches(batches) assert len(table) == 1_000 - assert len(table) < timdex_dataset_multi_source.row_count + assert len(table) < timdex_dataset_multi_source.dataset.count_rows() # assert loaded dataset is unchanged by filtering for a read method - assert timdex_dataset_multi_source.row_count == 5_000 + assert timdex_dataset_multi_source.dataset.count_rows() == 5_000 def test_read_dataframe_batches_yields_dataframes(timdex_dataset_multi_source): @@ -62,7 +64,7 @@ def test_read_dataframe_reads_all_dataset_rows_after_filtering( ): df = timdex_dataset_multi_source.read_dataframe() assert isinstance(df, pd.DataFrame) - assert len(df) == timdex_dataset_multi_source.row_count + assert len(df) == timdex_dataset_multi_source.dataset.count_rows() def test_read_dicts_yields_dictionary_for_each_dataset_record( @@ -90,3 +92,142 @@ def test_read_transformed_records_yields_parsed_dictionary(timdex_dataset_multi_ transformed_record = next(batches) assert isinstance(transformed_record, dict) assert transformed_record == {"title": ["Hello World."]} + + +@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") +def test_dataset_all_current_records_deduped(timdex_dataset_with_runs): + timdex_dataset_with_runs.load(current_records=True) + all_records_df = timdex_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 + + +@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") +def test_dataset_source_current_records_deduped(timdex_dataset_with_runs): + timdex_dataset_with_runs.load(current_records=True, source="alma") + alma_records_df = timdex_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) + + +@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") +def test_dataset_all_read_methods_get_deduplication( + timdex_dataset_with_runs, +): + timdex_dataset_with_runs.load(current_records=True, source="alma") + + full_df = timdex_dataset_with_runs.read_dataframe() + all_records = list(timdex_dataset_with_runs.read_dicts_iter()) + transformed_records = list(timdex_dataset_with_runs.read_transformed_records_iter()) + + assert len(full_df) == len(all_records) == len(transformed_records) + + +@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") +def test_dataset_current_records_no_additional_filtering_accurate_records_yielded( + timdex_dataset_with_runs, +): + timdex_dataset_with_runs.load(current_records=True, source="alma") + df = timdex_dataset_with_runs.read_dataframe() + assert df.action.value_counts().to_dict() == {"index": 99, "delete": 1} + + +@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") +def test_dataset_current_records_action_filtering_accurate_records_yielded( + timdex_dataset_with_runs, +): + timdex_dataset_with_runs.load(current_records=True, source="alma") + df = timdex_dataset_with_runs.read_dataframe(action="index") + assert df.action.value_counts().to_dict() == {"index": 99} + + +@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") +def test_dataset_current_records_index_filtering_accurate_records_yielded( + timdex_dataset_with_runs, +): + """This is a somewhat complex test, but demonstrates that only 'current' records + are yielded when .load(current_records=True) is applied. + + Given these runs from the fixture: + [ + ... + (25, "alma", "2025-01-03", "daily", "index", "run-5"), <---- filtered to + (10, "alma", "2025-01-04", "daily", "delete", "run-6"), <---- influences current + ... + ] + + Though we are filtering to run-5, which has 25 total records to-index, we see only 15 + records yielded. Why? This is because while we have filtered to only yield from + run-5, run-6 had 10 deletes which made records alma:0|9 no longer "current" in run-5. + As we yielded records reverse chronologically, the deletes from run-6 (alma:0-alma:9) + "influenced" what records we would see as we continue backwards in time. + """ + # with current_records=False, we get all 25 records from run-5 + timdex_dataset_with_runs.load(current_records=False, source="alma") + df = timdex_dataset_with_runs.read_dataframe(run_id="run-5") + assert len(df) == 25 + + # with current_records=True, we only get 15 records from run-5 + # because newer run-6 influenced what records are current for older run-5 + timdex_dataset_with_runs.load(current_records=True, source="alma") + df = timdex_dataset_with_runs.read_dataframe(run_id="run-5") + assert len(df) == 15 + assert list(df.timdex_record_id) == [ + "alma:10", + "alma:11", + "alma:12", + "alma:13", + "alma:14", + "alma:15", + "alma:16", + "alma:17", + "alma:18", + "alma:19", + "alma:20", + "alma:21", + "alma:22", + "alma:23", + "alma:24", + ] + + +@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") +def test_dataset_load_current_records_gets_correct_same_day_full_run( + timdex_dataset_same_day_runs, +): + """Two full runs were performed on the same day, but 'run-2' was performed most + recently. current_records=True should discover the more recent of the two 'run-2', + not 'run-1'.""" + timdex_dataset_same_day_runs.load(current_records=True, run_type="full") + df = timdex_dataset_same_day_runs.read_dataframe() + + assert list(df.run_id.unique()) == ["run-2"] + + +@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") +def test_dataset_load_current_records_gets_correct_same_day_daily_runs_ordering( + timdex_dataset_same_day_runs, +): + """Two runs were performed on 2025-01-02, but the most recent records should be from + run 'run-5' which are action='delete', not 'run-4' with action='index'.""" + timdex_dataset_same_day_runs.load(current_records=True, run_type="daily") + first_record = next(timdex_dataset_same_day_runs.read_dicts_iter()) + + assert first_record["run_id"] == "run-5" + assert first_record["action"] == "delete" diff --git a/tests/test_write.py b/tests/test_write.py index 13b43c5..3710989 100644 --- a/tests/test_write.py +++ b/tests/test_write.py @@ -6,12 +6,10 @@ import pyarrow.dataset as ds import pyarrow.parquet as pq -import pytest from tests.utils import generate_sample_records from timdex_dataset_api.dataset import ( TIMDEX_DATASET_SCHEMA, - TIMDEXDataset, ) from timdex_dataset_api.metadata import ORDERED_METADATA_COLUMN_NAMES @@ -20,11 +18,10 @@ def test_dataset_write_records_to_timdex_dataset_empty( timdex_dataset_empty, sample_records_generator ): written_files = timdex_dataset_empty.write(sample_records_generator(10_000)) - timdex_dataset_empty.load() assert len(written_files) == 1 assert os.path.exists(timdex_dataset_empty.location) - assert timdex_dataset_empty.row_count == 10_000 + assert timdex_dataset_empty.dataset.count_rows() == 10_000 def test_dataset_write_default_max_rows_per_file( @@ -36,9 +33,8 @@ def test_dataset_write_default_max_rows_per_file( total_records = 200_033 timdex_dataset_empty.write(sample_records_generator(total_records)) - timdex_dataset_empty.load() - assert timdex_dataset_empty.row_count == total_records + assert timdex_dataset_empty.dataset.count_rows() == total_records assert len(timdex_dataset_empty.dataset.files) == math.ceil( total_records / default_max_rows_per_file ) @@ -59,20 +55,6 @@ def test_dataset_write_record_batches_uses_batch_size( ) -@pytest.mark.skip( - reason="Test unneeded soon when list[str] not supported for dataset location." -) -def test_dataset_write_to_multiple_locations_raise_error(sample_records_generator): - 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(sample_records_generator(10)) - - def test_dataset_write_schema_applied_to_dataset( timdex_dataset_empty, sample_records_generator ): @@ -103,20 +85,18 @@ def test_dataset_write_partition_for_multiple_sources( ): # perform write for source="alma" and run_date="2024-12-01" written_files_source_a = timdex_dataset_empty.write(sample_records_generator(10)) - timdex_dataset_empty.load() assert os.path.exists(written_files_source_a[0].path) - assert timdex_dataset_empty.row_count == 10 + assert timdex_dataset_empty.dataset.count_rows() == 10 # perform write for source="libguides" and run_date="2024-12-01" written_files_source_b = timdex_dataset_empty.write( generate_sample_records(num_records=7, source="libguides") ) - timdex_dataset_empty.load() assert os.path.exists(written_files_source_b[0].path) assert os.path.exists(written_files_source_a[0].path) - assert timdex_dataset_empty.row_count == 17 + assert timdex_dataset_empty.dataset.count_rows() == 17 def test_dataset_write_partition_ignore_existing_data( @@ -125,12 +105,11 @@ def test_dataset_write_partition_ignore_existing_data( # perform two (2) writes for source="alma" and run_date="2024-12-01" written_files_source_a0 = timdex_dataset_empty.write(sample_records_generator(10)) written_files_source_a1 = timdex_dataset_empty.write(sample_records_generator(10)) - timdex_dataset_empty.load() # assert that both files exist and no overwriting occurs assert os.path.exists(written_files_source_a0[0].path) assert os.path.exists(written_files_source_a1[0].path) - assert timdex_dataset_empty.row_count == 20 + assert timdex_dataset_empty.dataset.count_rows() == 20 @patch("timdex_dataset_api.dataset.uuid.uuid4") @@ -148,11 +127,10 @@ def test_dataset_write_partition_overwrite_files_with_same_name( # perform two (2) writes for source="alma" and run_date="2024-12-01" _ = timdex_dataset_empty.write(sample_records_generator(10)) written_files_source_a1 = timdex_dataset_empty.write(sample_records_generator(7)) - timdex_dataset_empty.load() # assert that only the second file exists and overwriting occurs assert os.path.exists(written_files_source_a1[0].path) - assert timdex_dataset_empty.row_count == 7 + assert timdex_dataset_empty.dataset.count_rows() == 7 def test_dataset_write_single_append_delta_success( diff --git a/timdex_dataset_api/dataset.py b/timdex_dataset_api/dataset.py index 1cb29a1..d7295d2 100644 --- a/timdex_dataset_api/dataset.py +++ b/timdex_dataset_api/dataset.py @@ -21,7 +21,6 @@ from pyarrow import fs from timdex_dataset_api.config import configure_logger -from timdex_dataset_api.exceptions import DatasetNotLoadedError from timdex_dataset_api.metadata import TIMDEXDatasetMetadata if TYPE_CHECKING: @@ -108,14 +107,13 @@ class TIMDEXDataset: def __init__( self, - location: str | list[str], + location: 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. + location (str ): Local filesystem path or an S3 URI to a parquet dataset. """ self.config = config or TIMDEXDatasetConfig() self.location = location @@ -123,17 +121,16 @@ def __init__( self.create_data_structure() # pyarrow dataset - self.filesystem, self.paths = self.parse_location(self.data_records_root) - self.dataset: ds.Dataset = None # type: ignore[assignment] self.schema = TIMDEX_DATASET_SCHEMA self.partition_columns = TIMDEX_DATASET_PARTITION_COLUMNS + self.dataset = self.load_pyarrow_dataset() # dataset metadata - self.metadata = TIMDEXDatasetMetadata(location) # type: ignore[arg-type] + self.metadata = TIMDEXDatasetMetadata(location) @property def location_scheme(self) -> Literal["file", "s3"]: - scheme = urlparse(self.location).scheme # type: ignore[arg-type] + scheme = urlparse(self.location).scheme if scheme == "": return "file" if scheme == "s3": @@ -152,147 +149,51 @@ def create_data_structure(self) -> None: exist_ok=True, ) - @property - def row_count(self) -> int: - """Get row count from loaded dataset.""" - if not self.dataset: - raise DatasetNotLoadedError - return self.dataset.count_rows() - - def load( - self, - **filters: Unpack[DatasetFilters], - ) -> None: - """Lazy load a pyarrow.dataset.Dataset and set to self.dataset. - - Loading is comprised of two main steps: - - - load: Lazily load full dataset. PyArrow will "discover" full dataset. - Note: This step may take a couple of seconds but leans on PyArrow's - parquet reading processes. - - filter: Lazily filter rows in the PyArrow dataset by conditions on - TIMDEX_DATASET_FILTER_COLUMNS. + def load_pyarrow_dataset(self, parquet_files: list[str] | None = None) -> ds.Dataset: + """Lazy load a pyarrow.dataset.Dataset. The dataset is loaded via the expected schema as defined by module constant TIMDEX_DATASET_SCHEMA. If the target dataset differs in any way, errors may be raised when reading or writing data. Args: - - filters: kwargs typed via DatasetFilters TypedDict - - Filters passed directly in method call, e.g. source="alma", - run_date="2024-12-20", etc., but are typed according to DatasetFilters. + parquet_files: explicit list of parquet files to construct pyarrow dataset """ start_time = time.perf_counter() - # reset paths from original location before load - _, self.paths = self.parse_location(self.data_records_root) + # get pyarrow filesystem and dataset path basesd on self.location + filesystem, path = self.parse_location(self.data_records_root) - # perform initial load of full dataset - self.dataset = self._load_pyarrow_dataset() - - # filter dataset - self.dataset = self._get_filtered_dataset(**filters) - - logger.info( - f"Dataset successfully loaded: '{self.data_records_root}', " - f"{round(time.perf_counter()-start_time, 2)}s" - ) + # set source for pyarrow dataset + source: str | list[str] = parquet_files or path - def _load_pyarrow_dataset(self) -> ds.Dataset: - """Load the pyarrow dataset per local filesystem and paths attributes.""" - return ds.dataset( - self.paths, + dataset = ds.dataset( + source, schema=self.schema, format="parquet", partitioning="hive", - filesystem=self.filesystem, + filesystem=filesystem, ) - def _get_filtered_dataset( - self, - **filters: Unpack[DatasetFilters], - ) -> ds.Dataset: - """Lazy filter self.dataset and return a new pyarrow Dataset object. - - This method will construct a single pyarrow.compute.Expression - that is combined from individual equality comparison predicates - using the provided filters. - - Args: - - filters: kwargs typed via DatasetFilters TypedDict - - Filters passed directly in method call, e.g. source="alma", - run_date="2024-12-20", etc., but are typed according to DatasetFilters. - - Raises: - DatasetNotLoadedError: Raised if `self.dataset` is None. - TIMDEXDataset.load must be called before any filter method calls. - ValueError: Raised if provided 'run_date' is an invalid type or - cannot be parsed. - - Returns: - ds.Dataset: Original pyarrow.dataset.Dataset (if no filters applied) - or new pyarrow.dataset.Dataset with applied filters. - """ - if not self.dataset: - raise DatasetNotLoadedError - - # if run_date provided, derive year, month, and day partition filters and set - if filters.get("run_date"): - filters.update(self._parse_date_filters(filters["run_date"])) - - # create filter expressions for element-wise equality comparisons - expressions = [] - for field, value in filters.items(): # noqa: F402 - if isinstance(value, list): - expressions.append(ds.field(field).isin(value)) - else: - expressions.append(ds.field(field) == value) - - # if filter expressions not found, return original dataset - if not expressions: - return self.dataset - - # combine filter expressions as a single predicate - combined_expressions = reduce(operator.and_, expressions) - logger.debug( - "Filtering dataset based on the following column-value pairs: " - f"{combined_expressions}" + logger.info( + f"Dataset successfully loaded: '{self.data_records_root}', " + f"{round(time.perf_counter()-start_time, 2)}s" ) - return self.dataset.filter(combined_expressions) - - def _parse_date_filters(self, run_date: str | date | None) -> DatasetFilters: - """Parse date filters from 'run_date'. + return dataset - Args: - run_date (str | date | None): If str, the value must match the - date format "%Y-%m-%d"; if date, ymd values are extracted - as str. - - Raises: - TypeError: Raised when 'run_date' is an invalid type. - ValueError: Raised when either a datetime.date object cannot be parsed - from a provided 'run_date' str. - - Returns: - DatasetFilters[dict]: values for run_date, year, month, and day - """ - if isinstance(run_date, str): - run_date_obj = datetime.strptime(run_date, "%Y-%m-%d").astimezone(UTC).date() - elif isinstance(run_date, date): - run_date_obj = run_date + def parse_location( + self, + location: str, + ) -> tuple[fs.FileSystem, str]: + """Parse and return a pyarrow filesystem and normalized parquet path(s).""" + if self.location_scheme == "s3": + filesystem = TIMDEXDataset.get_s3_filesystem() + source = location.removeprefix("s3://") else: - raise TypeError( - "Provided 'run_date' value must be a string matching format " - "'%Y-%m-%d' or a datetime.date." - ) - - return { - "run_date": run_date_obj, - "year": run_date_obj.strftime("%Y"), - "month": run_date_obj.strftime("%m"), - "day": run_date_obj.strftime("%d"), - } + filesystem = fs.LocalFileSystem() + source = location + return filesystem, source @staticmethod def get_s3_filesystem() -> fs.FileSystem: @@ -307,7 +208,7 @@ def get_s3_filesystem() -> fs.FileSystem: raise RuntimeError("Could not locate AWS credentials") if os.getenv("MINIO_S3_ENDPOINT_URL"): - return fs.S3FileSystem( + return fs.S3FileSystem( # pragma: nocover access_key=os.environ["MINIO_USERNAME"], secret_key=os.environ["MINIO_PASSWORD"], endpoint_override=os.environ["MINIO_S3_ENDPOINT_URL"], @@ -320,54 +221,6 @@ def get_s3_filesystem() -> fs.FileSystem: session_token=credentials.token, ) - # NOTE: WIP: this will be heavily reworked in upcoming .load() updates - @classmethod - def parse_location( - cls, - location: str | list[str], - ) -> tuple[fs.FileSystem, str | list[str]]: - """Parse and return the filesystem and normalized source location(s). - - Handles both single location strings and lists of Parquet file paths. - """ - match location: - case str(): - return cls._parse_single_location(location) - case list(): - return cls._parse_multiple_locations(location) - case _: - raise TypeError("Location type must be str or list[str].") - - # NOTE: WIP: these will be removed in upcoming .load() updates - @classmethod - def _parse_single_location( - cls, location: str - ) -> tuple[fs.FileSystem, str | list[str]]: - """Get filesystem and normalized location for single location.""" - if location.startswith("s3://"): - filesystem = TIMDEXDataset.get_s3_filesystem() - source = location.removeprefix("s3://") - else: - filesystem = fs.LocalFileSystem() - source = location - return filesystem, source - - # NOTE: WIP: these will be removed in upcoming .load() updates - @classmethod - def _parse_multiple_locations( - cls, location: list[str] - ) -> tuple[fs.FileSystem, str | list[str]]: - """Get filesystem and normalized location for multiple locations.""" - if all(loc.startswith("s3://") for loc in location): - filesystem = TIMDEXDataset.get_s3_filesystem() - source = [loc.removeprefix("s3://") for loc in location] - elif all(not loc.startswith("s3://") for loc in location): - filesystem = fs.LocalFileSystem() - source = location - else: - raise ValueError("Mixed S3 and local paths are not supported.") - return filesystem, source - def write( self, records_iter: Iterator["DatasetRecord"], @@ -402,20 +255,16 @@ def write( start_time = time.perf_counter() written_files: list[ds.WrittenFile] = [] - dataset_filesystem, dataset_path = self.parse_location(self.data_records_root) - if isinstance(dataset_path, list): - raise TypeError( - "Dataset location must be the root of a single dataset for writing" - ) + filesystem, path = self.parse_location(self.data_records_root) # write ETL parquet records record_batches_iter = self.create_record_batches(records_iter) ds.write_dataset( record_batches_iter, - base_dir=dataset_path, + base_dir=path, basename_template="%s-{i}.parquet" % (str(uuid.uuid4())), # noqa: UP031 existing_data_behavior="overwrite_or_ignore", - filesystem=dataset_filesystem, + filesystem=filesystem, file_visitor=lambda written_file: written_files.append(written_file), # type: ignore[arg-type] format="parquet", max_open_files=500, @@ -427,6 +276,9 @@ def write( use_threads=use_threads, ) + # refresh dataset files + self.dataset = self.load_pyarrow_dataset() + # write metadata append deltas if write_append_deltas: for written_file in written_files: @@ -478,6 +330,87 @@ def log_write_statistics( f"total size: {total_size}" ) + def _get_filtered_dataset( + self, + **filters: Unpack[DatasetFilters], + ) -> ds.Dataset: + """Lazy filter self.dataset and return a new pyarrow Dataset object. + + This method will construct a single pyarrow.compute.Expression + that is combined from individual equality comparison predicates + using the provided filters. + + Args: + - filters: kwargs typed via DatasetFilters TypedDict + - Filters passed directly in method call, e.g. source="alma", + run_date="2024-12-20", etc., but are typed according to DatasetFilters. + + Raises: + ValueError: Raised if provided 'run_date' is an invalid type or + cannot be parsed. + + Returns: + ds.Dataset: Original pyarrow.dataset.Dataset (if no filters applied) + or new pyarrow.dataset.Dataset with applied filters. + """ + # if run_date provided, derive year, month, and day partition filters and set + if filters.get("run_date"): + filters.update(self._parse_date_filters(filters["run_date"])) + + # create filter expressions for element-wise equality comparisons + expressions = [] + for field, value in filters.items(): # noqa: F402 + if isinstance(value, list): + expressions.append(ds.field(field).isin(value)) + else: + expressions.append(ds.field(field) == value) + + # if filter expressions not found, return original dataset + if not expressions: + return self.dataset + + # combine filter expressions as a single predicate + combined_expressions = reduce(operator.and_, expressions) + logger.debug( + "Filtering dataset based on the following column-value pairs: " + f"{combined_expressions}" + ) + + return self.dataset.filter(combined_expressions) + + def _parse_date_filters(self, run_date: str | date | None) -> DatasetFilters: + """Parse date filters from 'run_date'. + + Args: + run_date (str | date | None): If str, the value must match the + date format "%Y-%m-%d"; if date, ymd values are extracted + as str. + + Raises: + TypeError: Raised when 'run_date' is an invalid type. + ValueError: Raised when either a datetime.date object cannot be parsed + from a provided 'run_date' str. + + Returns: + DatasetFilters[dict]: values for run_date, year, month, and day + """ + if isinstance(run_date, str): + run_date_obj = datetime.strptime(run_date, "%Y-%m-%d").astimezone(UTC).date() + elif isinstance(run_date, date): + run_date_obj = run_date + else: + raise TypeError( + "Provided 'run_date' value must be a string matching format " + "'%Y-%m-%d' or a datetime.date." + ) + + return { + "run_date": run_date_obj, + "year": run_date_obj.strftime("%Y"), + "month": run_date_obj.strftime("%m"), + "day": run_date_obj.strftime("%d"), + } + def read_batches_iter( self, columns: list[str] | None = None, @@ -492,10 +425,6 @@ def read_batches_iter( - columns: list[str], list of columns to return from the dataset - filters: 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) batches = dataset.to_batches( diff --git a/timdex_dataset_api/exceptions.py b/timdex_dataset_api/exceptions.py index 2ccbd77..deadb4a 100644 --- a/timdex_dataset_api/exceptions.py +++ b/timdex_dataset_api/exceptions.py @@ -1,9 +1,5 @@ """timdex_dataset_api/exceptions.py""" -class DatasetNotLoadedError(Exception): - """Custom exception for accessing methods requiring a loaded dataset.""" - - class InvalidDatasetRecordError(Exception): """Custom exception for invalid DatasetRecord instances."""