From 25d443019b03af503d99857264f60ca0fad8aacf Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Fri, 22 Aug 2025 14:28:40 -0400 Subject: [PATCH 1/7] current_records view relies on temp table Why these changes are being introduced: The current_records view has always been a performance and resource bottleneck. Moving to metadata and SQL has helped, but there was a little kink left in relation to using that metadata for data retreival. We often would materialize a query to a pandas dataframe for to use to drive data retrieval. In that moment, we do not benefit from having current_records be a view, when we're going to materialize the data anyhow. How this addresses that need: Utilizing a DuckDB temp table, we take a small performance hit on TIMDEXDatasetMetadata load, but then have near instant current_records queries thereafter. Additionally, we remove ordering in the metadata query for data retrieval and perform this in-memory with the pandas dataframe. Often this may be quite small, but even if large, it's more efficient here and already in python memory. This will also set the stage for performing just-in-time metadata queries as chunks before data retrieval, versus pulling all metadata rows in one query and then chunking that in memory. Side effects of this change: * Quicker metadata queries, small performance hit on load. Appears similarly memory intensive. Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-543 --- tests/test_read.py | 4 +- timdex_dataset_api/dataset.py | 7 +++ timdex_dataset_api/metadata.py | 83 +++++++++++++++++++++++----------- 3 files changed, 66 insertions(+), 28 deletions(-) diff --git a/tests/test_read.py b/tests/test_read.py index 207086c..f13112b 100644 --- a/tests/test_read.py +++ b/tests/test_read.py @@ -125,8 +125,8 @@ def test_read_batches_where_and_dataset_filters_are_combined(timdex_dataset_mult [ "SELECT * FROM current_records WHERE source = 'libguides'", "FROM records WHERE source = 'libguides'", - "source = 'libguides';", - " run_date = '2024-12-01'; ", + "ORDER BY timdex_record_id", + "LIMIT 3", ], ) def test_read_batches_where_rejects_non_predicate_sql( diff --git a/timdex_dataset_api/dataset.py b/timdex_dataset_api/dataset.py index 8e19685..ab3631d 100644 --- a/timdex_dataset_api/dataset.py +++ b/timdex_dataset_api/dataset.py @@ -375,10 +375,13 @@ def read_batches_iter( key/value DatasetFilters - filters: simple filtering based on key/value pairs from DatasetFilters """ + start_time = time.perf_counter() + # build and execute metadata query metadata_time = time.perf_counter() meta_query = self.metadata.build_meta_query(table, where, **filters) meta_df = self.metadata.conn.query(meta_query).to_df() + meta_df = meta_df.sort_values(by=["filename", "run_record_offset"]) logger.debug( f"Metadata query identified {len(meta_df)} rows, " f"across {len(meta_df.filename.unique())} parquet files, " @@ -410,6 +413,10 @@ def read_batches_iter( f"@ {batch_rps} records/second, total yielded: {total_yield_count}" ) + logger.debug( + f"read_batches_iter() elapsed: {round(time.perf_counter()-start_time, 2)}s" + ) + def _iter_meta_chunks(self, meta_df: pd.DataFrame) -> Iterator[pd.DataFrame]: """Utility method to yield chunks of metadata query results.""" for start in range(0, len(meta_df), self.config.duckdb_join_batch_size): diff --git a/timdex_dataset_api/metadata.py b/timdex_dataset_api/metadata.py index 227f3e0..cc45f92 100644 --- a/timdex_dataset_api/metadata.py +++ b/timdex_dataset_api/metadata.py @@ -334,6 +334,7 @@ def setup_duckdb_context(self) -> DuckDBPyConnection: start_time = time.perf_counter() conn = duckdb.connect() + conn.execute("""SET enable_progress_bar = false;""") self.configure_duckdb_connection(conn) if not self.database_exists(): @@ -436,36 +437,67 @@ def _create_current_records_view(self, conn: DuckDBPyConnection) -> None: This view builds on the table `records`. - This view includes only the most current version of each record in the dataset. - Because it includes the `timdex_record_id` and `run_id`, it makes yielding the - current version of a record via a TIMDEXDataset instance trivial: for any given - `timdex_record_id` if the `run_id` doesn't match, it's not the current version. + This metadata view includes only the most current version of each record in the + dataset. With the metadata provided from this view, we can streamline data + retrievals in TIMDEXDataset read methods. """ logger.info("creating view of current records metadata") - query = f""" - create or replace view metadata.current_records as - with ranked_records as ( + conn.execute( + """ + set temp_directory = '/tmp'; + """ + ) + + conn.execute( + """ + -- create temp table with current records using CTEs + create or replace temp table temp.main.current_records as + with + -- CTE of run_timestamp for last source full run + cr_source_last_full as ( + select + source, + max(run_timestamp) as last_full_ts + from metadata.records + where run_type = 'full' + group by source + ), + + -- CTE of all records, per source, on or after last full run + cr_since_last_full as ( + select + r.* + from metadata.records r + join cr_source_last_full f using (source) + where r.run_timestamp >= f.last_full_ts + ), + + -- CTE of records ranked by run_timestamp, with tie breaker + cr_ranked_records as ( + select + r.*, + row_number() over ( + partition by r.source, r.timdex_record_id + order by + r.run_timestamp desc nulls last, + r.run_id desc nulls last, + r.run_record_offset desc nulls last + ) as rn + from cr_since_last_full r + ) + + -- final select for current records (rn = 1) select - r.*, - row_number() over ( - partition by r.timdex_record_id - order by r.run_timestamp desc - ) as rn - from metadata.records r - where r.run_timestamp >= ( - select max(r2.run_timestamp) - from metadata.records r2 - where r2.source = r.source - and r2.run_type = 'full' - ) + * exclude (rn) + from cr_ranked_records + where rn = 1; + + -- create view in metadata schema + create or replace view metadata.current_records as + select * from temp.main.current_records; + """ ) - select - {','.join(ORDERED_METADATA_COLUMN_NAMES)} - from ranked_records - where rn = 1; - """ - conn.execute(query) def merge_append_deltas(self) -> None: """Merge append deltas into the static metadata database file.""" @@ -602,7 +634,6 @@ def build_meta_query( ).select_from(sa_table) if combined is not None: stmt = stmt.where(combined) - stmt = stmt.order_by(sa_table.c.filename, sa_table.c.run_record_offset) # using DuckDB dialect, compile to SQL string compiled = stmt.compile( From 4b2f2d2bd59512c5513bf1bda2acb4423285d87e Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Fri, 22 Aug 2025 15:18:40 -0400 Subject: [PATCH 2/7] Rename metadata rebuild method and improve refresh methods Why these changes are being introduced: The former TIMDEXDatasetMetadata method name recreate_static_database_file() was too narrowly focused. This method is responsible for rebuilding the entire dataset metadata structure. How this addresses that need: * Renames method * Updates refresh() methods on both TIMDEXDataset and TIMDEXDatasetMetadata to be more fully inclusive Side effects of this change: * None Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-543 --- tests/conftest.py | 6 +++--- tests/test_metadata.py | 2 +- tests/test_read.py | 4 ++-- timdex_dataset_api/dataset.py | 4 ++++ timdex_dataset_api/metadata.py | 25 ++++++++++++++----------- 5 files changed, 24 insertions(+), 17 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 6f89fe1..304d84e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -113,7 +113,7 @@ def timdex_dataset_multi_source(tmp_path_factory) -> TIMDEXDataset: ) # ensure static metadata database exists for read methods - dataset.metadata.recreate_static_database_file() + dataset.metadata.rebuild_dataset_metadata() dataset.metadata.refresh() return dataset @@ -223,7 +223,7 @@ def timdex_dataset_same_day_runs(tmp_path) -> TIMDEXDataset: def timdex_metadata(timdex_dataset_with_runs) -> TIMDEXDatasetMetadata: """TIMDEXDatasetMetadata with static database file created.""" metadata = TIMDEXDatasetMetadata(timdex_dataset_with_runs.location) - metadata.recreate_static_database_file() + metadata.rebuild_dataset_metadata() metadata.refresh() return metadata @@ -233,7 +233,7 @@ def timdex_dataset_with_runs_with_metadata( timdex_dataset_with_runs, ) -> TIMDEXDataset: """TIMDEXDataset with runs and static metadata created for read tests.""" - timdex_dataset_with_runs.metadata.recreate_static_database_file() + timdex_dataset_with_runs.metadata.rebuild_dataset_metadata() timdex_dataset_with_runs.metadata.refresh() return timdex_dataset_with_runs diff --git a/tests/test_metadata.py b/tests/test_metadata.py index d63144c..8f98bb9 100644 --- a/tests/test_metadata.py +++ b/tests/test_metadata.py @@ -43,7 +43,7 @@ def test_tdm_s3_dataset_structure_properties(s3_bucket_mocked): def test_tdm_create_metadata_database_file_success(caplog, timdex_metadata_empty): caplog.set_level("DEBUG") - timdex_metadata_empty.recreate_static_database_file() + timdex_metadata_empty.rebuild_dataset_metadata() def test_tdm_init_metadata_file_found_success(timdex_metadata): diff --git a/tests/test_read.py b/tests/test_read.py index f13112b..2bd31fd 100644 --- a/tests/test_read.py +++ b/tests/test_read.py @@ -254,7 +254,7 @@ def test_dataset_load_current_records_gets_correct_same_day_full_run( timdex_dataset_same_day_runs, ): # ensure metadata exists for this dataset - timdex_dataset_same_day_runs.metadata.recreate_static_database_file() + timdex_dataset_same_day_runs.metadata.rebuild_dataset_metadata() timdex_dataset_same_day_runs.metadata.refresh() df = timdex_dataset_same_day_runs.read_dataframe( table="current_records", run_type="full" @@ -265,7 +265,7 @@ def test_dataset_load_current_records_gets_correct_same_day_full_run( def test_dataset_load_current_records_gets_correct_same_day_daily_runs_ordering( timdex_dataset_same_day_runs, ): - timdex_dataset_same_day_runs.metadata.recreate_static_database_file() + timdex_dataset_same_day_runs.metadata.rebuild_dataset_metadata() timdex_dataset_same_day_runs.metadata.refresh() first_record = next( timdex_dataset_same_day_runs.read_dicts_iter( diff --git a/timdex_dataset_api/dataset.py b/timdex_dataset_api/dataset.py index ab3631d..774e57a 100644 --- a/timdex_dataset_api/dataset.py +++ b/timdex_dataset_api/dataset.py @@ -143,6 +143,10 @@ def location_scheme(self) -> Literal["file", "s3"]: def data_records_root(self) -> str: return f"{self.location.removesuffix('/')}/data/records" # type: ignore[union-attr] + def refresh(self) -> None: + """Fully reload TIMDEXDataset instance.""" + self.__init__(self.location) # type: ignore[misc] + def create_data_structure(self) -> None: """Ensure ETL records data structure exists in TIMDEX dataset.""" if self.location_scheme == "file": diff --git a/timdex_dataset_api/metadata.py b/timdex_dataset_api/metadata.py index cc45f92..2258aab 100644 --- a/timdex_dataset_api/metadata.py +++ b/timdex_dataset_api/metadata.py @@ -249,16 +249,14 @@ def refresh(self) -> None: self.conn = self.setup_duckdb_context() self._sa_metadata = sa_reflect_duckdb_conn(self.conn, schema="metadata") - def recreate_static_database_file(self) -> None: - """Create/recreate the static metadata database file. - - The following work is performed: - 1. Create a local working directory - 2. Open a DuckDB connection with a database file in this local working dir - 3. Create tables and views by scanning ETL data in dataset/data/records - 4. Close DuckDB connection ensuring a fully formed, local database file - 5. Upload DuckDB database file to target destination, making that the new - static metadata database file + def rebuild_dataset_metadata(self) -> None: + """Fully rebuild dataset metadata. + + Work includes: + - remove any append deltas, understanding a full metadata rebuild + will pickup that data from the ETL records themselves + - build a local, temporary static metadata database file, then overwrite the + canonical version in the dataset (e.g. in S3) """ if self.location_scheme == "s3": s3_client = S3Client() @@ -272,7 +270,6 @@ def recreate_static_database_file(self) -> None: with duckdb.connect(local_db_path) as conn: self.configure_duckdb_connection(conn) - conn.execute("""SET threads = 64;""") self._create_full_dataset_table(conn) @@ -299,6 +296,9 @@ def _create_full_dataset_table(self, conn: DuckDBPyConnection) -> None: start_time = time.perf_counter() logger.info("creating table of full dataset metadata") + # temporarily increase thread count + conn.execute("""SET threads = 64;""") + query = f""" create or replace table records as ( select @@ -312,6 +312,9 @@ def _create_full_dataset_table(self, conn: DuckDBPyConnection) -> None: """ conn.execute(query) + # reset thread count + conn.execute(f"""SET threads = {self.config.duckdb_connection_threads};""") + row_count = conn.query("""select count(*) from records;""").fetchone()[0] # type: ignore[index] logger.info( f"'records' table created - rows: {row_count}, " From 29d1e0e73f4fd8df19c0dc7de8df142476cf91f4 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Fri, 22 Aug 2025 16:00:00 -0400 Subject: [PATCH 3/7] Add LIMIT clause to read methods Why these changes are being introduced: Sometimes it can be helpful to limit the results from a read method. How this addresses that need: Adds optional limit= arg to all read methods which is passed along to the metadata query. By limiting the metadata results, we limit the data records retrieved. Side effects of this change: * None Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-543 --- tests/test_read.py | 6 ++++++ timdex_dataset_api/dataset.py | 32 +++++++++++++++++++++++++++----- timdex_dataset_api/metadata.py | 10 +++++++++- 3 files changed, 42 insertions(+), 6 deletions(-) diff --git a/tests/test_read.py b/tests/test_read.py index 2bd31fd..85fb085 100644 --- a/tests/test_read.py +++ b/tests/test_read.py @@ -276,3 +276,9 @@ def test_dataset_load_current_records_gets_correct_same_day_daily_runs_ordering( # just assert it's one of the daily runs assert first_record["run_id"] in {"run-4", "run-5"} assert first_record["action"] in {"index", "delete"} + + +def test_read_batches_iter_limit_returns_n_rows(timdex_dataset_multi_source): + batches = timdex_dataset_multi_source.read_batches_iter(limit=10) + table = pa.Table.from_batches(batches) + assert len(table) == 10 diff --git a/timdex_dataset_api/dataset.py b/timdex_dataset_api/dataset.py index 774e57a..3dd7443 100644 --- a/timdex_dataset_api/dataset.py +++ b/timdex_dataset_api/dataset.py @@ -358,6 +358,7 @@ def read_batches_iter( self, table: str = "records", columns: list[str] | None = None, + limit: int | None = None, where: str | None = None, **filters: Unpack[DatasetFilters], ) -> Iterator[pa.RecordBatch]: @@ -375,6 +376,7 @@ def read_batches_iter( Args: - table: an available DuckDB view or table - columns: list of columns to return + - limit: limit number of records yielded - where: raw SQL WHERE clause that can be used alone, or in combination with key/value DatasetFilters - filters: simple filtering based on key/value pairs from DatasetFilters @@ -383,7 +385,7 @@ def read_batches_iter( # build and execute metadata query metadata_time = time.perf_counter() - meta_query = self.metadata.build_meta_query(table, where, **filters) + meta_query = self.metadata.build_meta_query(table, limit, where, **filters) meta_df = self.metadata.conn.query(meta_query).to_df() meta_df = meta_df.sort_values(by=["filename", "run_record_offset"]) logger.debug( @@ -472,11 +474,16 @@ def read_dataframes_iter( self, table: str = "records", columns: list[str] | None = None, + limit: int | None = None, where: str | None = None, **filters: Unpack[DatasetFilters], ) -> Iterator[pd.DataFrame]: for record_batch in self.read_batches_iter( - table=table, columns=columns, where=where, **filters + table=table, + columns=columns, + limit=limit, + where=where, + **filters, ): yield record_batch.to_pandas() @@ -484,13 +491,18 @@ def read_dataframe( self, table: str = "records", columns: list[str] | None = None, + limit: int | None = None, where: str | None = None, **filters: Unpack[DatasetFilters], ) -> pd.DataFrame | None: df_batches = [ record_batch.to_pandas() for record_batch in self.read_batches_iter( - table=table, columns=columns, where=where, **filters + table=table, + columns=columns, + limit=limit, + where=where, + **filters, ) ] if not df_batches: @@ -501,22 +513,32 @@ def read_dicts_iter( self, table: str = "records", columns: list[str] | None = None, + limit: int | None = None, where: str | None = None, **filters: Unpack[DatasetFilters], ) -> Iterator[dict]: for record_batch in self.read_batches_iter( - table=table, columns=columns, where=where, **filters + table=table, + columns=columns, + limit=limit, + where=where, + **filters, ): yield from record_batch.to_pylist() def read_transformed_records_iter( self, table: str = "records", + limit: int | None = None, where: str | None = None, **filters: Unpack[DatasetFilters], ) -> Iterator[dict]: for record_dict in self.read_dicts_iter( - table=table, columns=["transformed_record"], where=where, **filters + table=table, + columns=["transformed_record"], + limit=limit, + where=where, + **filters, ): if transformed_record := record_dict["transformed_record"]: yield json.loads(transformed_record) diff --git a/timdex_dataset_api/metadata.py b/timdex_dataset_api/metadata.py index 2258aab..62beedc 100644 --- a/timdex_dataset_api/metadata.py +++ b/timdex_dataset_api/metadata.py @@ -612,7 +612,11 @@ def write_append_delta_duckdb(self, filepath: str) -> None: ) def build_meta_query( - self, table: str, where: str | None, **filters: Unpack["DatasetFilters"] + self, + table: str, + limit: int | None, + where: str | None, + **filters: Unpack["DatasetFilters"], ) -> str: """Build SQL query using SQLAlchemy against metadata schema tables and views.""" sa_table = self.get_sa_table(table) @@ -638,6 +642,10 @@ def build_meta_query( if combined is not None: stmt = stmt.where(combined) + # apply limit if present + if limit: + stmt = stmt.limit(limit) + # using DuckDB dialect, compile to SQL string compiled = stmt.compile( dialect=DuckDBDialect(), From 14c2b1b3f5204164f2982b0d8c6d312b8dce1a4a Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Mon, 25 Aug 2025 09:08:10 -0400 Subject: [PATCH 4/7] bump TDA version to 3.1 --- timdex_dataset_api/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/timdex_dataset_api/__init__.py b/timdex_dataset_api/__init__.py index fb1b437..bdd8bb8 100644 --- a/timdex_dataset_api/__init__.py +++ b/timdex_dataset_api/__init__.py @@ -4,7 +4,7 @@ from timdex_dataset_api.metadata import TIMDEXDatasetMetadata from timdex_dataset_api.record import DatasetRecord -__version__ = "3.0.0" +__version__ = "3.1.0" __all__ = [ "DatasetRecord", From 51a35ac6e814734eb511c744d6eb81a80584cd7b Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Mon, 25 Aug 2025 16:25:20 -0400 Subject: [PATCH 5/7] Move read ordering back to SQL query Why these changes are being introduced: The ordering of metadata records by filename + run_record_offset was moved into the python pandas context for a performance boost, but it was not ideal from the POV of keeping the majority of our logic in SQL. Upon learning that we could use `hash(filename)` to still order the filenames but with a dramatic speed and memory improvement, it makes sense to move this back into the SQL context. How this addresses that need: * Moves metadata query ordering back to SQL instead of python pandas context Side effects of this change: * None Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-543 --- timdex_dataset_api/dataset.py | 1 - timdex_dataset_api/metadata.py | 10 +++++++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/timdex_dataset_api/dataset.py b/timdex_dataset_api/dataset.py index 3dd7443..0198016 100644 --- a/timdex_dataset_api/dataset.py +++ b/timdex_dataset_api/dataset.py @@ -387,7 +387,6 @@ def read_batches_iter( metadata_time = time.perf_counter() meta_query = self.metadata.build_meta_query(table, limit, where, **filters) meta_df = self.metadata.conn.query(meta_query).to_df() - meta_df = meta_df.sort_values(by=["filename", "run_record_offset"]) logger.debug( f"Metadata query identified {len(meta_df)} rows, " f"across {len(meta_df.filename.unique())} parquet files, " diff --git a/timdex_dataset_api/metadata.py b/timdex_dataset_api/metadata.py index 62beedc..c8e6603 100644 --- a/timdex_dataset_api/metadata.py +++ b/timdex_dataset_api/metadata.py @@ -12,7 +12,7 @@ import duckdb from duckdb import DuckDBPyConnection from duckdb_engine import Dialect as DuckDBDialect -from sqlalchemy import Table, and_, select, text +from sqlalchemy import Table, and_, func, select, text from timdex_dataset_api.config import configure_logger from timdex_dataset_api.utils import ( @@ -642,6 +642,14 @@ def build_meta_query( if combined is not None: stmt = stmt.where(combined) + # order by filename + run_record_offset + # NOTE: we use a hash of the filename for ordering for a dramatic speedup, where + # we don't really care about the exact order, just that they are ordered + stmt = stmt.order_by( + func.hash(sa_table.c.filename), + sa_table.c.run_record_offset, + ) + # apply limit if present if limit: stmt = stmt.limit(limit) From 140a8d5ea7ea3d7b076088db21273d4eb739e900 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Tue, 26 Aug 2025 10:29:10 -0400 Subject: [PATCH 6/7] Remove imprecise 'tie breaker' language --- timdex_dataset_api/metadata.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/timdex_dataset_api/metadata.py b/timdex_dataset_api/metadata.py index c8e6603..92ae8e8 100644 --- a/timdex_dataset_api/metadata.py +++ b/timdex_dataset_api/metadata.py @@ -476,7 +476,7 @@ def _create_current_records_view(self, conn: DuckDBPyConnection) -> None: where r.run_timestamp >= f.last_full_ts ), - -- CTE of records ranked by run_timestamp, with tie breaker + -- CTE of records ranked by run_timestamp cr_ranked_records as ( select r.*, From c6ccd58ed72a22ca11cd64dcb6af6920a3f53885 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Tue, 26 Aug 2025 10:50:57 -0400 Subject: [PATCH 7/7] Additional current_records view docstrings --- timdex_dataset_api/metadata.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/timdex_dataset_api/metadata.py b/timdex_dataset_api/metadata.py index 92ae8e8..bea0d84 100644 --- a/timdex_dataset_api/metadata.py +++ b/timdex_dataset_api/metadata.py @@ -443,6 +443,14 @@ def _create_current_records_view(self, conn: DuckDBPyConnection) -> None: This metadata view includes only the most current version of each record in the dataset. With the metadata provided from this view, we can streamline data retrievals in TIMDEXDataset read methods. + + For performance reasons, the final view reads from a DuckDB temporary table that + is constructed, "temp.main.current_records". Because our connection is in memory, + the data in this temporary table is mostly in memory but has the ability to spill + to disk if we risk getting too close to our memory constraints. We explicitly + set the temporary location on disk for DuckDB at "/tmp" to play nice with contexts + like AWS ECS or Lambda, where sometimes the $HOME env var is missing; DuckDB + often tries to utilize the user's home directory and this works around that. """ logger.info("creating view of current records metadata")