TIMX 494 - TIMDEXRunManager for producing ETL run metadata

[ghukill]

Purpose and background context

This PR adds a new class TIMDEXRunManager that is a utility class for efficiently gathering metadata about all runs in the TIMDEX parquet dataset.

The original requirements of this library were to support efficient writing and reading from the dataset over the course of a single ETL run. For example Transmogrifier writes to the dataset where all records have a shared run_id and TIM reads from the dataset for that run_id and indexes those records to Opensearch. This was made very efficient by the year/month/day partitioning strategy, where TIM instantly narrows it down to parquet files under those partitions and then just reads the metadata from those few parquet files to get the file(s) that match a run_id.

However, it was assumed the dataset and this library would be capable of far more over time! One known use case: providing easy access to all the records for a given TIMDEX source that are currently in Opensearch. Let's think of this as the "current records for a source." To provide this functionality, we first need some understanding of the "runs" in the dataset specific to that source, including the run_date, run_type, run_id, associated parquet file, etc.

While technically we could load the whole dataset and use pyarrow dataset filtering or DuckDB, this is quite inefficient. This new class relies on some things we -- the dataset creators -- know about the dataset to support highly-parallelized run metadata retrieval:

• all rows in a given parquet file are from a single run, therefore metadata from a single row tell us about the run as a whole
• we can collect this information in parallel and each request is very lightweight
• we can group by run_id to create a full picture of a run from this collected metadata

And this is precisely how this utility class approaches producing dataset runs metadata:

1. Accept an instantiated and loaded TIMDEXDataset which has a list of all parquet files in the dataset
2. In parallel, read the metadata from the first row in each parquet file
3. Aggregate this and group by run_id
4. Returning a dataframe of run metadata across all runs in the dataset

This has proven to be very fast, around 1-2 seconds given even hundreds of parquet files. With that final run metadata in hand, it paves the way for future functionality like yielding only the currrent version of a record, for a source.

How can a reviewer manually see the effects of these changes?

1- For all tests below, set AWS Dev Admin credentials (using a test bucket with a static clone of prod data)

2- Start an ipython shell:

pipenv run ipython

Sometimes an alternate approach can illuminate the proposed one. Let's use DuckDB to get ETL run metadata.

Define a function that uses a SQL query to get at run data, grouping by run_id:

import time

import duckdb

def get_source_etl_runs(source:str):
    t0 = time.perf_counter()
    with duckdb.connect() as conn:

        # setup AWS credentials
        conn.execute("""
        CREATE OR REPLACE SECRET aws_secret (
            TYPE s3,
            PROVIDER credential_chain,
            CHAIN 'env;sso;process'
        );
        """)

        # perform query
        query = f"""
        with dataset as (
            select * from read_parquet(
                's3://ghukill-test/timdex-dataset/prod_2025_05_05/**/*.parquet',
                filename=true
            )
        ),
        etl_runs as (
            select
                source,
                run_date,
                run_type,
                run_id,
                count(*) as num_rows,
                array_agg(filename) as parquet_files
            from dataset
            where source='{source}'
            group by run_date, source, run_type, run_id  -- here is the expensive part of the query
        )
        select
            *
        from etl_runs
        ;
        """
        result = conn.query(query).to_df()
        print(f"elapsed: {time.perf_counter() - t0}")
        return result

Now, let's run this for a few different sources:

df1 = get_source_etl_runs('libguides')
100% ▕████████████████████████████████████████████████████████████▏ 
elapsed: 10.36918495898135

df2 = get_source_etl_runs('dspace')
100% ▕████████████████████████████████████████████████████████████▏ 
elapsed: 12.032363249920309

df3 = get_source_etl_runs('alma')
100% ▕████████████████████████████████████████████████████████████▏ 
elapsed: 120.38288041693158

The first run for libguides and even dspace looked promising: 10-12 seconds and low memory usage. But what happened with alma?

Unfortunately, DuckDB doesn't know some things about our dataset that it could leverage. It's amazing that it's this fast, but for alma we can see it's very briefly touching each row in the dataset to get that run ETL metadata, which is not needed. Given that alma is by far the largest, getting ETL run data for all sources would be about the same time (just a little slower).

We would see similar performance results using pyarrow.dataset.filter() and pyarrow.dataset.group_by().

Though we want run-level metadata, we also know that for each parquet file its associated ETL run metadata is present in every row, e.g. the first, single row. It's no shade on DuckDB or pyarrow, they simply don't know as much about dataset and thus touch much more data than needed.

This PR body probably isn't the best place for in-depth analysis of those approaches, but suffice to say, tried many angles and just couldn't get close to matching the performance of the custom class demonstrated next.

Now, let's use the TIMDEXRunManager to get the same data.

Configure the dev logger for more output:

from timdex_dataset_api.config import configure_dev_logger

configure_dev_logger()

Next, load a normal TIMDEXDataset instance with represents the full dataset:

from timdex_dataset_api import TIMDEXDataset

td = TIMDEXDataset("s3://ghukill-test/timdex-dataset/prod_2025_05_05")
td.load()

Note the quick load time, unchanged from before, discovering all parquet files in the dataset.

Observe how many parquet files comprise the dataset:

print(len(td.dataset.files))
# 370

Now, load a TIMDEXRunManager instance and retrieve ETL run data for all runs:

from timdex_dataset_api.run import TIMDEXRunManager

ts = TIMDEXRunManager(timdex_dataset=td)
runs_df = ts.get_runs_metadata()

Metadata about all ~270 ETL runs comes back in about 1-3 seconds! A fraction of a single, small source using DuckDB, and orders and orders magnitude faster than touching sources like alma. Furthermore, this will continue to scale as the number of parquet files grows.

Lastly, let's get an ordered list of all parquet files for a given source that represent the "current" records in TIMDEX:

parquet_files = ts.get_current_source_parquet_files('alma')

Note it's basically instant, given that we've cached the full dataframe of ETL run metadata. If we had not done that step yet, this would take that 1-3 seconds we saw above.

A valid question would now be, "what do we do with a list of reverse chronologically ordered parquet files?" The answer is we use that list to load a new TIMDEXDataset instance with those specific, ordered parquet files.

This use of the valuable list of ordered, specific parquet files leads into a future PR which will utilize this valuable information to provide an efficient way to yield only "current" records for a given TIMDEX source.

Includes new or updated dependencies?

NO

Changes expectations for external applications?

NO

What are the relevant tickets?

• https://mitlibraries.atlassian.net/browse/TIMX-493
• https://mitlibraries.atlassian.net/browse/TIMX-494

Developer

• All new ENV is documented in README
• All new ENV has been added to staging and production environments
• All related Jira tickets are linked in commit message(s)
• Stakeholder approval has been confirmed (or is not needed)

Code Reviewer(s)

• The commit message is clear and follows our guidelines (not just this PR message)
• There are appropriate tests covering any new functionality
• The provided documentation is sufficient for understanding any new functionality introduced
• Any manual tests have been performed or provided examples verified
• New dependencies are appropriate or there were no changes

[ghukill]

Maybe not obvious on a first pass, but it is here that we are reading the metadata from parquet files in parallel. max_workers is set to 250, which is quite high, but believed to be safe given the lightweight nature of the HTTP request and response (just a few bytes).

[ehanson8]

Thank you for clarifying this!

[ghukill]

Because a single ETL run may span multiple parquet files (we limit the parquet files to 100k rows), we must group by run_id to ensure that our final result is grouped at the ETL run level.

[ehanson8]

Great context for the new functionality, a few comments!

[ehanson8]

Thank you for clarifying this!

[ehanson8]

Maybe this will change later as more methods are added but it feels like the method order should flip given that each method is contained in the one below it? There are certainly exceptions to every rule, but I thought we were generally leading with the highest-level of abstraction and then descending into more specific methods. This also may be a good DataEng meeting topic to discuss in more detail to clarify our norms!

[ghukill]

I think I agree! Thanks for pushing on this. As I develop a class like this, I'm often writing the lower level / private methods first to build up to the higher-level / public / interface methods.

I'll reorder, and actually privatize a couple.

I'm happy to discuss, my proposal would be:

1. properties
2. "magic" methods that override the __X__ dunder methods
3. public methods, the primary interface
4. private methods, starting with leading underscores _X

Internal ordering of 3, the public methods, I'd vote for dealer's choice, whatever makes sense.

Bringing it back here, if I were to make the truly private methods private, then the ordering kind of takes care of itself via that approach!

Thanks again though, good suggestion.

[ehanson8]

This is so much more logical and readable to me, great change! I agree with your proposed order, @jonavellecuerdo should also weigh in. We should use private methods more frequently to achieve this type of order (I still forget to at times as well) but the behavior of this class is so much clearer with the public methods up top

[ehanson8]

Also, it seems the branch protection rules should be updated to block merging without an approval

[ghukill]

Also, it seems the branch protection rules should be updated to block merging without an approval

Good call. Based on new github permissions, I think we'll need to ask an admin.

[coveralls]

Pull Request Test Coverage Report for Build 15163424367

Details

• 63 of 69 (91.3%) changed or added relevant lines in 1 file are covered.
• No unchanged relevant lines lost coverage.
• Overall coverage decreased (-1.1%) to 94.565%

---

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
timdex_dataset_api/run.py	63	69	91.3%

Totals
Change from base Build 15144487872:	-1.1%
Covered Lines:	261
Relevant Lines:	276

---

💛 - Coveralls

[ehanson8]

Excellent changes! Consider this approved from me but letting @jonavellecuerdo handle the final approval since she self-assigned

[ehanson8]

This is so much more logical and readable to me, great change! I agree with your proposed order, @jonavellecuerdo should also weigh in. We should use private methods more frequently to achieve this type of order (I still forget to at times as well) but the behavior of this class is so much clearer with the public methods up top

[ehanson8]

Also, it seems the branch protection rules should be updated to block merging without an approval

Good call. Based on new github permissions, I think we'll need to ask an admin.

Ah, we should also ask if default rules can be created for new repos in this new permissions scheme so we don't have to bother an admin each time

[jonavellecuerdo]

Really cool stuff in this PR. Approved! 😃