limit pushdown

bmsdna/lakeapi/context/df_base.py

@@ -384,6 +384,7 @@ def register_datasource(
         file_type: FileTypes,
         filters: Optional[FilterType],
         meta_only: bool = False,
+        limit: int | None = None,
     ):
         ds = self.get_pyarrow_dataset(uri, file_type, filters)
         self.modified_dates[target_name] = self.get_modified_date(uri, file_type)

bmsdna/lakeapi/context/df_duckdb.py

@@ -343,6 +343,7 @@ def register_datasource(
         file_type: FileTypes,
         filters: Optional[FilterType],
         meta_only: bool = False,
+        limit: int | None = None,
     ):
         self.modified_dates[target_name] = self.get_modified_date(uri, file_type)
 
@@ -390,6 +391,7 @@ def register_datasource(
                 storage_options=uri_opts,
                 conditions=filters,
                 use_fsspec=os.getenv("DUCKDB_DELTA_USE_FSSPEC", "0") == "1",
+                limit=limit,
             )
             return
         if file_type == "duckdb":

bmsdna/lakeapi/context/df_odbc.py

@@ -186,6 +186,7 @@ def register_datasource(
         file_type: FileTypes,
         filters: Any,
         meta_only: bool = False,
+        limit: int | None = None,
     ):
         assert file_type == "odbc"
         assert uri.account is None

bmsdna/lakeapi/context/df_polars.py

@@ -231,6 +231,7 @@ def register_datasource(
         file_type: FileTypes,
         filters: Optional[FilterType],
         meta_only: bool = False,
+        limit: int | None = None,
     ):
         import polars as pl
 
@@ -245,14 +246,15 @@ def register_datasource(
                     db_uri, db_opts = uri.get_uri_options(flavor="original")
                     from deltalake2db import polars_scan_delta, get_polars_schema
 
-                    if meta_only:
+                    if meta_only or limit == 0:
                         schema = get_polars_schema(db_uri, storage_options=db_opts)
                         df = pl.DataFrame(schema=schema)
                     else:
                         df = polars_scan_delta(
                             db_uri,
                             storage_options=db_opts,
                             conditions=filters,
+                            limit=limit,
                         )
                 except DeltaProtocolError as de:
                     raise FileTypeNotSupportedError(

bmsdna/lakeapi/core/config.py

@@ -160,6 +160,12 @@ class DuckDBBackendConfig:
     index: List[str] | None = None
 
 
+def _to_dict(d: Union[list, dict]):
+    if isinstance(d, list):
+        return {i["key"]: i["value"] for i in d}
+    return d
+
+
 @dataclass
 class Config:
     name: str
@@ -225,7 +231,7 @@ def _from_dict(cls, config: Dict, basic_config: BasicConfig, accounts: dict):
                 meta = get_deltalake_meta(uri_obj)
                 assert meta.last_metadata is not None
                 cfg = json.loads(
-                    meta.last_metadata.get("configuration", {}).get(
+                    _to_dict(meta.last_metadata.get("configuration", {})).get(
                         "lakeapi.config", "{}"
                     )
                 )

bmsdna/lakeapi/core/datasource.py

@@ -28,7 +28,7 @@
 )
 from bmsdna.lakeapi.core.log import get_logger
 from bmsdna.lakeapi.core.model import get_param_def
-from bmsdna.lakeapi.core.types import DeltaOperatorTypes, OperatorType
+from bmsdna.lakeapi.core.types import OperatorType
 from deltalake2db.delta_meta_retrieval import (
     DataType,
     field_to_type,
@@ -233,6 +233,7 @@ def get_df(
         self,
         filters: Optional[FilterType] = None,
         endpoint: endpoints = "request",
+        limit: int | None = None,
     ) -> ResultData:
         if self.df is None:
             unique_table_name = (
@@ -254,6 +255,7 @@ def get_df(
                     self.config.file_type,
                     filters=filters,
                     meta_only=endpoint == "meta",
+                    limit=limit,
                 )
                 self.df = self.sql_context.execute_sql(self.query)
 
@@ -264,7 +266,7 @@ def get_partition_filter(
     param: tuple[str, str],
     deltaMeta: DeltaMetadata,
     param_def: Iterable[Param | str],
-) -> Optional[Tuple[str, OperatorType, Any]]:
+) -> Optional[Tuple[str, SupportedOperator, Any]]:
     operators = ("<=", ">=", "=", "==", "in", "not in")
     key, value = param
     if not key or not value or key in ("limit", "offset"):
@@ -274,7 +276,7 @@ def get_partition_filter(
     if prmdef_and_op is None:
         raise ValueError(f"thats not parameter: {key}")
     prmdef, op = prmdef_and_op
-    if op not in get_args(DeltaOperatorTypes):
+    if op not in get_args(SupportedOperator):
         return None
     colname = prmdef.colname or prmdef.name
 
@@ -337,7 +339,7 @@ def get_partition_filter(
     # partition value must be string
     return (
         col_for_partitioning,
-        op,
+        cast(SupportedOperator, op),
         value_for_partitioning,
     )
 

bmsdna/lakeapi/core/types.py

@@ -54,7 +54,6 @@
     "has",
     "startswith",
 ]
-DeltaOperatorTypes = Literal["<", "=", ">", ">=", "<=", "in", "not in"]
 PolaryTypeFunction = Literal[
     "sum",
     "count",

bmsdna/lakeapi/endpoint/endpoint.py

@@ -207,6 +207,8 @@ async def data(
             engine,
             chunk_size=real_chunk_size,
         )
+        if not (limit == -1 and config.allow_get_all_pages):
+            limit = (1000 if limit == -1 else limit) or 1000
         assert config.datasource is not None
         with Datasource(
             config.version_str,
@@ -225,7 +227,20 @@ async def data(
                 )
             else:
                 pre_filter = None
-            df = realdataframe.get_df(filters=pre_filter)
+            if config.datasource.file_type == "delta" and config.params is not None:
+                st = realdataframe.get_delta_table(True)
+                if st is not None and params is not None:
+                    part_filter = filter_partitions_based_on_params(
+                        st, params.model_dump(exclude_unset=True), config.params
+                    )
+                    if part_filter:
+                        pre_filter = list(pre_filter or []) + part_filter
+            df = realdataframe.get_df(
+                filters=pre_filter,
+                limit=limit
+                if not config.datasource.sortby and not offset and not limit == -1
+                else None,  # if sorted we need all data to sort correctly
+            )
             df_cols = df.columns()
             expr = get_params_filter_expr(  # this supports all kinds of filters, while the prefilter only supports equality
                 context,

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bmsdna-lakeapi"
-version = "0.25.1"
+version = "0.25.2"
 description = ""
 authors = [{ name = "DWH Team", email = "you@example.com" }]
 dependencies = [
@@ -16,10 +16,9 @@ dependencies = [
     "fastapi >=0.110.0",
     "arrow-odbc >=5.0.0",
     "expandvars >=0.12.0",
-    "pandas >=2.1.0,<3",
     "fsspec >=2024.2.0,<2025",
     "adlfs >=2024.2.0,<2025",
-    "deltalake2db >=1.1.1",
+    "deltalake2db >=1.1.4",
 ]
 requires-python = "~=3.10"
 
@@ -32,6 +31,7 @@ polars = ["fastexcel"]
 auth = ["argon2-cffi", "pyjwt"]
 useradd = ["ruamel.yaml"]
 odbc = ["arrow-odbc"]
+pandas = ["pandas >=2.1.0,<3"]
 
 [build-system]
 requires = ["hatchling"]

tests/test_app.py

@@ -77,21 +77,20 @@ def test_fruits_sort_desc(client: TestClient):
             ]
 
 
-def test_fruits_offset_1(client: TestClient):
+@pytest.mark.parametrize("engine", engines)
+def test_fruits_offset_1(engine, client: TestClient):
     for _ in range(2):
-        for e in engines:  # polars does not support offset
-            response = client.get(
-                f"/api/v1/test/fruits?limit=1&&offset=1&format=json&%24engine={e}",
-            )
-            assert response.status_code == 200
-            assert response.json() == [
-                {
-                    "A": 2,
-                    "fruits": "banana",
-                    "B": 4,
-                    "cars": "audi",
-                }
-            ]
+        response = client.get(
+            f"/api/v1/test/fruits?limit=1&&offset=0&format=json&%24engine={engine}",
+        )
+        assert response.status_code == 200
+        res = response.json()
+        response = client.get(
+            f"/api/v1/test/fruits?limit=1&&offset=1&format=json&%24engine={engine}",
+        )
+        assert response.status_code == 200
+        res2 = response.json()
+        assert res != res2
 
 
 def test_data_limit(client: TestClient):