From 97df92d1b32637d3d6899f2f0f1b65e40ba36da3 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Fri, 30 May 2025 11:27:36 -0400 Subject: [PATCH 1/3] Add migrations folder and run_timestamp migration Why these changes are being introduced: We have a need to perform a onetime bulk edit of the parquet dataset, which is very similar to a SQL database migration: addition of a column and backfilling of data. While we may not require the strictness of SQL schema migrations, which often can be "replayed" in order to upgrade or reverse order to downgrade, we nonetheless would benefit from some structure and code-as-documentation for bulk operations we perform on the actual parquet dataset. How this addresses that need: * Creates a new root level migrations/ folder * Includes a README with a proposed simple structure for migrations (mostly naming conventions) * Includes our first migration to backfill all pre-existing parquet files with a new 'run_timestamp' column * The run_timestamp will be pulled from the creation date of the parquet file in S3, which is close enough for our purposes Side effects of this change: * TDA becomes the location for storing code related to "migrations" for the parquet dataset * All pre-existing parquet files will now have a run_timestamp column Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-496 --- ...025_05_30_backfill_run_timestamp_column.py | 204 ++++++++++++++++++ migrations/README.md | 57 +++++ migrations/__init__.py | 0 3 files changed, 261 insertions(+) create mode 100644 migrations/001_2025_05_30_backfill_run_timestamp_column.py create mode 100644 migrations/README.md create mode 100644 migrations/__init__.py diff --git a/migrations/001_2025_05_30_backfill_run_timestamp_column.py b/migrations/001_2025_05_30_backfill_run_timestamp_column.py new file mode 100644 index 0000000..d508b74 --- /dev/null +++ b/migrations/001_2025_05_30_backfill_run_timestamp_column.py @@ -0,0 +1,204 @@ +# ruff: noqa: BLE001, D212, TRY300, TRY400 +""" +Date: 2025-05-30 + +Description: + +After the creation of a new run_timestamp column as part of Jira ticket TIMX-496, there +was a need to backfill a run timestamp for all parquet files in the dataset. + +This migration performs the following: +1. retrieves all parquet file from the dataset +2. for each parquet file: + a. if the run_timestamp column already exists, skip + b. retrieve the file creation date of the parquet file, this becomes the run_timestamp + c. rewrite the parquet file with a new run_timestamp column + +Usage: + +pipenv run python migrations/001_2025_05_30_backfill_run_timestamp_column.py \ + \ +--dry-run +""" + +import argparse +import json +import time +from datetime import UTC, datetime + +import pyarrow as pa +import pyarrow.dataset as ds +import pyarrow.parquet as pq +from pyarrow import fs + +from timdex_dataset_api.config import configure_dev_logger, configure_logger +from timdex_dataset_api.dataset import TIMDEX_DATASET_SCHEMA, TIMDEXDataset + +configure_dev_logger() + +logger = configure_logger(__name__) + + +def backfill_dataset(location: str, *, dry_run: bool = False) -> None: + """Main entrypoint for backfill script. + + Loop through all parquet files in the dataset and, if the run_timestamp column does + not exist, create it using the S3 object creation date. + """ + start_time = time.perf_counter() + td = TIMDEXDataset(location) + td.load() + + parquet_files = td.dataset.files # type: ignore[attr-defined] + logger.info(f"Found {len(parquet_files)} parquet files in dataset.") + + success_count = 0 + skip_count = 0 + error_count = 0 + + for i, parquet_file in enumerate(parquet_files): + logger.info( + f"Working on parquet file {i + 1}/{len(parquet_files)}: {parquet_file}" + ) + + success, result = backfill_parquet_file(parquet_file, td.dataset, dry_run=dry_run) + + if success: + if result and "skipped" in result: + skip_count += 1 + else: + success_count += 1 + else: + error_count += 1 + + logger.info(json.dumps(result)) + + logger.info( + f"Backfill complete. Elapsed: {time.perf_counter()-start_time}, " + f"Success: {success_count}, Skipped: {skip_count}, Errors: {error_count}" + ) + + +def get_s3_object_creation_date(file_path: str, filesystem: fs.FileSystem) -> datetime: + """Get the creation date of an S3 object. + + Args: + file_path: Path to the S3 object + filesystem: PyArrow S3 filesystem instance + + Returns: + datetime: Creation date of the S3 object in UTC + """ + try: + # Get creation date of S3 object + file_info = filesystem.get_file_info(file_path) + creation_date: datetime = file_info.mtime # type: ignore[assignment] + + # Ensure it's timezone-aware and in UTC + if creation_date.tzinfo is None: + creation_date = creation_date.replace(tzinfo=UTC) + elif creation_date.tzinfo != UTC: + creation_date = creation_date.astimezone(UTC) + + return creation_date + + except Exception as e: + logger.error(f"Error getting S3 object creation date for {file_path}: {e}") + raise + + +def backfill_parquet_file( + parquet_filepath: str, + dataset: ds.Dataset, + *, + dry_run: bool = False, +) -> tuple[bool, dict]: + """Backfill a single parquet file with run_timestamp column. + + Args: + parquet_filepath: Path to the parquet file + dataset: PyArrow dataset instance + dry_run: If True, don't actually write changes + + Returns: + Tuple of (success: bool, result: dict) + """ + start_time = time.perf_counter() + try: + parquet_file = pq.ParquetFile(parquet_filepath, filesystem=dataset.filesystem) # type: ignore[attr-defined] + + # Check if run_timestamp column already exists + if "run_timestamp" in parquet_file.schema.names: + logger.info( + f"Parquet already has 'run_timestamp', skipping: {parquet_filepath}" + ) + return True, {"file_path": parquet_filepath, "skipped": True} + + # Read all rows from the parquet file into a pyarrow Table + # NOTE: memory intensive for very large parquet files, though suitable for onetime + # migration work. + table = parquet_file.read() + + # Get S3 object creation date + creation_date = get_s3_object_creation_date(parquet_filepath, dataset.filesystem) # type: ignore[attr-defined] + + # Create run_timestamp column using the exact schema definition + num_rows = len(table) + run_timestamp_field = TIMDEX_DATASET_SCHEMA.field("run_timestamp") + run_timestamp_array = pa.array( + [creation_date] * num_rows, type=run_timestamp_field.type + ) + + # Add the run_timestamp column to the table + table_with_timestamp = table.append_column("run_timestamp", run_timestamp_array) + + # Write the updated table back to the same file + if not dry_run: + pq.write_table( + table_with_timestamp, # type: ignore[attr-defined] + parquet_filepath, + filesystem=dataset.filesystem, # type: ignore[attr-defined] + ) + logger.info(f"Successfully updated file: {parquet_filepath}") + else: + logger.info(f"DRY RUN: Would update file: {parquet_filepath}") + + update_details = { + "file_path": parquet_filepath, + "rows_updated": num_rows, + "run_timestamp_added": creation_date.isoformat(), + "elapsed": time.perf_counter() - start_time, + "dry_run": dry_run, + } + + return True, update_details + + except Exception as e: + logger.error(f"Error processing parquet file {parquet_filepath}: {e}") + return False, { + "file_path": parquet_filepath, + "error": str(e), + "elapsed": time.perf_counter() - start_time, + "dry_run": dry_run, + } + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description=( + "Backfill run_timestamp column in TIMDEX parquet files " + "using S3 creation dates" + ) + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Scan files and report what would be done without making changes", + ) + parser.add_argument( + "dataset_location", help="Path to the dataset (local path or s3://bucket/path)" + ) + + args = parser.parse_args() + + backfill_dataset(args.dataset_location, dry_run=args.dry_run) diff --git a/migrations/README.md b/migrations/README.md new file mode 100644 index 0000000..8bd66da --- /dev/null +++ b/migrations/README.md @@ -0,0 +1,57 @@ +# TIMDEX Dataset Migrations + +This directory includes manual, bulk migrations of data and schema in the TIMDEX parquet dataset. Consider it like migrations for a SQL database, except a bit more unstructured and ad-hoc. + +## Structure + +Each migration is either a single python file, or a dedicated directory, with that follows the naming convention: + + - `###_`: incrementing migration sequence number + - `YYYY_MM_DD_`: approximate date of migration creation and run + - `short_name.py` (file) or `short_name` (directory): short migration name + +Examples: + + - `001_2025_05_30_backfill_run_timestamp_column.py` --> single file + - `002_2025_06_15_remove_errant_parquet_files` --> directory that contains 1+ files + +The entrypoint for each migration should contain a docstring at the root of the file with a structure like: + +```python +""" +Date: YYYY-MM-DD + +Description: + +Description here about the nature of the migration... + +Usage: + +Explanation here for how to run it... +""" +``` + +Example: +```python +""" +Date: 2025-05-30 + +Description: + +After the creation of a new run_timestamp column as part of Jira ticket TIMX-496, there +was a need to backfill a run timestamp for all parquet files in the dataset. + +This migration performs the following: +1. retrieves all parquet file from the dataset +2. for each parquet file: + a. if the run_timestamp column already exists, skip + b. retrieve the file creation date of the parquet file, this becomes the run_timestamp + c. rewrite the parquet file with a new run_timestamp column + +Usage: +PYTHONPATH=. \ +pipenv run python migrations/001_2025_05_30_backfill_run_timestamp_column.py \ + \ +--dry-run +""" +``` \ No newline at end of file diff --git a/migrations/__init__.py b/migrations/__init__.py new file mode 100644 index 0000000..e69de29 From 2fa84df2c2cbd013247e9bc0ac0c9c6505f7a714 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Tue, 3 Jun 2025 11:58:36 -0400 Subject: [PATCH 2/3] Reorder functions in run_timestamp backfill migration --- ...025_05_30_backfill_run_timestamp_column.py | 56 +++++++++---------- 1 file changed, 28 insertions(+), 28 deletions(-) diff --git a/migrations/001_2025_05_30_backfill_run_timestamp_column.py b/migrations/001_2025_05_30_backfill_run_timestamp_column.py index d508b74..26fb732 100644 --- a/migrations/001_2025_05_30_backfill_run_timestamp_column.py +++ b/migrations/001_2025_05_30_backfill_run_timestamp_column.py @@ -79,34 +79,6 @@ def backfill_dataset(location: str, *, dry_run: bool = False) -> None: ) -def get_s3_object_creation_date(file_path: str, filesystem: fs.FileSystem) -> datetime: - """Get the creation date of an S3 object. - - Args: - file_path: Path to the S3 object - filesystem: PyArrow S3 filesystem instance - - Returns: - datetime: Creation date of the S3 object in UTC - """ - try: - # Get creation date of S3 object - file_info = filesystem.get_file_info(file_path) - creation_date: datetime = file_info.mtime # type: ignore[assignment] - - # Ensure it's timezone-aware and in UTC - if creation_date.tzinfo is None: - creation_date = creation_date.replace(tzinfo=UTC) - elif creation_date.tzinfo != UTC: - creation_date = creation_date.astimezone(UTC) - - return creation_date - - except Exception as e: - logger.error(f"Error getting S3 object creation date for {file_path}: {e}") - raise - - def backfill_parquet_file( parquet_filepath: str, dataset: ds.Dataset, @@ -183,6 +155,34 @@ def backfill_parquet_file( } +def get_s3_object_creation_date(file_path: str, filesystem: fs.FileSystem) -> datetime: + """Get the creation date of an S3 object. + + Args: + file_path: Path to the S3 object + filesystem: PyArrow S3 filesystem instance + + Returns: + datetime: Creation date of the S3 object in UTC + """ + try: + # Get creation date of S3 object + file_info = filesystem.get_file_info(file_path) + creation_date: datetime = file_info.mtime # type: ignore[assignment] + + # Ensure it's timezone-aware and in UTC + if creation_date.tzinfo is None: + creation_date = creation_date.replace(tzinfo=UTC) + elif creation_date.tzinfo != UTC: + creation_date = creation_date.astimezone(UTC) + + return creation_date + + except Exception as e: + logger.error(f"Error getting S3 object creation date for {file_path}: {e}") + raise + + if __name__ == "__main__": parser = argparse.ArgumentParser( description=( From 243c7e1c35942a96e956770f7c9dc41f27ebad4e Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Tue, 3 Jun 2025 13:47:14 -0400 Subject: [PATCH 3/3] Update migrations README and migration 001 --- .../001_2025_05_30_backfill_run_timestamp_column.py | 13 +++++++++++++ migrations/README.md | 12 ++++++++++-- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/migrations/001_2025_05_30_backfill_run_timestamp_column.py b/migrations/001_2025_05_30_backfill_run_timestamp_column.py index 26fb732..b80ded3 100644 --- a/migrations/001_2025_05_30_backfill_run_timestamp_column.py +++ b/migrations/001_2025_05_30_backfill_run_timestamp_column.py @@ -14,6 +14,15 @@ b. retrieve the file creation date of the parquet file, this becomes the run_timestamp c. rewrite the parquet file with a new run_timestamp column +Side effects: + +1- Loss of "Last Modified" date in S3 + +This migration is using the original "Last Modified" date in S3 that was minted when the +parquet file was written. It is storing that data in a `run_timestamp` column and thus +will persist, but the actual parquet file will LOSE this "Last Modified" date when it is +recreated. + Usage: pipenv run python migrations/001_2025_05_30_backfill_run_timestamp_column.py \ @@ -158,6 +167,10 @@ def backfill_parquet_file( def get_s3_object_creation_date(file_path: str, filesystem: fs.FileSystem) -> datetime: """Get the creation date of an S3 object. + This function assumes that all datetimes coming back are coming from the same source + and will be formatted similarly, which means either all values are timezone aware or + not. + Args: file_path: Path to the S3 object filesystem: PyArrow S3 filesystem instance diff --git a/migrations/README.md b/migrations/README.md index 8bd66da..01d57e0 100644 --- a/migrations/README.md +++ b/migrations/README.md @@ -1,10 +1,16 @@ # TIMDEX Dataset Migrations -This directory includes manual, bulk migrations of data and schema in the TIMDEX parquet dataset. Consider it like migrations for a SQL database, except a bit more unstructured and ad-hoc. +This directory stores data and/or schema modifications that were made to the TIMDEX parquet dataset. Consider them like ["migrations"](https://en.wikipedia.org/wiki/Schema_migration) for a SQL database, but -- at least at the time of this writing -- considerably more informal and ad-hoc. + +Unless otherwise noted, it assumed that these migrations were: + + * manually run by a developer, either on a local machine or some cloud operations + * have been performed already, should not be performed again + * the migration script does not contain a way to rollback the changes ## Structure -Each migration is either a single python file, or a dedicated directory, with that follows the naming convention: +Each migration is either a single python file, or a dedicated directory, that follow this naming convention: - `###_`: incrementing migration sequence number - `YYYY_MM_DD_`: approximate date of migration creation and run @@ -15,6 +21,8 @@ Examples: - `001_2025_05_30_backfill_run_timestamp_column.py` --> single file - `002_2025_06_15_remove_errant_parquet_files` --> directory that contains 1+ files +Files inside a migration directory like `002_2025_06_15_remove_errant_parquet_files` are _not_ expected to follow any particular format (though a `README.md` is encourage to inform future developers how it was performed!). + The entrypoint for each migration should contain a docstring at the root of the file with a structure like: ```python