polars-runtime-compat 1.34.0b2__cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

polars/io/database/__init__.py

@@ -0,0 +1,6 @@
+from polars.io.database.functions import read_database, read_database_uri
+
+__all__ = [
+    "read_database",
+    "read_database_uri",
+]

polars/io/database/_arrow_registry.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+from typing import TypedDict
+
+
+class ArrowDriverProperties(TypedDict):
+    # name of the method that fetches all arrow data; tuple form
+    # calls the fetch_all method with the given chunk size (int)
+    fetch_all: str
+    # name of the method that fetches arrow data in batches
+    fetch_batches: str | None
+    # indicate whether the given batch size is respected exactly
+    exact_batch_size: bool | None
+    # repeat batch calls (if False, the batch call is a generator)
+    repeat_batch_calls: bool
+    # if arrow/polars functionality requires a minimum module version
+    minimum_version: str | None
+
+
+ARROW_DRIVER_REGISTRY: dict[str, ArrowDriverProperties] = {
+    "adbc_.*": {
+        "fetch_all": "fetch_arrow_table",
+        "fetch_batches": None,
+        "exact_batch_size": None,
+        "repeat_batch_calls": False,
+        "minimum_version": None,
+    },
+    "arrow_odbc_proxy": {
+        "fetch_all": "fetch_arrow_table",
+        "fetch_batches": "fetch_record_batches",
+        "exact_batch_size": True,
+        "repeat_batch_calls": False,
+        "minimum_version": None,
+    },
+    "databricks": {
+        "fetch_all": "fetchall_arrow",
+        "fetch_batches": "fetchmany_arrow",
+        "exact_batch_size": True,
+        "repeat_batch_calls": True,
+        "minimum_version": None,
+    },
+    "duckdb": {
+        "fetch_all": "fetch_arrow_table",
+        "fetch_batches": "fetch_record_batch",
+        "exact_batch_size": True,
+        "repeat_batch_calls": False,
+        "minimum_version": None,
+    },
+    "kuzu": {
+        "fetch_all": "get_as_pl",
+        "fetch_batches": None,
+        "exact_batch_size": None,
+        "repeat_batch_calls": False,
+        "minimum_version": "0.3.2",
+    },
+    "snowflake": {
+        "fetch_all": "fetch_arrow_all",
+        "fetch_batches": "fetch_arrow_batches",
+        "exact_batch_size": False,
+        "repeat_batch_calls": False,
+        "minimum_version": None,
+    },
+    "turbodbc": {
+        "fetch_all": "fetchallarrow",
+        "fetch_batches": "fetcharrowbatches",
+        "exact_batch_size": False,
+        "repeat_batch_calls": False,
+        "minimum_version": None,
+    },
+}

polars/io/database/_cursor_proxies.py

@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from polars._dependencies import import_optional
+from polars.io.database._utils import _run_async
+
+if TYPE_CHECKING:
+    import sys
+    from collections.abc import Coroutine, Iterable
+
+    import pyarrow as pa
+
+    if sys.version_info >= (3, 11):
+        from typing import Self
+    else:
+        from typing_extensions import Self
+
+
+class ODBCCursorProxy:
+    """Cursor proxy for ODBC connections (requires `arrow-odbc`)."""
+
+    def __init__(self, connection_string: str) -> None:
+        self.connection_string = connection_string
+        self.execute_options: dict[str, Any] = {}
+        self.query: str | None = None
+
+    def close(self) -> None:
+        """Close the cursor."""
+        # n/a: nothing to close
+
+    def execute(self, query: str, **execute_options: Any) -> None:
+        """Execute a query (n/a: just store query for the fetch* methods)."""
+        self.execute_options = execute_options
+        self.query = query
+
+    def fetch_arrow_table(
+        self, batch_size: int = 10_000, *, fetch_all: bool = False
+    ) -> pa.Table:
+        """Fetch all results as a pyarrow Table."""
+        from pyarrow import Table
+
+        return Table.from_batches(
+            self.fetch_record_batches(batch_size=batch_size, fetch_all=True)
+        )
+
+    def fetch_record_batches(
+        self, batch_size: int = 10_000, *, fetch_all: bool = False
+    ) -> Iterable[pa.RecordBatch]:
+        """Fetch results as an iterable of RecordBatches."""
+        from arrow_odbc import read_arrow_batches_from_odbc
+        from pyarrow import RecordBatch
+
+        n_batches = 0
+        batch_reader = read_arrow_batches_from_odbc(
+            query=self.query,
+            batch_size=batch_size,
+            connection_string=self.connection_string,
+            **self.execute_options,
+        )
+        for batch in batch_reader:
+            yield batch
+            n_batches += 1
+
+        if n_batches == 0 and fetch_all:
+            # empty result set; return empty batch with accurate schema
+            yield RecordBatch.from_pylist([], schema=batch_reader.schema)
+
+    # note: internally arrow-odbc always reads batches
+    fetchall = fetch_arrow_table
+    fetchmany = fetch_record_batches
+
+
+class SurrealDBCursorProxy:
+    """Cursor proxy for both SurrealDB and AsyncSurrealDB connections."""
+
+    _cached_result: list[dict[str, Any]] | None = None
+
+    def __init__(self, client: Any) -> None:
+        surrealdb = import_optional("surrealdb")
+        self.is_async = isinstance(client, surrealdb.AsyncSurrealDB)
+        self.execute_options: dict[str, Any] = {}
+        self.client = client
+        self.query: str = None  # type: ignore[assignment]
+
+    @staticmethod
+    async def _unpack_result_async(
+        result: Coroutine[Any, Any, list[dict[str, Any]]],
+    ) -> Coroutine[Any, Any, list[dict[str, Any]]]:
+        """Unpack the async query result."""
+        response = (await result)[0]
+        if response["status"] != "OK":
+            raise RuntimeError(response["result"])
+        return response["result"]
+
+    @staticmethod
+    def _unpack_result(
+        result: list[dict[str, Any]],
+    ) -> list[dict[str, Any]]:
+        """Unpack the query result."""
+        response = result[0]
+        if response["status"] != "OK":
+            raise RuntimeError(response["result"])
+        return response["result"]
+
+    def close(self) -> None:
+        """Close the cursor."""
+        # no-op; never close a user's Surreal session
+
+    def execute(self, query: str, **execute_options: Any) -> Self:
+        """Execute a query (n/a: just store query for the fetch* methods)."""
+        self._cached_result = None
+        self.execute_options = execute_options
+        self.query = query
+        return self
+
+    def fetchall(self) -> list[dict[str, Any]]:
+        """Fetch all results (as a list of dictionaries)."""
+        return (
+            _run_async(
+                self._unpack_result_async(
+                    result=self.client.query(
+                        query=self.query,
+                        variables=(self.execute_options or None),
+                    ),
+                )
+            )
+            if self.is_async
+            else self._unpack_result(
+                result=self.client.query(
+                    query=self.query,
+                    variables=(self.execute_options or None),
+                ),
+            )
+        )
+
+    def fetchmany(self, size: int) -> list[dict[str, Any]]:
+        """Fetch results in batches (simulated)."""
+        # first 'fetchmany' call acquires/caches the result object
+        if self._cached_result is None:
+            self._cached_result = self.fetchall()
+
+        # return batches from the result, actively removing from the cache
+        # as we go, so as not to hold on to additional copies when done
+        result = self._cached_result[:size]
+        del self._cached_result[:size]
+        return result