From 791335f18e4d08f9c86098318fee7dfe9b8b1118 Mon Sep 17 00:00:00 2001 From: Gil Forsyth Date: Mon, 3 Oct 2022 12:12:29 -0400 Subject: [PATCH] feat: add experimental decorator --- ibis/backends/base/__init__.py | 45 ++++++++++++++++++++++++++++++ ibis/backends/base/sql/__init__.py | 24 ++++++++++++++++ ibis/expr/types/core.py | 17 +++++------ ibis/util.py | 28 +++++++++++++++++++ 4 files changed, 106 insertions(+), 8 deletions(-) diff --git a/ibis/backends/base/__init__.py b/ibis/backends/base/__init__.py index f5a36e3e0d46..d64b57255ba6 100644 --- a/ibis/backends/base/__init__.py +++ b/ibis/backends/base/__init__.py @@ -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, @@ -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 @@ -269,6 +290,7 @@ def to_pyarrow( else: raise ValueError + @util.experimental def to_pyarrow_batches( self, expr: ir.Expr, @@ -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 diff --git a/ibis/backends/base/sql/__init__.py b/ibis/backends/base/sql/__init__.py index 87e2e56021bb..31f204260cb6 100644 --- a/ibis/backends/base/sql/__init__.py +++ b/ibis/backends/base/sql/__init__.py @@ -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, @@ -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 diff --git a/ibis/expr/types/core.py b/ibis/expr/types/core.py index 1b9277920fb3..15fefb8c1592 100644 --- a/ibis/expr/types/core.py +++ b/ibis/expr/types/core.py @@ -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 @@ -308,6 +308,7 @@ def compile( self, limit=limit, timecontext=timecontext, params=params ) + @experimental def to_pyarrow_batches( self, *, @@ -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 ---------- @@ -345,6 +346,7 @@ def to_pyarrow_batches( **kwargs, ) + @experimental def to_pyarrow( self, *, @@ -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 ------- diff --git a/ibis/util.py b/ibis/util.py index 747a6f225eeb..c02519acf54b 100644 --- a/ibis/util.py +++ b/ibis/util.py @@ -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."""