feat: add Arrow type system foundation for select_to_arrow() (#154)

Add foundational type system and utilities for Apache Arrow integration.

Author: cofin

pyproject.toml

@@ -104,6 +104,7 @@ lint = [
   "types-protobuf",
   "asyncpg-stubs",
   "pyarrow-stubs",
+  "pandas-stubs",
 
 ]
 test = [

sqlspec/_typing.py

@@ -377,18 +377,101 @@ def slice(self, offset: int = 0, length: "int | None" = None) -> Any:
         return None
 
 
+@runtime_checkable
+class ArrowSchemaProtocol(Protocol):
+    """Typed shim for pyarrow.Schema."""
+
+    def field(self, i: int) -> Any:
+        """Get field by index."""
+        ...
+
+    @property
+    def names(self) -> "list[str]":
+        """Get list of field names."""
+        ...
+
+    def __len__(self) -> int:
+        """Get number of fields."""
+        return 0
+
+
+@runtime_checkable
+class ArrowRecordBatchReaderProtocol(Protocol):
+    """Typed shim for pyarrow.RecordBatchReader."""
+
+    def read_all(self) -> Any:
+        """Read all batches into a table."""
+        ...
+
+    def read_next_batch(self) -> Any:
+        """Read next batch."""
+        ...
+
+    def __iter__(self) -> "Iterable[Any]":
+        """Iterate over batches."""
+        ...
+
+
 try:
     from pyarrow import RecordBatch as ArrowRecordBatch
+    from pyarrow import RecordBatchReader as ArrowRecordBatchReader
+    from pyarrow import Schema as ArrowSchema
     from pyarrow import Table as ArrowTable
 
     PYARROW_INSTALLED = True
 except ImportError:
     ArrowTable = ArrowTableResult  # type: ignore[assignment,misc]
     ArrowRecordBatch = ArrowRecordBatchResult  # type: ignore[assignment,misc]
+    ArrowSchema = ArrowSchemaProtocol  # type: ignore[assignment,misc]
+    ArrowRecordBatchReader = ArrowRecordBatchReaderProtocol  # type: ignore[assignment,misc]
 
     PYARROW_INSTALLED = False  # pyright: ignore[reportConstantRedefinition]
 
 
+@runtime_checkable
+class PandasDataFrameProtocol(Protocol):
+    """Typed shim for pandas.DataFrame."""
+
+    def __len__(self) -> int:
+        """Get number of rows."""
+        ...
+
+    def __getitem__(self, key: Any) -> Any:
+        """Get column or row."""
+        ...
+
+
+@runtime_checkable
+class PolarsDataFrameProtocol(Protocol):
+    """Typed shim for polars.DataFrame."""
+
+    def __len__(self) -> int:
+        """Get number of rows."""
+        ...
+
+    def __getitem__(self, key: Any) -> Any:
+        """Get column or row."""
+        ...
+
+
+try:
+    from pandas import DataFrame as PandasDataFrame
+
+    PANDAS_INSTALLED = True
+except ImportError:
+    PandasDataFrame = PandasDataFrameProtocol  # type: ignore[assignment,misc]
+    PANDAS_INSTALLED = False
+
+
+try:
+    from polars import DataFrame as PolarsDataFrame
+
+    POLARS_INSTALLED = True
+except ImportError:
+    PolarsDataFrame = PolarsDataFrameProtocol  # type: ignore[assignment,misc]
+    POLARS_INSTALLED = False
+
+
 @runtime_checkable
 class NumpyArrayStub(Protocol):
     """Protocol stub for numpy.ndarray when numpy is not installed.
@@ -639,7 +722,9 @@ async def insert_returning(self, conn: Any, query_name: str, sql: str, parameter
     "OBSTORE_INSTALLED",
     "OPENTELEMETRY_INSTALLED",
     "ORJSON_INSTALLED",
+    "PANDAS_INSTALLED",
     "PGVECTOR_INSTALLED",
+    "POLARS_INSTALLED",
     "PROMETHEUS_INSTALLED",
     "PYARROW_INSTALLED",
     "PYDANTIC_INSTALLED",
@@ -650,7 +735,9 @@ async def insert_returning(self, conn: Any, query_name: str, sql: str, parameter
     "AiosqlSQLOperationType",
     "AiosqlSyncProtocol",
     "ArrowRecordBatch",
+    "ArrowRecordBatchReader",
     "ArrowRecordBatchResult",
+    "ArrowSchema",
     "ArrowTable",
     "ArrowTableResult",
     "AttrsInstance",
@@ -670,6 +757,10 @@ async def insert_returning(self, conn: Any, query_name: str, sql: str, parameter
     "Histogram",
     "NumpyArray",
     "NumpyArrayStub",
+    "PandasDataFrame",
+    "PandasDataFrameProtocol",
+    "PolarsDataFrame",
+    "PolarsDataFrameProtocol",
     "Span",
     "Status",
     "StatusCode",

sqlspec/adapters/adbc/driver.py

@@ -23,7 +23,6 @@
     DataError,
     ForeignKeyViolationError,
     IntegrityError,
-    MissingDependencyError,
     NotNullViolationError,
     SQLParsingError,
     SQLSpecError,
@@ -507,18 +506,6 @@ def __init__(
         self.dialect = statement_config.dialect
         self._data_dictionary: SyncDataDictionaryBase | None = None
 
-    @staticmethod
-    def _ensure_pyarrow_installed() -> None:
-        """Ensure PyArrow is installed.
-
-        Raises:
-            MissingDependencyError: If PyArrow is not installed
-        """
-        from sqlspec.typing import PYARROW_INSTALLED
-
-        if not PYARROW_INSTALLED:
-            raise MissingDependencyError(package="pyarrow", install_package="arrow")
-
     @staticmethod
     def _get_dialect(connection: "AdbcConnection") -> str:
         """Detect database dialect from connection information.

sqlspec/core/result.py

@@ -16,13 +16,14 @@
 from typing_extensions import TypeVar
 
 from sqlspec.core.compiler import OperationType
+from sqlspec.utils.module_loader import ensure_pandas, ensure_polars, ensure_pyarrow
 from sqlspec.utils.schema import to_schema
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
 
     from sqlspec.core.statement import SQL
-    from sqlspec.typing import SchemaT
+    from sqlspec.typing import ArrowTable, PandasDataFrame, PolarsDataFrame, SchemaT
 
 
 __all__ = ("ArrowResult", "SQLResult", "StatementResult")
@@ -618,18 +619,27 @@ def is_success(self) -> bool:
         """
         return self.data is not None
 
-    def get_data(self) -> Any:
+    def get_data(self) -> "ArrowTable":
         """Get the Apache Arrow Table from the result.
 
         Returns:
             The Arrow table containing the result data.
 
         Raises:
             ValueError: If no Arrow table is available.
+            TypeError: If data is not an Arrow Table.
         """
         if self.data is None:
             msg = "No Arrow table available for this result"
             raise ValueError(msg)
+
+        ensure_pyarrow()
+
+        import pyarrow as pa
+
+        if not isinstance(self.data, pa.Table):
+            msg = f"Expected an Arrow Table, but got {type(self.data).__name__}"
+            raise TypeError(msg)
         return self.data
 
     @property
@@ -680,6 +690,127 @@ def num_columns(self) -> int:
 
         return cast("int", self.data.num_columns)
 
+    def to_pandas(self) -> "PandasDataFrame":
+        """Convert Arrow data to pandas DataFrame.
+
+        Returns:
+            pandas DataFrame containing the result data.
+
+        Raises:
+            ValueError: If no Arrow table is available.
+
+        Examples:
+            >>> result = session.select_to_arrow("SELECT * FROM users")
+            >>> df = result.to_pandas()
+            >>> print(df.head())
+        """
+        if self.data is None:
+            msg = "No Arrow table available"
+            raise ValueError(msg)
+
+        ensure_pandas()
+
+        import pandas as pd
+
+        result = self.data.to_pandas()
+        if not isinstance(result, pd.DataFrame):
+            msg = f"Expected a pandas DataFrame, but got {type(result).__name__}"
+            raise TypeError(msg)
+        return result
+
+    def to_polars(self) -> "PolarsDataFrame":
+        """Convert Arrow data to Polars DataFrame.
+
+        Returns:
+            Polars DataFrame containing the result data.
+
+        Raises:
+            ValueError: If no Arrow table is available.
+
+        Examples:
+            >>> result = session.select_to_arrow("SELECT * FROM users")
+            >>> df = result.to_polars()
+            >>> print(df.head())
+        """
+        if self.data is None:
+            msg = "No Arrow table available"
+            raise ValueError(msg)
+
+        ensure_polars()
+
+        import polars as pl
+
+        result = pl.from_arrow(self.data)
+        if not isinstance(result, pl.DataFrame):
+            msg = f"Expected a Polars DataFrame, but got {type(result).__name__}"
+            raise TypeError(msg)
+        return result
+
+    def to_dict(self) -> "list[dict[str, Any]]":
+        """Convert Arrow data to list of dictionaries.
+
+        Returns:
+            List of dictionaries, one per row.
+
+        Raises:
+            ValueError: If no Arrow table is available.
+
+        Examples:
+            >>> result = session.select_to_arrow(
+            ...     "SELECT id, name FROM users"
+            ... )
+            >>> rows = result.to_dict()
+            >>> print(rows[0])
+            {'id': 1, 'name': 'Alice'}
+        """
+        if self.data is None:
+            msg = "No Arrow table available"
+            raise ValueError(msg)
+
+        return cast("list[dict[str, Any]]", self.data.to_pylist())
+
+    def __len__(self) -> int:
+        """Return number of rows in the Arrow table.
+
+        Returns:
+            Number of rows.
+
+        Raises:
+            ValueError: If no Arrow table is available.
+
+        Examples:
+            >>> result = session.select_to_arrow("SELECT * FROM users")
+            >>> print(len(result))
+            100
+        """
+        if self.data is None:
+            msg = "No Arrow table available"
+            raise ValueError(msg)
+
+        return cast("int", self.data.num_rows)
+
+    def __iter__(self) -> "Iterator[dict[str, Any]]":
+        """Iterate over rows as dictionaries.
+
+        Yields:
+            Dictionary for each row.
+
+        Raises:
+            ValueError: If no Arrow table is available.
+
+        Examples:
+            >>> result = session.select_to_arrow(
+            ...     "SELECT id, name FROM users"
+            ... )
+            >>> for row in result:
+            ...     print(row["name"])
+        """
+        if self.data is None:
+            msg = "No Arrow table available"
+            raise ValueError(msg)
+
+        yield from self.data.to_pylist()
+
 
 def create_sql_result(
     statement: "SQL",

sqlspec/extensions/aiosql/adapter.py

@@ -13,14 +13,8 @@
 from sqlspec.core.result import SQLResult
 from sqlspec.core.statement import SQL, StatementConfig
 from sqlspec.driver import AsyncDriverAdapterBase, SyncDriverAdapterBase
-from sqlspec.exceptions import MissingDependencyError
-from sqlspec.typing import (
-    AIOSQL_INSTALLED,
-    AiosqlAsyncProtocol,
-    AiosqlParamType,
-    AiosqlSQLOperationType,
-    AiosqlSyncProtocol,
-)
+from sqlspec.typing import AiosqlAsyncProtocol, AiosqlParamType, AiosqlSQLOperationType, AiosqlSyncProtocol
+from sqlspec.utils.module_loader import ensure_aiosql
 
 logger = logging.getLogger("sqlspec.extensions.aiosql")
 
@@ -58,12 +52,6 @@ def fetchone(self) -> Any | None:
         return rows[0] if rows else None
 
 
-def _check_aiosql_available() -> None:
-    if not AIOSQL_INSTALLED:
-        msg = "aiosql"
-        raise MissingDependencyError(msg, "aiosql")
-
-
 def _normalize_dialect(dialect: "str | Any | None") -> str:
     """Normalize dialect name for SQLGlot compatibility.
 
@@ -105,7 +93,7 @@ def __init__(self, driver: DriverT) -> None:
         Args:
             driver: SQLSpec driver to use for execution.
         """
-        _check_aiosql_available()
+        ensure_aiosql()
         self.driver: DriverT = driver
 
     def process_sql(self, query_name: str, op_type: "AiosqlSQLOperationType", sql: str) -> str: