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

Author: ghukill

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

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,25 +474,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:
@@ -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)

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(),