Merge pull request #168 from MITLibraries/TIMX-543-cr-optimize-v2

TIMX 543 - Optimize current records metadata queries and data retreival

Author: ghukill

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
 

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):

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(
@@ -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(
@@ -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

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",

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":
@@ -354,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]:
@@ -371,13 +376,16 @@ 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
         """
+        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_query = self.metadata.build_meta_query(table, limit, where, **filters)
         meta_df = self.metadata.conn.query(meta_query).to_df()
         logger.debug(
             f"Metadata query identified {len(meta_df)} rows, "
@@ -410,6 +418,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):
@@ -461,25 +473,35 @@ 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()
 
     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:
@@ -490,22 +512,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)