Skip to content

Commit

Permalink
simplify!
Browse files Browse the repository at this point in the history
  • Loading branch information
MarcoGorelli committed Feb 24, 2024
1 parent 24d6d55 commit f868abe
Show file tree
Hide file tree
Showing 10 changed files with 84 additions and 74 deletions.
4 changes: 2 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -51,8 +51,8 @@ def my_agnostic_function(
suppliers_native: AnyDataFrame,
parts_native: AnyDataFrame,
) -> AnyDataFrame:
suppliers, pl = to_polars_api(suppliers_native, version="0.20")
parts, _ = to_polars_api(parts_native, version="0.20")
suppliers, pl = to_polars_api(suppliers_native, lazy_only=True)
parts, _ = to_polars_api(parts_native, lazy_only=True)
result = (
suppliers.join(parts, left_on="city", right_on="city")
.filter(
Expand Down
16 changes: 5 additions & 11 deletions narwhals/pandas_like/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,20 +32,19 @@ def __init__(
self,
dataframe: Any,
*,
api_version: str,
implementation: str,
) -> None:
self._validate_columns(dataframe.columns)
self._dataframe = dataframe.reset_index(drop=True)
self._api_version = api_version
self._implementation = implementation

def _dispatch_to_lazy(self, method: str, *args: Any, **kwargs: Any) -> Self:
return getattr(self.lazy(), method)(*args, **kwargs).collect() # type: ignore[no-any-return]

def __repr__(self) -> str: # pragma: no cover
header = f" Narwhals DataFrame (api_version={self._api_version}) "
header = " Narwhals DataFrame"
length = len(header)
# TODO: use `to_original_object` instead of `._dataframe`
return (
"┌"
+ "─" * length
Expand Down Expand Up @@ -125,7 +124,6 @@ def sort(
def lazy(self) -> LazyFrame:
return LazyFrame(
self._dataframe,
api_version=self._api_version,
implementation=self._implementation,
)

Expand All @@ -151,7 +149,7 @@ def to_dict(self, *, as_series: bool = True) -> dict[str, Any]:
def group_by(self, *keys: str | Iterable[str]) -> GroupBy:
from narwhals.pandas_like.group_by import GroupBy

return GroupBy(self, flatten_str(*keys), api_version=self._api_version)
return GroupBy(self, flatten_str(*keys))

def join(
self,
Expand Down Expand Up @@ -179,16 +177,14 @@ def __init__(
self,
dataframe: Any,
*,
api_version: str,
implementation: str,
) -> None:
self._validate_columns(dataframe.columns)
self._dataframe = dataframe.reset_index(drop=True)
self._api_version = api_version
self._implementation = implementation

def __repr__(self) -> str: # pragma: no cover
header = f" Narwhals DataFrame (api_version={self._api_version}) "
header = " Narwhals DataFrame"
length = len(header)
return (
"┌"
Expand Down Expand Up @@ -222,7 +218,6 @@ def _validate_booleanness(self) -> None:
def _from_dataframe(self, df: Any) -> Self:
return self.__class__(
df,
api_version=self._api_version,
implementation=self._implementation,
)

Expand Down Expand Up @@ -298,15 +293,14 @@ def sort(
def collect(self) -> DataFrame:
return DataFrame(
self._dataframe,
api_version=self._api_version,
implementation=self._implementation,
)

# --- actions ---
def group_by(self, *keys: str | Iterable[str]) -> LazyGroupBy:
from narwhals.pandas_like.group_by import LazyGroupBy

return LazyGroupBy(self, flatten_str(*keys), api_version=self._api_version)
return LazyGroupBy(self, flatten_str(*keys))

def join(
self,
Expand Down
4 changes: 0 additions & 4 deletions narwhals/pandas_like/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,6 @@ def __init__( # noqa: PLR0913
implementation: str,
) -> None:
self._call = call
self._api_version = "0.20.0" # todo
self._depth = depth
self._function_name = function_name
self._root_names = root_names
Expand All @@ -54,7 +53,6 @@ def from_column_names(
lambda df: [
Series(
df._dataframe.loc[:, column_name],
api_version=df._api_version,
implementation=implementation,
)
for column_name in column_names
Expand Down Expand Up @@ -217,7 +215,6 @@ def ends_with(self, suffix: str) -> Expr:
lambda df: [
Series(
series.series.str.endswith(suffix),
api_version=df._api_version,
implementation=df._implementation,
)
for series in self._expr._call(df)
Expand All @@ -234,7 +231,6 @@ def strip_chars(self, characters: str = " ") -> Expr:
lambda df: [
Series(
series.series.str.strip(characters),
api_version=df._api_version,
implementation=df._implementation,
)
for series in self._expr._call(df)
Expand Down
16 changes: 4 additions & 12 deletions narwhals/pandas_like/group_by.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,28 +28,22 @@


class GroupBy(GroupByProtocol):
def __init__(self, df: DataFrame, keys: list[str], api_version: str) -> None:
def __init__(self, df: DataFrame, keys: list[str]) -> None:
self._df = df
self._keys = list(keys)
self._api_version = api_version

def agg(
self,
*aggs: IntoExpr | Iterable[IntoExpr],
**named_aggs: IntoExpr,
) -> DataFrame:
return (
LazyGroupBy(self._df.lazy(), self._keys, self._api_version)
.agg(*aggs, **named_aggs)
.collect()
)
return LazyGroupBy(self._df.lazy(), self._keys).agg(*aggs, **named_aggs).collect()


class LazyGroupBy(LazyGroupByProtocol):
def __init__(self, df: LazyFrame, keys: list[str], api_version: str) -> None:
def __init__(self, df: LazyFrame, keys: list[str]) -> None:
self._df = df
self._keys = list(keys)
self._api_version = api_version

def agg(
self,
Expand Down Expand Up @@ -99,9 +93,7 @@ def agg(
def _from_dataframe(self, df: DataFrame) -> LazyFrame:
from narwhals.pandas_like.dataframe import LazyFrame

return LazyFrame(
df, api_version=self._api_version, implementation=self._df._implementation
)
return LazyFrame(df, implementation=self._df._implementation)


def agg_pandas(
Expand Down
9 changes: 1 addition & 8 deletions narwhals/pandas_like/namespace.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,9 +38,7 @@ class Namespace(NamespaceProtocol):
String = dtypes.String

# --- not in spec ---
def __init__(self, *, api_version: str, implementation: str) -> None:
self.__dataframeapi_version__ = api_version
self.api_version = api_version
def __init__(self, *, implementation: str) -> None:
self._implementation = implementation

def _create_expr_from_callable( # noqa: PLR0913
Expand Down Expand Up @@ -69,7 +67,6 @@ def _create_series_from_scalar(self, value: Any, series: Series) -> Series:
index=series.series.index[0:1],
implementation=self._implementation,
),
api_version=self.api_version,
implementation=self._implementation,
)

Expand All @@ -94,7 +91,6 @@ def all(self) -> Expr:
lambda df: [
Series(
df._dataframe.loc[:, column_name],
api_version=df._api_version,
implementation=self._implementation,
)
for column_name in df.columns
Expand Down Expand Up @@ -137,7 +133,6 @@ def len(self) -> Expr:
index=[0],
implementation=self._implementation,
),
api_version=df._api_version,
implementation=self._implementation,
),
],
Expand Down Expand Up @@ -173,11 +168,9 @@ def concat(self, items: Iterable[AnyDataFrame], *, how: str) -> AnyDataFrame: #
if kind[0] is DataFrame:
return DataFrame( # type: ignore[return-value]
horizontal_concat(dfs, implementation=self._implementation),
api_version=self.api_version,
implementation=self._implementation,
)
return LazyFrame( # type: ignore[return-value]
horizontal_concat(dfs, implementation=self._implementation),
api_version=self.api_version,
implementation=self._implementation,
)
5 changes: 1 addition & 4 deletions narwhals/pandas_like/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@ def __init__(
self,
series: Any,
*,
api_version: str,
implementation: str,
) -> None:
"""Parameters
Expand All @@ -33,11 +32,10 @@ def __init__(
self._name = series.name
assert self._name is not None
self._series = series.reset_index(drop=True)
self.api_version = api_version
self._implementation = implementation

def __repr__(self) -> str: # pragma: no cover
header = f" Narwhals Series (api_version={self.api_version}) "
header = " Narwhals Series"
length = len(header)
return (
"┌"
Expand All @@ -53,7 +51,6 @@ def __repr__(self) -> str: # pragma: no cover
def _from_series(self, series: Any) -> Self:
return self.__class__(
series.rename(series.name, copy=False),
api_version=self.api_version,
implementation=self._implementation,
)

Expand Down
33 changes: 21 additions & 12 deletions narwhals/pandas_like/translate.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,20 +15,20 @@
def translate(
df: Any,
implementation: str,
api_version: str,
*,
eager: Literal[True],
) -> tuple[DataFrame, Namespace]:
eager_only: Literal[False],
lazy_only: Literal[False],
) -> tuple[LazyFrame, Namespace]:
...


@overload
def translate(
df: Any,
implementation: str,
api_version: str,
*,
eager: Literal[False],
eager_only: Literal[False],
lazy_only: Literal[True],
) -> tuple[LazyFrame, Namespace]:
...

Expand All @@ -37,34 +37,43 @@ def translate(
def translate(
df: Any,
implementation: str,
api_version: str,
*,
eager: bool,
eager_only: Literal[True],
lazy_only: Literal[False],
) -> tuple[DataFrame, Namespace]:
...


@overload
def translate(
df: Any,
implementation: str,
*,
eager_only: bool,
lazy_only: bool,
) -> tuple[DataFrame | LazyFrame, Namespace]:
...


def translate(
df: Any,
implementation: str,
api_version: str,
*,
eager: bool,
eager_only: bool,
lazy_only: bool,
) -> tuple[LazyFrame | DataFrame, Namespace]:
from narwhals.pandas_like.dataframe import DataFrame
from narwhals.pandas_like.dataframe import LazyFrame
from narwhals.pandas_like.utils import get_namespace

if eager:
if eager_only and not lazy_only:
df = DataFrame(
df,
api_version=api_version,
implementation=implementation,
)
else:
df = LazyFrame(
df,
api_version=api_version,
implementation=implementation,
)
return df, get_namespace(df)
2 changes: 1 addition & 1 deletion narwhals/pandas_like/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,7 @@ def maybe_evaluate_expr(df: DataFrame | LazyFrame, arg: Any) -> Any:
def get_namespace(obj: Any) -> Namespace:
from narwhals.pandas_like.namespace import Namespace

return Namespace(api_version="0.20.0", implementation=obj._implementation)
return Namespace(implementation=obj._implementation)


def parse_into_exprs(
Expand Down
Loading

0 comments on commit f868abe

Please sign in to comment.