feat: improve typing and docstring for to_pyarrow

NickCrews · cpcloud · commit 2600c65f72ec · 2025-05-22T17:48:39.000+08:00
diff --git a/ibis/backends/__init__.py b/ibis/backends/__init__.py
@@ -12,7 +12,7 @@
 import weakref
 from collections import Counter
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple
+from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple, overload
 
 import ibis
 import ibis.common.exceptions as exc
@@ -179,6 +179,39 @@ def to_pandas_batches(
             )
         )
 
+    @overload
+    def to_pyarrow(
+        self,
+        expr: ir.Table,
+        /,
+        *,
+        params: Mapping[ir.Scalar, Any] | None = None,
+        limit: int | str | None = None,
+        **kwargs: Any,
+    ) -> pa.Table: ...
+
+    @overload
+    def to_pyarrow(
+        self,
+        expr: ir.Column,
+        /,
+        *,
+        params: Mapping[ir.Scalar, Any] | None = None,
+        limit: int | str | None = None,
+        **kwargs: Any,
+    ) -> pa.Array: ...
+
+    @overload
+    def to_pyarrow(
+        self,
+        expr: ir.Scalar,
+        /,
+        *,
+        params: Mapping[ir.Scalar, Any] | None = None,
+        limit: int | str | None = None,
+        **kwargs: Any,
+    ) -> pa.Scalar: ...
+
     @util.experimental
     def to_pyarrow(
         self,
@@ -188,8 +221,8 @@ def to_pyarrow(
         params: Mapping[ir.Scalar, Any] | None = None,
         limit: int | str | None = None,
         **kwargs: Any,
-    ) -> pa.Table:
-        """Execute expression and return results in as a pyarrow table.
+    ) -> pa.Table | pa.Array | pa.Scalar:
+        """Execute expression to a pyarrow object.
 
         This method is eager and will execute the associated expression
         immediately.
@@ -208,9 +241,10 @@ def to_pyarrow(
 
         Returns
         -------
-        Table
-            A pyarrow table holding the results of the executed expression.
-
+        result
+            If the passed expression is a Table, a pyarrow table is returned.
+            If the passed expression is a Column, a pyarrow array is returned.
+            If the passed expression is a Scalar, a pyarrow scalar is returned.
         """
         pa = self._import_pyarrow()
         self._run_pre_execute_hooks(expr)
diff --git a/ibis/expr/types/core.py b/ibis/expr/types/core.py
@@ -555,8 +555,8 @@ def to_pyarrow(
         params: Mapping[ir.Scalar, Any] | None = None,
         limit: int | str | None = None,
         **kwargs: Any,
-    ) -> pa.Table:
-        """Execute expression and return results in as a pyarrow table.
+    ) -> pa.Table | pa.Array | pa.Scalar:
+        """Execute expression to a pyarrow object.
 
         This method is eager and will execute the associated expression
         immediately.
@@ -567,14 +567,16 @@ def to_pyarrow(
             Mapping of scalar parameter expressions to value.
         limit
             An integer to effect a specific row limit. A value of `None` means
-            "no limit". The default is in `ibis/config.py`.
+            no limit. The default is in `ibis/config.py`.
         kwargs
             Keyword arguments
 
         Returns
         -------
-        Table
-            A pyarrow table holding the results of the executed expression.
+        result
+            If the passed expression is a Table, a pyarrow table is returned.
+            If the passed expression is a Column, a pyarrow array is returned.
+            If the passed expression is a Scalar, a pyarrow scalar is returned.
         """
         return self._find_backend(use_default=True).to_pyarrow(
             self, params=params, limit=limit, **kwargs
diff --git a/ibis/expr/types/generic.py b/ibis/expr/types/generic.py
@@ -14,7 +14,7 @@
 from ibis.common.grounds import Singleton
 from ibis.expr.rewrites import rewrite_window_input
 from ibis.expr.types.core import Expr, _binop, _FixedTextJupyterMixin
-from ibis.util import deprecated, promote_list
+from ibis.util import deprecated, experimental, promote_list
 
 if TYPE_CHECKING:
     import datetime
@@ -1385,6 +1385,17 @@ def to_pandas(
 
 @public
 class Scalar(Value):
+    # overriding Expr's implementation just for typing
+    @experimental
+    def to_pyarrow(
+        self,
+        *,
+        params: Mapping[ir.Scalar, Any] | None = None,
+        limit: int | str | None = None,
+        **kwargs: Any,
+    ) -> pa.Scalar:
+        return super().to_pyarrow(params=params, limit=limit, **kwargs)
+
     def __pyarrow_result__(
         self,
         table: pa.Table,
@@ -1575,6 +1586,17 @@ def preview(
             console_width=console_width,
         )
 
+    # overriding Expr's implementation just for typing
+    @experimental
+    def to_pyarrow(
+        self,
+        *,
+        params: Mapping[ir.Scalar, Any] | None = None,
+        limit: int | str | None = None,
+        **kwargs: Any,
+    ) -> pa.Array:
+        return super().to_pyarrow(params=params, limit=limit, **kwargs)
+
     def __pyarrow_result__(
         self,
         table: pa.Table,
diff --git a/ibis/expr/types/relations.py b/ibis/expr/types/relations.py
@@ -24,7 +24,7 @@
 from ibis.expr.types.core import Expr, _FixedTextJupyterMixin
 from ibis.expr.types.generic import Value, literal
 from ibis.expr.types.temporal import TimestampColumn
-from ibis.util import deprecated
+from ibis.util import deprecated, experimental
 
 if TYPE_CHECKING:
     import pandas as pd
@@ -242,6 +242,17 @@ def __polars_result__(self, df: pl.DataFrame) -> Any:
 
         return PolarsData.convert_table(df, self.schema())
 
+    # overriding Expr's implementation just for typing
+    @experimental
+    def to_pyarrow(
+        self,
+        *,
+        params: Mapping[ir.Scalar, Any] | None = None,
+        limit: int | str | None = None,
+        **kwargs: Any,
+    ) -> pa.Table:
+        return super().to_pyarrow(params=params, limit=limit, **kwargs)
+
     def _fast_bind(self, *args, **kwargs):
         # allow the first argument to be either a dictionary or a list of values
         if len(args) == 1:
