Skip to content

Commit

Permalink
Merge branch 'main' into feat/improving-array-casting
Browse files Browse the repository at this point in the history
  • Loading branch information
@FBruzzesi
FBruzzesi committed Jan 27, 2025
2 parents 31dd90f + 9b6a355 commit 044db09
Show file tree
Hide file tree
Showing 5 changed files with 109 additions and 301 deletions.
101 changes: 26 additions & 75 deletions narwhals/_arrow/expr_name.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

from typing import TYPE_CHECKING
from typing import Callable
from typing import Sequence

if TYPE_CHECKING:
from typing_extensions import Self
Expand All @@ -14,107 +15,59 @@ def __init__(self: Self, expr: ArrowExpr) -> None:
self._compliant_expr = expr

def keep(self: Self) -> ArrowExpr:
return self._compliant_expr.__class__(
lambda df: [
series.alias(name)
for series, name in zip(
self._compliant_expr._call(df),
self._compliant_expr._evaluate_output_names(df),
)
],
depth=self._compliant_expr._depth,
function_name=self._compliant_expr._function_name,
evaluate_output_names=self._compliant_expr._evaluate_output_names,
return self._from_colname_func_and_alias_output_names(
name_mapping_func=lambda name: name,
alias_output_names=None,
backend_version=self._compliant_expr._backend_version,
version=self._compliant_expr._version,
kwargs=self._compliant_expr._kwargs,
)

def map(self: Self, function: Callable[[str], str]) -> ArrowExpr:
return self._compliant_expr.__class__(
lambda df: [
series.alias(function(str(name)))
for series, name in zip(
self._compliant_expr._call(df),
self._compliant_expr._evaluate_output_names(df),
)
],
depth=self._compliant_expr._depth,
function_name=self._compliant_expr._function_name,
evaluate_output_names=self._compliant_expr._evaluate_output_names,
return self._from_colname_func_and_alias_output_names(
name_mapping_func=lambda name: function(str(name)),
alias_output_names=lambda output_names: [
function(str(name)) for name in output_names
],
backend_version=self._compliant_expr._backend_version,
version=self._compliant_expr._version,
kwargs={**self._compliant_expr._kwargs, "function": function},
)

def prefix(self: Self, prefix: str) -> ArrowExpr:
return self._compliant_expr.__class__(
lambda df: [
series.alias(f"{prefix}{name}")
for series, name in zip(
self._compliant_expr._call(df),
self._compliant_expr._evaluate_output_names(df),
)
],
depth=self._compliant_expr._depth,
function_name=self._compliant_expr._function_name,
evaluate_output_names=self._compliant_expr._evaluate_output_names,
return self._from_colname_func_and_alias_output_names(
name_mapping_func=lambda name: f"{prefix}{name}",
alias_output_names=lambda output_names: [
f"{prefix}{output_name}" for output_name in output_names
],
backend_version=self._compliant_expr._backend_version,
version=self._compliant_expr._version,
kwargs={**self._compliant_expr._kwargs, "prefix": prefix},
)

def suffix(self: Self, suffix: str) -> ArrowExpr:
return self._compliant_expr.__class__(
lambda df: [
series.alias(f"{name}{suffix}")
for series, name in zip(
self._compliant_expr._call(df),
self._compliant_expr._evaluate_output_names(df),
)
],
depth=self._compliant_expr._depth,
function_name=self._compliant_expr._function_name,
evaluate_output_names=self._compliant_expr._evaluate_output_names,
return self._from_colname_func_and_alias_output_names(
name_mapping_func=lambda name: f"{name}{suffix}",
alias_output_names=lambda output_names: [
f"{output_name}{suffix}" for output_name in output_names
],
backend_version=self._compliant_expr._backend_version,
version=self._compliant_expr._version,
kwargs={**self._compliant_expr._kwargs, "suffix": suffix},
)

def to_lowercase(self: Self) -> ArrowExpr:
return self._compliant_expr.__class__(
lambda df: [
series.alias(str(name).lower())
for series, name in zip(
self._compliant_expr._call(df),
self._compliant_expr._evaluate_output_names(df),
)
],
depth=self._compliant_expr._depth,
function_name=self._compliant_expr._function_name,
evaluate_output_names=self._compliant_expr._evaluate_output_names,
return self._from_colname_func_and_alias_output_names(
name_mapping_func=lambda name: str(name).lower(),
alias_output_names=lambda output_names: [
str(name).lower() for name in output_names
],
backend_version=self._compliant_expr._backend_version,
version=self._compliant_expr._version,
kwargs=self._compliant_expr._kwargs,
)

def to_uppercase(self: Self) -> ArrowExpr:
return self._from_colname_func_and_alias_output_names(
name_mapping_func=lambda name: str(name).upper(),
alias_output_names=lambda output_names: [
str(name).upper() for name in output_names
],
)

def _from_colname_func_and_alias_output_names(
self: Self,
name_mapping_func: Callable[[str], str],
alias_output_names: Callable[[Sequence[str]], Sequence[str]] | None,
) -> ArrowExpr:
return self._compliant_expr.__class__(
lambda df: [
series.alias(str(name).upper())
call=lambda df: [
series.alias(name_mapping_func(name))
for series, name in zip(
self._compliant_expr._call(df),
self._compliant_expr._evaluate_output_names(df),
Expand All @@ -123,9 +76,7 @@ def to_uppercase(self: Self) -> ArrowExpr:
depth=self._compliant_expr._depth,
function_name=self._compliant_expr._function_name,
evaluate_output_names=self._compliant_expr._evaluate_output_names,
alias_output_names=lambda output_names: [
str(name).upper() for name in output_names
],
alias_output_names=alias_output_names,
backend_version=self._compliant_expr._backend_version,
version=self._compliant_expr._version,
kwargs=self._compliant_expr._kwargs,
Expand Down
71 changes: 19 additions & 52 deletions narwhals/_dask/expr_name.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

from typing import TYPE_CHECKING
from typing import Callable
from typing import Sequence

if TYPE_CHECKING:
from typing_extensions import Self
Expand All @@ -14,87 +15,53 @@ def __init__(self: Self, expr: DaskExpr) -> None:
self._compliant_expr = expr

def keep(self: Self) -> DaskExpr:
return self._compliant_expr.__class__(
self._compliant_expr._call,
depth=self._compliant_expr._depth,
function_name=self._compliant_expr._function_name,
evaluate_output_names=self._compliant_expr._evaluate_output_names,
alias_output_names=None,
returns_scalar=self._compliant_expr._returns_scalar,
backend_version=self._compliant_expr._backend_version,
version=self._compliant_expr._version,
kwargs=self._compliant_expr._kwargs,
)
return self._from_alias_output_names(alias_output_names=None)

def map(self: Self, function: Callable[[str], str]) -> DaskExpr:
return self._compliant_expr.__class__(
self._compliant_expr._call,
depth=self._compliant_expr._depth,
function_name=self._compliant_expr._function_name,
evaluate_output_names=self._compliant_expr._evaluate_output_names,
return self._from_alias_output_names(
alias_output_names=lambda output_names: [
function(str(name)) for name in output_names
],
returns_scalar=self._compliant_expr._returns_scalar,
backend_version=self._compliant_expr._backend_version,
version=self._compliant_expr._version,
kwargs={**self._compliant_expr._kwargs, "function": function},
)

def prefix(self: Self, prefix: str) -> DaskExpr:
return self._compliant_expr.__class__(
self._compliant_expr._call,
depth=self._compliant_expr._depth,
function_name=self._compliant_expr._function_name,
evaluate_output_names=self._compliant_expr._evaluate_output_names,
return self._from_alias_output_names(
alias_output_names=lambda output_names: [
f"{prefix}{output_name}" for output_name in output_names
],
returns_scalar=self._compliant_expr._returns_scalar,
backend_version=self._compliant_expr._backend_version,
version=self._compliant_expr._version,
kwargs={**self._compliant_expr._kwargs, "prefix": prefix},
)

def suffix(self: Self, suffix: str) -> DaskExpr:
return self._compliant_expr.__class__(
self._compliant_expr._call,
depth=self._compliant_expr._depth,
function_name=self._compliant_expr._function_name,
evaluate_output_names=self._compliant_expr._evaluate_output_names,
return self._from_alias_output_names(
alias_output_names=lambda output_names: [
f"{output_name}{suffix}" for output_name in output_names
],
returns_scalar=self._compliant_expr._returns_scalar,
backend_version=self._compliant_expr._backend_version,
version=self._compliant_expr._version,
kwargs={**self._compliant_expr._kwargs, "suffix": suffix},
]
)

def to_lowercase(self: Self) -> DaskExpr:
return self._compliant_expr.__class__(
self._compliant_expr._call,
depth=self._compliant_expr._depth,
function_name=self._compliant_expr._function_name,
evaluate_output_names=self._compliant_expr._evaluate_output_names,
return self._from_alias_output_names(
alias_output_names=lambda output_names: [
str(name).lower() for name in output_names
],
returns_scalar=self._compliant_expr._returns_scalar,
backend_version=self._compliant_expr._backend_version,
version=self._compliant_expr._version,
kwargs=self._compliant_expr._kwargs,
)

def to_uppercase(self: Self) -> DaskExpr:
return self._from_alias_output_names(
alias_output_names=lambda output_names: [
str(name).upper() for name in output_names
]
)

def _from_alias_output_names(
self: Self,
alias_output_names: Callable[[Sequence[str]], Sequence[str]] | None,
) -> DaskExpr:
return self._compliant_expr.__class__(
self._compliant_expr._call,
call=self._compliant_expr._call,
depth=self._compliant_expr._depth,
function_name=self._compliant_expr._function_name,
evaluate_output_names=self._compliant_expr._evaluate_output_names,
alias_output_names=lambda output_names: [
str(name).upper() for name in output_names
],
alias_output_names=alias_output_names,
returns_scalar=self._compliant_expr._returns_scalar,
backend_version=self._compliant_expr._backend_version,
version=self._compliant_expr._version,
Expand Down
66 changes: 19 additions & 47 deletions narwhals/_duckdb/expr_name.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

from typing import TYPE_CHECKING
from typing import Callable
from typing import Sequence

if TYPE_CHECKING:
from typing_extensions import Self
Expand All @@ -14,82 +15,53 @@ def __init__(self: Self, expr: DuckDBExpr) -> None:
self._compliant_expr = expr

def keep(self: Self) -> DuckDBExpr:
return self._compliant_expr.__class__(
self._compliant_expr._call,
depth=self._compliant_expr._depth,
function_name=self._compliant_expr._function_name,
evaluate_output_names=self._compliant_expr._evaluate_output_names,
alias_output_names=None,
returns_scalar=self._compliant_expr._returns_scalar,
backend_version=self._compliant_expr._backend_version,
version=self._compliant_expr._version,
)
return self._from_alias_output_names(alias_output_names=None)

def map(self: Self, function: Callable[[str], str]) -> DuckDBExpr:
return self._compliant_expr.__class__(
self._compliant_expr._call,
depth=self._compliant_expr._depth,
function_name=self._compliant_expr._function_name,
evaluate_output_names=self._compliant_expr._evaluate_output_names,
return self._from_alias_output_names(
alias_output_names=lambda output_names: [
function(str(name)) for name in output_names
],
returns_scalar=self._compliant_expr._returns_scalar,
backend_version=self._compliant_expr._backend_version,
version=self._compliant_expr._version,
)

def prefix(self: Self, prefix: str) -> DuckDBExpr:
return self._compliant_expr.__class__(
self._compliant_expr._call,
depth=self._compliant_expr._depth,
function_name=self._compliant_expr._function_name,
evaluate_output_names=self._compliant_expr._evaluate_output_names,
return self._from_alias_output_names(
alias_output_names=lambda output_names: [
f"{prefix}{output_name}" for output_name in output_names
],
returns_scalar=self._compliant_expr._returns_scalar,
backend_version=self._compliant_expr._backend_version,
version=self._compliant_expr._version,
)

def suffix(self: Self, suffix: str) -> DuckDBExpr:
return self._compliant_expr.__class__(
self._compliant_expr._call,
depth=self._compliant_expr._depth,
function_name=self._compliant_expr._function_name,
evaluate_output_names=self._compliant_expr._evaluate_output_names,
return self._from_alias_output_names(
alias_output_names=lambda output_names: [
f"{output_name}{suffix}" for output_name in output_names
],
returns_scalar=self._compliant_expr._returns_scalar,
backend_version=self._compliant_expr._backend_version,
version=self._compliant_expr._version,
]
)

def to_lowercase(self: Self) -> DuckDBExpr:
return self._compliant_expr.__class__(
self._compliant_expr._call,
depth=self._compliant_expr._depth,
function_name=self._compliant_expr._function_name,
evaluate_output_names=self._compliant_expr._evaluate_output_names,
return self._from_alias_output_names(
alias_output_names=lambda output_names: [
str(name).lower() for name in output_names
],
returns_scalar=self._compliant_expr._returns_scalar,
backend_version=self._compliant_expr._backend_version,
version=self._compliant_expr._version,
)

def to_uppercase(self: Self) -> DuckDBExpr:
return self._from_alias_output_names(
alias_output_names=lambda output_names: [
str(name).upper() for name in output_names
]
)

def _from_alias_output_names(
self: Self,
alias_output_names: Callable[[Sequence[str]], Sequence[str]] | None,
) -> DuckDBExpr:
return self._compliant_expr.__class__(
self._compliant_expr._call,
call=self._compliant_expr._call,
depth=self._compliant_expr._depth,
function_name=self._compliant_expr._function_name,
evaluate_output_names=self._compliant_expr._evaluate_output_names,
alias_output_names=lambda output_names: [
str(name).upper() for name in output_names
],
alias_output_names=alias_output_names,
returns_scalar=self._compliant_expr._returns_scalar,
backend_version=self._compliant_expr._backend_version,
version=self._compliant_expr._version,
Expand Down
Loading

0 comments on commit 044db09

Please sign in to comment.