Skip to content

Commit

Permalink
modernize expression_lib
Browse files Browse the repository at this point in the history
  • Loading branch information
MarcoGorelli committed Oct 2, 2024
1 parent 7a353c3 commit ddd8706
Show file tree
Hide file tree
Showing 7 changed files with 59 additions and 87 deletions.
20 changes: 12 additions & 8 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -56,20 +56,24 @@ fn pig_latinnify(inputs: &[Series], kwargs: PigLatinKwargs) -> PolarsResult<Seri
This can then be exposed on the Python side:

```python
from __future__ import annotations

from typing import TYPE_CHECKING

import polars as pl
from polars.type_aliases import IntoExpr
from polars.utils.udfs import _get_shared_lib_location
from polars.plugins import register_plugin_function

from expression_lib.utils import parse_into_expr
from expression_lib._utils import LIB

lib = _get_shared_lib_location(__file__)
if TYPE_CHECKING:
from expression_lib._typing import IntoExpr


def pig_latinnify(expr: IntoExpr, capitalize: bool = False) -> pl.Expr:
expr = parse_into_expr(expr)
return expr.register_plugin(
lib=lib,
symbol="pig_latinnify",
return register_plugin_function(
plugin_path=LIB,
args=[expr],
function_name="pig_latinnify",
is_elementwise=True,
kwargs={"capitalize": capitalize},
)
Expand Down
Original file line number Diff line number Diff line change
@@ -1,15 +1,19 @@
from __future__ import annotations

from typing import TYPE_CHECKING

import polars as pl
from polars.type_aliases import IntoExpr
from polars.plugins import register_plugin_function
from pathlib import Path

from expression_lib.utils import parse_into_expr
from expression_lib._utils import LIB

if TYPE_CHECKING:
from expression_lib._typing import IntoExpr


def is_leap_year(expr: IntoExpr) -> pl.Expr:
expr = parse_into_expr(expr)
return register_plugin_function(
plugin_path=Path(__file__).parent,
plugin_path=LIB,
args=[expr],
function_name="is_leap_year",
is_elementwise=True,
Expand All @@ -19,9 +23,10 @@ def is_leap_year(expr: IntoExpr) -> pl.Expr:
# Note that this already exists in Polars. It is just for explanatory
# purposes.
def change_time_zone(expr: IntoExpr, tz: str = "Europe/Amsterdam") -> pl.Expr:
expr = parse_into_expr(expr)
return register_plugin_function(
plugin_path=Path(__file__).parent,
plugin_path=LIB,
args=[expr],
function_name="change_time_zone", is_elementwise=True, kwargs={"tz": tz}
function_name="change_time_zone",
is_elementwise=True,
kwargs={"tz": tz},
)
20 changes: 11 additions & 9 deletions example/derive_expression/expression_lib/expression_lib/dist.py
Original file line number Diff line number Diff line change
@@ -1,25 +1,28 @@
from __future__ import annotations

from typing import TYPE_CHECKING

import polars as pl
from polars.type_aliases import IntoExpr
from polars.plugins import register_plugin_function
from pathlib import Path

from expression_lib._utils import LIB

if TYPE_CHECKING:
from expression_lib._typing import IntoExpr

from expression_lib.utils import parse_into_expr

def hamming_distance(expr: IntoExpr, other: IntoExpr) -> pl.Expr:
expr = parse_into_expr(expr)
return register_plugin_function(
plugin_path=Path(__file__).parent,
plugin_path=LIB,
args=[expr, other],
function_name="hamming_distance",
is_elementwise=True,
)


def jaccard_similarity(expr: IntoExpr, other: IntoExpr) -> pl.Expr:
expr = parse_into_expr(expr)
return register_plugin_function(
plugin_path=Path(__file__).parent,
plugin_path=LIB,
args=[expr, other],
function_name="jaccard_similarity",
is_elementwise=True,
Expand All @@ -32,9 +35,8 @@ def haversine(
end_lat: IntoExpr,
end_long: IntoExpr,
) -> pl.Expr:
start_lat = parse_into_expr(start_lat)
return register_plugin_function(
plugin_path=Path(__file__).parent,
plugin_path=LIB,
args=[start_lat, start_long, end_lat, end_long],
function_name="haversine",
is_elementwise=True,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,11 +16,14 @@
- static typing will not recognise your custom namespace. Errors such
as `"Expr" has no attribute "dist" [attr-defined]`.
"""

from __future__ import annotations

import polars as pl
from typing import Any, Callable
from expression_lib import date_util, dist, language, utils, panic

import polars as pl

from expression_lib import date_util, dist, language, panic


@pl.api.register_expr_namespace("language")
Expand Down
Original file line number Diff line number Diff line change
@@ -1,16 +1,19 @@
from __future__ import annotations

from typing import TYPE_CHECKING

import polars as pl
from polars.type_aliases import IntoExpr
from polars.plugins import register_plugin_function
from pathlib import Path

from expression_lib.utils import parse_into_expr
from expression_lib._utils import LIB

if TYPE_CHECKING:
from expression_lib._typing import IntoExpr


def pig_latinnify(expr: IntoExpr, capitalize: bool = False) -> pl.Expr:
expr = parse_into_expr(expr)
return register_plugin_function(
plugin_path=Path(__file__).parent,
plugin_path=LIB,
args=[expr],
function_name="pig_latinnify",
is_elementwise=True,
Expand All @@ -28,9 +31,8 @@ def append_args(
"""
This example shows how arguments other than `Series` can be used.
"""
expr = parse_into_expr(expr)
return register_plugin_function(
plugin_path=Path(__file__).parent,
plugin_path=LIB,
args=[expr],
kwargs={
"float_arg": float_arg,
Expand Down
14 changes: 9 additions & 5 deletions example/derive_expression/expression_lib/expression_lib/panic.py
Original file line number Diff line number Diff line change
@@ -1,15 +1,19 @@
from __future__ import annotations

from typing import TYPE_CHECKING

import polars as pl
from polars.type_aliases import IntoExpr
from polars.plugins import register_plugin_function

from expression_lib.utils import parse_into_expr
from pathlib import Path
from expression_lib._utils import LIB

if TYPE_CHECKING:
from expression_lib._typing import IntoExpr


def panic(expr: IntoExpr) -> pl.Expr:
expr = parse_into_expr(expr)
return register_plugin_function(
plugin_path=Path(__file__).parent,
plugin_path=LIB,
args=[expr],
function_name="panic",
)
48 changes: 0 additions & 48 deletions example/derive_expression/expression_lib/expression_lib/utils.py

This file was deleted.

0 comments on commit ddd8706

Please sign in to comment.