Skip to content

Commit

Permalink
feat: add experimental decorator
Browse files Browse the repository at this point in the history
  • Loading branch information
gforsyth authored and kszucs committed Oct 5, 2022
1 parent 3a88170 commit 791335f
Show file tree
Hide file tree
Showing 4 changed files with 106 additions and 8 deletions.
45 changes: 45 additions & 0 deletions ibis/backends/base/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -231,6 +231,7 @@ def _table_or_column_schema(expr: ir.Expr) -> sch.Schema:
# ColumnExpr has no schema method, define single-column schema
return sch.schema([(expr.get_name(), expr.type())])

@util.experimental
def to_pyarrow(
self,
expr: ir.Expr,
Expand All @@ -239,6 +240,26 @@ def to_pyarrow(
limit: int | str | None = None,
**kwargs: Any,
) -> pa.Table:
"""Execute expression and return results in as a pyarrow table.
This method is eager and will execute the associated expression
immediately.
Parameters
----------
expr
Ibis expression to export to pyarrow
params
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`.
Returns
-------
Table
A pyarrow table holding the results of the executed expression.
"""
pa = self._import_pyarrow()
try:
# Can't construct an array from record batches
Expand Down Expand Up @@ -269,6 +290,7 @@ def to_pyarrow(
else:
raise ValueError

@util.experimental
def to_pyarrow_batches(
self,
expr: ir.Expr,
Expand All @@ -278,6 +300,29 @@ def to_pyarrow_batches(
chunk_size: int = 1_000_000,
**kwargs: Any,
) -> pa.RecordBatchReader:
"""Execute expression and return results in an iterator of pyarrow
record batches.
This method is eager and will execute the associated expression
immediately.
Parameters
----------
expr
Ibis expression to export to pyarrow
limit
An integer to effect a specific row limit. A value of `None` means
"no limit". The default is in `ibis/config.py`.
params
Mapping of scalar parameter expressions to value.
chunk_size
Number of rows in each returned record batch.
Returns
-------
record_batches
An iterator of pyarrow record batches.
"""
raise NotImplementedError


Expand Down
24 changes: 24 additions & 0 deletions ibis/backends/base/sql/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -158,6 +158,7 @@ def _cursor_batches(
while batch := cursor.fetchmany(chunk_size):
yield batch

@util.experimental
def to_pyarrow_batches(
self,
expr: ir.Expr,
Expand All @@ -167,6 +168,29 @@ def to_pyarrow_batches(
chunk_size: int = 1_000_000,
**kwargs: Any,
) -> pa.RecordBatchReader:
"""Execute expression and return results in an iterator of pyarrow
record batches.
This method is eager and will execute the associated expression
immediately.
Parameters
----------
expr
Ibis expression to export to pyarrow
limit
An integer to effect a specific row limit. A value of `None` means
"no limit". The default is in `ibis/config.py`.
params
Mapping of scalar parameter expressions to value.
chunk_size
Number of rows in each returned record batch.
Returns
-------
record_batches
An iterator of pyarrow record batches.
"""
pa = self._import_pyarrow()

from ibis.backends.pyarrow.datatypes import ibis_to_pyarrow_struct
Expand Down
17 changes: 9 additions & 8 deletions ibis/expr/types/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
from ibis.common.grounds import Immutable
from ibis.config import _default_backend, options
from ibis.expr.typing import TimeContext
from ibis.util import UnnamedMarker
from ibis.util import UnnamedMarker, experimental

if TYPE_CHECKING:
import pyarrow as pa
Expand Down Expand Up @@ -308,6 +308,7 @@ def compile(
self, limit=limit, timecontext=timecontext, params=params
)

@experimental
def to_pyarrow_batches(
self,
*,
Expand All @@ -319,8 +320,8 @@ def to_pyarrow_batches(
"""Execute expression and return results in an iterator of pyarrow
record batches.
**Warning**: This method is eager and will execute the associated
expression immediately. This API is experimental and subject to change.
This method is eager and will execute the associated expression
immediately.
Parameters
----------
Expand All @@ -345,6 +346,7 @@ def to_pyarrow_batches(
**kwargs,
)

@experimental
def to_pyarrow(
self,
*,
Expand All @@ -354,17 +356,16 @@ def to_pyarrow(
) -> pa.Table:
"""Execute expression and return results in as a pyarrow table.
**Warning**: This method is eager and will execute the associated
expression immediately. This API is experimental and subject to change.
This method is eager and will execute the associated expression
immediately.
Parameters
----------
params
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`.
params
Mapping of scalar parameter expressions to value.
Returns
-------
Expand Down
28 changes: 28 additions & 0 deletions ibis/util.py
Original file line number Diff line number Diff line change
Expand Up @@ -448,6 +448,34 @@ def wrapper(*args, **kwargs):
return decorator


def experimental(func):
"""Decorate experimental function to add warning about potential API
instability in docstring."""

msg = "This API is experimental and subject to change."

if docstr := func.__doc__:
preamble, *rest = docstr.split("\n\n", maxsplit=1)

leading_spaces = " " * sum(
1
for _ in itertools.takewhile(str.isspace, rest[0] if rest else [])
)

warning_doc = f'{leading_spaces}!!! warning "{msg}"'

docstr = "\n\n".join([preamble, warning_doc, *rest])
else:
docstr = f'!!! warning "{msg}"'
func.__doc__ = docstr

@functools.wraps(func)
def wrapper(*args, **kwargs):
return func(*args, **kwargs)

return wrapper


class ToFrame(abc.ABC):
"""Interface for in-memory objects that can be converted to a DataFrame."""

Expand Down

0 comments on commit 791335f

Please sign in to comment.