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 207086c..85fb085 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( @@ -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 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", diff --git a/timdex_dataset_api/dataset.py b/timdex_dataset_api/dataset.py index 8e19685..0198016 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": @@ -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,11 +473,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() @@ -473,13 +490,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: @@ -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) diff --git a/timdex_dataset_api/metadata.py b/timdex_dataset_api/metadata.py index 227f3e0..bea0d84 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 ( @@ -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. + def rebuild_dataset_metadata(self) -> None: + """Fully rebuild dataset metadata. - 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 + 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}, " @@ -334,6 +337,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 +440,75 @@ 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. + + 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") - 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 + 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.""" @@ -577,7 +620,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) @@ -602,7 +649,18 @@ 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) + + # 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) # using DuckDB dialect, compile to SQL string compiled = stmt.compile(