-
Notifications
You must be signed in to change notification settings - Fork 120
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
5df15be
commit e6d5c67
Showing
5 changed files
with
100 additions
and
3 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,92 @@ | ||
from __future__ import annotations | ||
|
||
from narwhals.pandas_like.utils import evaluate_into_exprs | ||
from narwhals.translate import get_pandas | ||
from narwhals.translate import get_polars | ||
|
||
|
||
def extract_native(obj: Any, implementation) -> Any: | ||
from narwhals.expression import NarwhalsExpr | ||
|
||
# if isinstance(obj, NarwhalsExpr): | ||
# return obj._call(pl.col) | ||
if isinstance(obj, NarwhalsExpr): | ||
if implementation == "polars": | ||
import polars as pl | ||
|
||
return obj._call(pl.col) | ||
# if isinstance(obj, DType): | ||
# return obj._dtype | ||
if isinstance(obj, NarwhalsFrame): | ||
return obj._dataframe | ||
# if isinstance(obj, PolarsSeries): | ||
# return obj._series | ||
return obj | ||
|
||
|
||
class NarwhalsFrame: | ||
def __init__( | ||
self, df, *, is_eager=False, is_lazy=False, implementation: str | None = None | ||
): | ||
self._is_eager = is_eager | ||
self._is_lazy = is_lazy | ||
if implementation is not None: | ||
self._dataframe = df | ||
self._implementation = implementation | ||
return | ||
if (pl := get_polars()) is not None: | ||
if isinstance(df, pl.DataFrame): | ||
if is_lazy: | ||
raise ValueError( | ||
"can't instantiate with `is_lazy` if you pass a polars DataFrame" | ||
) | ||
self._dataframe = df | ||
self._implementation = "polars" | ||
return | ||
elif isinstance(df, pl.LazyFrame): | ||
if is_eager: | ||
raise ValueError( | ||
"can't instantiate with `is_eager` if you pass a polars LazyFrame" | ||
) | ||
self._dataframe = df | ||
self._implementation = "polars" | ||
return | ||
if (pd := get_pandas()) is not None and isinstance(df, pd.DataFrame): | ||
self._dataframe = df | ||
self._implementation = "pandas" | ||
return | ||
raise TypeError( | ||
f"Expected pandas or Polars dataframe or lazyframe, got: {type(df)}" | ||
) | ||
|
||
def _from_dataframe(self, df: Any) -> Self: | ||
# construct, preserving properties | ||
return self.__class__( | ||
df, | ||
is_eager=self._is_eager, | ||
is_lazy=self._is_lazy, | ||
implementation=self._implementation, | ||
) | ||
|
||
def _extract_native(self, obj): | ||
return extract_native(obj, implementation=self._implementation) | ||
|
||
def with_columns( | ||
self, *exprs: IntoExpr | Iterable[IntoExpr], **named_exprs: IntoExpr | ||
) -> Self: | ||
if self._implementation == "polars": | ||
return self._from_dataframe( | ||
self._dataframe.with_columns( | ||
*[self._extract_native(v) for v in exprs], | ||
**{ | ||
key: self._extract_native(value) | ||
for key, value in named_exprs.items() | ||
}, | ||
) | ||
) | ||
elif self._implementation == "pandas": | ||
new_series = evaluate_into_exprs(self, *exprs, **named_exprs) | ||
df = self._dataframe.assign( | ||
**{series.name: series._series for series in new_series} | ||
) | ||
return self._from_dataframe(df) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters