TIMX 543 - keyset pagination for read methods

tests/conftest.py

@@ -114,7 +114,7 @@ def timdex_dataset_multi_source(tmp_path_factory) -> TIMDEXDataset:
 
     # ensure static metadata database exists for read methods
     dataset.metadata.rebuild_dataset_metadata()
-    dataset.metadata.refresh()
+    dataset.refresh()
 
     return dataset
 
@@ -234,7 +234,7 @@ def timdex_dataset_with_runs_with_metadata(
 ) -> TIMDEXDataset:
     """TIMDEXDataset with runs and static metadata created for read tests."""
     timdex_dataset_with_runs.metadata.rebuild_dataset_metadata()
-    timdex_dataset_with_runs.metadata.refresh()
+    timdex_dataset_with_runs.refresh()
     return timdex_dataset_with_runs
 
 

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.1.0"
+__version__ = "3.2.0"
 
 __all__ = [
     "DatasetRecord",

timdex_dataset_api/dataset.py

@@ -364,7 +364,8 @@ def read_batches_iter(
     ) -> Iterator[pa.RecordBatch]:
         """Yield ETL records as pyarrow.RecordBatches.
 
-        This method performs a two step process:
+        This is the base read method.  All read methods eventually drop down and use this
+        for streaming batches of records.  This method performs a two-step process:
 
             1. Perform a "metadata" query that narrows down records and physical parquet
             files to read from.
@@ -383,34 +384,43 @@ def read_batches_iter(
         """
         start_time = time.perf_counter()
 
-        # build and execute metadata query
-        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()
-        logger.debug(
-            f"Metadata query identified {len(meta_df)} rows, "
-            f"across {len(meta_df.filename.unique())} parquet files, "
-            f"elapsed: {round(time.perf_counter()-metadata_time,2)}s"
-        )
-
-        # execute data queries in batches and yield results
+        temp_table_name = "read_meta_chunk"
         total_yield_count = 0
-        for i, meta_chunk_df in enumerate(self._iter_meta_chunks(meta_df)):
+
+        for i, meta_chunk_df in enumerate(
+            self._iter_meta_chunks(
+                table,
+                limit=limit,
+                where=where,
+                **filters,
+            )
+        ):
             batch_time = time.perf_counter()
             batch_yield_count = len(meta_chunk_df)
             total_yield_count += batch_yield_count
 
-            if batch_yield_count == 0:
-                continue
-
-            self.conn.register("meta_chunk", meta_chunk_df)
-            data_query = self._build_data_query_for_chunk(
-                columns,
-                meta_chunk_df,
-                registered_metadata_chunk="meta_chunk",
+            self.conn.register(
+                temp_table_name,
+                meta_chunk_df[
+                    [
+                        "timdex_record_id",
+                        "run_id",
+                        "run_record_offset",
+                    ]
+                ],
             )
-            yield from self._stream_data_query_batches(data_query)
-            self.conn.unregister("meta_chunk")
+
+            # build and perform data query, yield records
+            # set in try/finally block to ensure we always deregister the meta table
+            try:
+                data_query = self._build_data_query_for_chunk(
+                    columns,
+                    meta_chunk_df,
+                    registered_metadata_chunk=temp_table_name,
+                )
+                yield from self._iter_data_chunks(data_query)
+            finally:
+                self.conn.unregister(temp_table_name)
 
             batch_rps = int(batch_yield_count / (time.perf_counter() - batch_time))
             logger.debug(
@@ -422,32 +432,104 @@ def read_batches_iter(
             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):
-            yield meta_df.iloc[start : start + self.config.duckdb_join_batch_size]
+    def _iter_meta_chunks(
+        self,
+        table: str = "records",
+        limit: int | None = None,
+        where: str | None = None,
+        **filters: Unpack[DatasetFilters],
+    ) -> Iterator[pd.DataFrame]:
+        """Utility method to yield pandas Dataframe chunks of metadata query results.
 
-    def _build_parquet_file_list(self, meta_chunk_df: pd.DataFrame) -> str:
-        """Build SQL list of parquet filepaths."""
-        filenames = meta_chunk_df["filename"].unique().tolist()
-        if self.location_scheme == "s3":
-            filenames = [f"s3://{f.removeprefix('s3://')}" for f in filenames]
-        return "[" + ",".join((f"'{f}'") for f in filenames) + "]"
+        The approach here is to use "keyset" pagination, which means each paged result
+        is a greater-than (>) check against a tuple of ordered values from the previous
+        chunk.  This is more performant than a LIMIT + OFFSET.
+        """
+        # use duckdb_join_batch_size as the chunk size for keyset pagination
+        chunk_size = self.config.duckdb_join_batch_size
+
+        # init keyset value of zeros to begin with
+        keyset_value = (0, 0, 0)
+
+        total_yielded = 0
+        while True:
+
+            # enforce limit if passed
+            if limit is not None:
+                remaining = limit - total_yielded
+                if remaining <= 0:
+                    break
+                chunk_limit = min(chunk_size, remaining)
+            else:
+                chunk_limit = chunk_size
+
+            # perform chunk query and convert to pyarrow Table
+            meta_query = self.metadata.build_keyset_paginated_metadata_query(
+                table,
+                limit=chunk_limit,  # pass chunk_limit instead of limit
+                where=where,
+                keyset_value=keyset_value,
+                **filters,
+            )
+            meta_chunk_df = self.metadata.conn.query(meta_query).to_df()
+
+            meta_chunk_count = len(meta_chunk_df)
+
+            # an empty chunk signals end of pagination
+            if meta_chunk_count == 0:
+                break
+
+            # yield this chunk of data
+            total_yielded += meta_chunk_count
+            yield meta_chunk_df[
+                [
+                    "timdex_record_id",
+                    "run_id",
+                    "run_record_offset",
+                    "filename",
+                ]
+            ]
+
+            # update keyset value using the last row from this chunk
+            last_row = meta_chunk_df.iloc[-1]
+            keyset_value = (
+                int(last_row.filename_hash),
+                int(last_row.run_id_hash),
+                int(last_row.run_record_offset),
+            )
 
     def _build_data_query_for_chunk(
         self,
         columns: list[str] | None,
         meta_chunk_df: pd.DataFrame,
         registered_metadata_chunk: str = "meta_chunk",
     ) -> str:
-        """Build SQL query used for data retrieval, joining on metadata data."""
-        parquet_list_sql = self._build_parquet_file_list(meta_chunk_df)
-        rro_list_sql = ",".join(
-            str(rro) for rro in meta_chunk_df["run_record_offset"].unique()
-        )
+        """Build SQL query used for data retrieval, joining on passed metadata data."""
+        # build select columns
         select_cols = ",".join(
             [f"ds.{col}" for col in (columns or TIMDEX_DATASET_SCHEMA.names)]
         )
+
+        # build list of explicit parquet files to read from
+        filenames = list(meta_chunk_df["filename"].unique())
+        if self.location_scheme == "s3":
+            filenames = [
+                f"s3://{f.removeprefix('s3://')}" for f in filenames  # type: ignore[union-attr]
+            ]
+        parquet_list_sql = "[" + ",".join(f"'{f}'" for f in filenames) + "]"
+
+        # build run_record_offset WHERE clause to leverage row group pruning
+        rro_values = meta_chunk_df["run_record_offset"].unique()
+        rro_values.sort()
+        if len(rro_values) <= 1_000:  # noqa: PLR2004
+            rro_clause = (
+                f"and run_record_offset in ({','.join(str(rro) for rro in rro_values)})"
+            )
+        else:
+            rro_clause = (
+                f"and run_record_offset between {rro_values[0]} and {rro_values[-1]}"
+            )
+
         return f"""
             select
                 {select_cols}
@@ -459,15 +541,24 @@ def _build_data_query_for_chunk(
             inner join {registered_metadata_chunk} mc using (
                 timdex_record_id, run_id, run_record_offset
             )
-            where ds.run_record_offset in ({rro_list_sql});
+            where true
+            {rro_clause};
             """
 
-    def _stream_data_query_batches(self, data_query: str) -> Iterator[pa.RecordBatch]:
-        """Yield pyarrow RecordBatches from a SQL query."""
-        self.conn.execute("set enable_progress_bar = false;")
-        cursor = self.conn.execute(data_query)
-        yield from cursor.fetch_record_batch(rows_per_batch=self.config.read_batch_size)
-        self.conn.execute("set enable_progress_bar = true;")
+    def _iter_data_chunks(self, data_query: str) -> Iterator[pa.RecordBatch]:
+        """Perform a query to retrieve data and stream chunks."""
+        if self.location_scheme == "s3":
+            self.conn.execute("""set threads=16;""")
+        try:
+            cursor = self.conn.execute(data_query)
+            yield from cursor.fetch_record_batch(
+                rows_per_batch=self.config.read_batch_size
+            )
+        finally:
+            if self.location_scheme == "s3":
+                self.conn.execute(
+                    f"""set threads={self.metadata.config.duckdb_connection_threads};"""
+                )
 
     def read_dataframes_iter(
         self,

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_, func, select, text
+from sqlalchemy import Table, func, literal, select, text, tuple_
 
 from timdex_dataset_api.config import configure_logger
 from timdex_dataset_api.utils import (
@@ -619,42 +619,56 @@ def write_append_delta_duckdb(self, filepath: str) -> None:
             f"Append delta written: {output_path}, {time.perf_counter()-start_time}s"
         )
 
-    def build_meta_query(
+    def build_keyset_paginated_metadata_query(
         self,
         table: str,
-        limit: int | None,
-        where: str | None,
+        *,
+        limit: int | None = None,
+        where: str | None = None,
+        keyset_value: tuple[int, int, int] = (0, 0, 0),
         **filters: Unpack["DatasetFilters"],
     ) -> str:
         """Build SQL query using SQLAlchemy against metadata schema tables and views."""
         sa_table = self.get_sa_table(table)
 
-        # build WHERE clause filter expression based on any passed key/value filters
-        # and/or an explicit WHERE string
-        filter_expr = build_filter_expr_sa(sa_table, **filters)
-        if where is not None and where.strip():
-            text_where = text(where)
-            combined = (
-                and_(filter_expr, text_where) if filter_expr is not None else text_where
-            )
-        else:
-            combined = filter_expr
-
         # create SQL statement object
         stmt = select(
             sa_table.c.timdex_record_id,
             sa_table.c.run_id,
+            func.hash(sa_table.c.run_id).label("run_id_hash"),
             sa_table.c.run_record_offset,
             sa_table.c.filename,
+            func.hash(sa_table.c.filename).label("filename_hash"),
         ).select_from(sa_table)
-        if combined is not None:
-            stmt = stmt.where(combined)
+
+        # filter expressions from key/value filters (may return None)
+        filter_expr = build_filter_expr_sa(sa_table, **filters)
+        if filter_expr is not None:
+            stmt = stmt.where(filter_expr)
+
+        # explicit raw WHERE string
+        if where is not None and where.strip():
+            stmt = stmt.where(text(where))
+
+        # keyset pagination
+        filename_has, run_id_hash, run_record_offset_ = keyset_value
+        stmt = stmt.where(
+            tuple_(
+                func.hash(sa_table.c.filename),
+                func.hash(sa_table.c.run_id),
+                sa_table.c.run_record_offset,
+            )
+            > tuple_(
+                literal(filename_has),
+                literal(run_id_hash),
+                literal(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),
+            func.hash(sa_table.c.run_id),
             sa_table.c.run_record_offset,
         )
 
@@ -667,7 +681,4 @@ def build_meta_query(
             dialect=DuckDBDialect(),
             compile_kwargs={"literal_binds": True},
         )
-        compiled_str = str(compiled)
-        logger.debug(compiled_str)
-
-        return compiled_str
+        return str(compiled)