Skip to content

Commit

Permalink
docs(udfs): improve udf docstring examples (#9079)
Browse files Browse the repository at this point in the history
UDF arguments need to be Ibis expressions. In the existing docstring
examples
for the UDF decorators, it can be unclear that the integers being passed
to
`add_one` are implicitly converted to Ibis literals.

To make it clearer, I've added memtables to all the examples, and used
pandas-
and pyarrow-specific functions to the examples.
  • Loading branch information
gforsyth authored Apr 29, 2024
1 parent a9aa8a7 commit fead40b
Showing 1 changed file with 76 additions and 17 deletions.
93 changes: 76 additions & 17 deletions ibis/expr/operations/udf.py
Original file line number Diff line number Diff line change
Expand Up @@ -296,13 +296,44 @@ def python(
Examples
--------
>>> import ibis
>>> ibis.options.interactive = True
>>> t = ibis.memtable(dict(int_col=[1, 2, 3], str_col=["a", "b", "c"]))
>>> t
┏━━━━━━━━━┳━━━━━━━━━┓
┃ int_col ┃ str_col ┃
┡━━━━━━━━━╇━━━━━━━━━┩
│ int64 │ string │
├─────────┼─────────┤
│ 1 │ a │
│ 2 │ b │
│ 3 │ c │
└─────────┴─────────┘
>>> @ibis.udf.scalar.python
... def str_magic(x: str) -> str:
... return f"{x}_magic"
>>> @ibis.udf.scalar.python
... def add_one(x: int) -> int:
... def add_one_py(x: int) -> int:
... return x + 1
>>> expr = add_one(2)
>>> con = ibis.connect("duckdb://")
>>> con.execute(expr)
3
>>> str_magic(t.str_col)
┏━━━━━━━━━━━━━━━━━━━━━━┓
┃ str_magic_0(str_col) ┃
┡━━━━━━━━━━━━━━━━━━━━━━┩
│ string │
├──────────────────────┤
│ a_magic │
│ b_magic │
│ c_magic │
└──────────────────────┘
>>> add_one_py(t.int_col)
┏━━━━━━━━━━━━━━━━━━━━━━━┓
┃ add_one_py_0(int_col) ┃
┡━━━━━━━━━━━━━━━━━━━━━━━┩
│ int64 │
├───────────────────────┤
│ 2 │
│ 3 │
│ 4 │
└───────────────────────┘
See Also
--------
Expand Down Expand Up @@ -366,13 +397,32 @@ def pandas(
--------
```python
>>> import ibis
>>> ibis.options.interactive = True
>>> t = ibis.memtable(dict(int_col=[1, 2, 3], str_col=["a", "b", "c"]))
>>> t
┏━━━━━━━━━┳━━━━━━━━━┓
┃ int_col ┃ str_col ┃
┡━━━━━━━━━╇━━━━━━━━━┩
│ int64 │ string │
├─────────┼─────────┤
│ 1 │ a │
│ 2 │ b │
│ 3 │ c │
└─────────┴─────────┘
>>> @ibis.udf.scalar.pandas
... def add_one(x: int) -> int:
... return x + 1
>>> expr = add_one(2)
>>> con = ibis.connect(os.environ["SNOWFLAKE_URL"]) # doctest: +SKIP
>>> con.execute(expr) # doctest: +SKIP
3
... def str_cap(x: str) -> str:
... # note usage of pandas `str` method
... return x.str.capitalize()
>>> str_cap(t.str_col) # doctest: +SKIP
┏━━━━━━━━━━━━━━━━━━━━━━━┓
┃ string_cap_0(str_col) ┃
┡━━━━━━━━━━━━━━━━━━━━━━━┩
│ string │
├───────────────────────┤
│ A │
│ B │
│ C │
└───────────────────────┘
```
See Also
Expand Down Expand Up @@ -437,13 +487,22 @@ def pyarrow(
--------
>>> import ibis
>>> import pyarrow.compute as pc
>>> from datetime import date
>>> ibis.options.interactive = True
>>> t = ibis.memtable(
... dict(start_col=[date(2024, 4, 29)], end_col=[date(2025, 4, 29)]),
... )
>>> @ibis.udf.scalar.pyarrow
... def add_one(x: int) -> int:
... return pc.add(x, 1)
>>> expr = add_one(2)
>>> con = ibis.connect("duckdb://")
>>> con.execute(expr)
3
... def weeks_between(start: date, end: date) -> int:
... return pc.weeks_between(start, end)
>>> weeks_between(t.start_col, t.end_col)
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ weeks_between_0(start_col, end_col) ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│ int64 │
├─────────────────────────────────────┤
│ 52 │
└─────────────────────────────────────┘
See Also
--------
Expand Down

0 comments on commit fead40b

Please sign in to comment.