chatExcel

from __future__ import annotations from typing import ( TYPE_CHECKING, Any, Iterator, Literal, Mapping, Sequence, Sized, cast, overload, ) import polars as pl from narwhals._polars.namespace import PolarsNamespace from narwhals._polars.series import PolarsSeries from narwhals._polars.utils import ( catch_polars_exception, extract_args_kwargs, native_to_narwhals_dtype, ) from narwhals._utils import ( Implementation, _into_arrow_table, check_columns_exist, convert_str_slice_to_int_slice, is_compliant_series, is_index_selector, is_range, is_sequence_like, is_slice_index, is_slice_none, parse_columns_to_drop, parse_version, requires, validate_backend_version, ) from narwhals.dependencies import is_numpy_array_1d from narwhals.exceptions import ColumnNotFoundError if TYPE_CHECKING: from types import ModuleType from typing import Callable, TypeVar import pandas as pd import pyarrow as pa from typing_extensions import Self, TypeAlias, TypeIs from narwhals._compliant.typing import CompliantDataFrameAny, CompliantLazyFrameAny from narwhals._polars.expr import PolarsExpr from narwhals._polars.group_by import PolarsGroupBy, PolarsLazyGroupBy from narwhals._translate import IntoArrowTable from narwhals._utils import Version, _FullContext from narwhals.dataframe import DataFrame, LazyFrame from narwhals.dtypes import DType from narwhals.schema import Schema from narwhals.typing import ( JoinStrategy, MultiColSelector, MultiIndexSelector, PivotAgg, SingleIndexSelector, _2DArray, ) T = TypeVar("T") R = TypeVar("R") Method: TypeAlias = "Callable[..., R]" """Generic alias representing all methods implemented via `__getattr__`. Where `R` is the return type. """ # DataFrame methods where PolarsDataFrame just defers to Polars.DataFrame directly. INHERITED_METHODS = frozenset( [ "clone", "drop_nulls", "estimated_size", "explode", "filter", "gather_every", "head", "is_unique", "item", "iter_rows", "join_asof", "rename", "row", "rows", "sample", "select", "sort", "tail", "to_arrow", "to_pandas", "unique", "with_columns", "write_csv", "write_parquet", ] ) class PolarsDataFrame: clone: Method[Self] collect: Method[CompliantDataFrameAny] drop_nulls: Method[Self] estimated_size: Method[int | float] explode: Method[Self] filter: Method[Self] gather_every: Method[Self] item: Method[Any] iter_rows: Method[Iterator[tuple[Any, ...]] | Iterator[Mapping[str, Any]]] is_unique: Method[PolarsSeries] join_asof: Method[Self] rename: Method[Self] row: Method[tuple[Any, ...]] rows: Method[Sequence[tuple[Any, ...]] | Sequence[Mapping[str, Any]]] sample: Method[Self] select: Method[Self] sort: Method[Self] to_arrow: Method[pa.Table] to_pandas: Method[pd.DataFrame] unique: Method[Self] with_columns: Method[Self] # NOTE: `write_csv` requires an `@overload` for `str | None` # Can't do that here 😟 write_csv: Method[Any] write_parquet: Method[None] # CompliantDataFrame _evaluate_aliases: Any def __init__( self, df: pl.DataFrame, *, backend_version: tuple[int, ...], version: Version ) -> None: self._native_frame = df self._backend_version = backend_version self._implementation = Implementation.POLARS self._version = version validate_backend_version(self._implementation, self._backend_version) @classmethod def from_arrow(cls, data: IntoArrowTable, /, *, context: _FullContext) -> Self: if context._backend_version >= (1, 3): native = pl.DataFrame(data) else: native = cast("pl.DataFrame", pl.from_arrow(_into_arrow_table(data, context))) return cls.from_native(native, context=context) @classmethod def from_dict( cls, data: Mapping[str, Any], /, *, context: _FullContext, schema: Mapping[str, DType] | Schema | None, ) -> Self: from narwhals.schema import Schema pl_schema = Schema(schema).to_polars() if schema is not None else schema return cls.from_native(pl.from_dict(data, pl_schema), context=context) @staticmethod def _is_native(obj: pl.DataFrame | Any) -> TypeIs[pl.DataFrame]: return isinstance(obj, pl.DataFrame) @classmethod def from_native(cls, data: pl.DataFrame, /, *, context: _FullContext) -> Self: return cls( data, backend_version=context._backend_version, version=context._version ) @classmethod def from_numpy( cls, data: _2DArray, /, *, context: _FullContext, # NOTE: Maybe only `Implementation`? schema: Mapping[str, DType] | Schema | Sequence[str] | None, ) -> Self: from narwhals.schema import Schema pl_schema = ( Schema(schema).to_polars() if isinstance(schema, (Mapping, Schema)) else schema ) return cls.from_native(pl.from_numpy(data, pl_schema), context=context) def to_narwhals(self) -> DataFrame[pl.DataFrame]: return self._version.dataframe(self, level="full") @property def native(self) -> pl.DataFrame: return self._native_frame def __repr__(self) -> str: # pragma: no cover return "PolarsDataFrame" def __narwhals_dataframe__(self) -> Self: return self def __narwhals_namespace__(self) -> PolarsNamespace: return PolarsNamespace( backend_version=self._backend_version, version=self._version ) def __native_namespace__(self) -> ModuleType: if self._implementation is Implementation.POLARS: return self._implementation.to_native_namespace() msg = f"Expected polars, got: {type(self._implementation)}" # pragma: no cover raise AssertionError(msg) def _with_version(self, version: Version) -> Self: return self.__class__( self.native, backend_version=self._backend_version, version=version ) def _with_native(self, df: pl.DataFrame) -> Self: return self.__class__( df, backend_version=self._backend_version, version=self._version ) @overload def _from_native_object(self, obj: pl.Series) -> PolarsSeries: ... @overload def _from_native_object(self, obj: pl.DataFrame) -> Self: ... @overload def _from_native_object(self, obj: T) -> T: ... def _from_native_object( self, obj: pl.Series | pl.DataFrame | T ) -> Self | PolarsSeries | T: if isinstance(obj, pl.Series): return PolarsSeries.from_native(obj, context=self) if self._is_native(obj): return self._with_native(obj) # scalar return obj def __len__(self) -> int: return len(self.native) def head(self, n: int) -> Self: return self._with_native(self.native.head(n)) def tail(self, n: int) -> Self: return self._with_native(self.native.tail(n)) def __getattr__(self, attr: str) -> Any: if attr not in INHERITED_METHODS: # pragma: no cover msg = f"{self.__class__.__name__} has not attribute '{attr}'." raise AttributeError(msg) def func(*args: Any, **kwargs: Any) -> Any: pos, kwds = extract_args_kwargs(args, kwargs) try: return self._from_native_object(getattr(self.native, attr)(*pos, **kwds)) except pl.exceptions.ColumnNotFoundError as e: # pragma: no cover msg = f"{e!s}\n\nHint: Did you mean one of these columns: {self.columns}?" raise ColumnNotFoundError(msg) from e except Exception as e: # noqa: BLE001 raise catch_polars_exception(e, self._backend_version) from None return func def __array__( self, dtype: Any | None = None, *, copy: bool | None = None ) -> _2DArray: if self._backend_version < (0, 20, 28) and copy is not None: msg = "`copy` in `__array__` is only supported for 'polars>=0.20.28'" raise NotImplementedError(msg) if self._backend_version < (0, 20, 28): return self.native.__array__(dtype) return self.native.__array__(dtype) def to_numpy(self, dtype: Any = None, *, copy: bool | None = None) -> _2DArray: return self.native.to_numpy() def collect_schema(self) -> dict[str, DType]: if self._backend_version < (1,): return { name: native_to_narwhals_dtype( dtype, self._version, self._backend_version ) for name, dtype in self.native.schema.items() } else: collected_schema = self.native.collect_schema() return { name: native_to_narwhals_dtype( dtype, self._version, self._backend_version ) for name, dtype in collected_schema.items() } @property def shape(self) -> tuple[int, int]: return self.native.shape def __getitem__( # noqa: C901, PLR0912 self, item: tuple[ SingleIndexSelector | MultiIndexSelector[PolarsSeries], MultiColSelector[PolarsSeries], ], ) -> Any: rows, columns = item if self._backend_version > (0, 20, 30): rows_native = rows.native if is_compliant_series(rows) else rows columns_native = columns.native if is_compliant_series(columns) else columns selector = rows_native, columns_native selected = self.native.__getitem__(selector) # type: ignore[index] return self._from_native_object(selected) else: # pragma: no cover # TODO(marco): we can delete this branch after Polars==0.20.30 becomes the minimum # Polars version we support # This mostly mirrors the logic in `EagerDataFrame.__getitem__`. rows = list(rows) if isinstance(rows, tuple) else rows columns = list(columns) if isinstance(columns, tuple) else columns if is_numpy_array_1d(columns): columns = columns.tolist() native = self.native if not is_slice_none(columns): if isinstance(columns, Sized) and len(columns) == 0: return self.select() if is_index_selector(columns): if is_slice_index(columns) or is_range(columns): native = native.select( self.columns[slice(columns.start, columns.stop, columns.step)] ) elif is_compliant_series(columns): native = native[:, columns.native.to_list()] else: native = native[:, columns] elif isinstance(columns, slice): native = native.select( self.columns[ slice(*convert_str_slice_to_int_slice(columns, self.columns)) ] ) elif is_compliant_series(columns): native = native.select(columns.native.to_list()) elif is_sequence_like(columns): native = native.select(columns) else: msg = f"Unreachable code, got unexpected type: {type(columns)}" raise AssertionError(msg) if not is_slice_none(rows): if isinstance(rows, int): native = native[[rows], :] elif isinstance(rows, (slice, range)): native = native[rows, :] elif is_compliant_series(rows): native = native[rows.native, :] elif is_sequence_like(rows): native = native[rows, :] else: msg = f"Unreachable code, got unexpected type: {type(rows)}" raise AssertionError(msg) return self._with_native(native) def simple_select(self, *column_names: str) -> Self: return self._with_native(self.native.select(*column_names)) def aggregate(self, *exprs: Any) -> Self: return self.select(*exprs) def get_column(self, name: str) -> PolarsSeries: return PolarsSeries.from_native(self.native.get_column(name), context=self) def iter_columns(self) -> Iterator[PolarsSeries]: for series in self.native.iter_columns(): yield PolarsSeries.from_native(series, context=self) @property def columns(self) -> list[str]: return self.native.columns @property def schema(self) -> dict[str, DType]: return { name: native_to_narwhals_dtype(dtype, self._version, self._backend_version) for name, dtype in self.native.schema.items() } def lazy(self, *, backend: Implementation | None = None) -> CompliantLazyFrameAny: if backend is None or backend is Implementation.POLARS: return PolarsLazyFrame.from_native(self.native.lazy(), context=self) elif backend is Implementation.DUCKDB: import duckdb # ignore-banned-import from narwhals._duckdb.dataframe import DuckDBLazyFrame # NOTE: (F841) is a false positive df = self.native # noqa: F841 return DuckDBLazyFrame( duckdb.table("df"), backend_version=parse_version(duckdb), version=self._version, ) elif backend is Implementation.DASK: import dask # ignore-banned-import import dask.dataframe as dd # ignore-banned-import from narwhals._dask.dataframe import DaskLazyFrame return DaskLazyFrame( dd.from_pandas(self.native.to_pandas()), backend_version=parse_version(dask), version=self._version, ) raise AssertionError # pragma: no cover @overload def to_dict(self, *, as_series: Literal[True]) -> dict[str, PolarsSeries]: ... @overload def to_dict(self, *, as_series: Literal[False]) -> dict[str, list[Any]]: ... def to_dict( self, *, as_series: bool ) -> dict[str, PolarsSeries] | dict[str, list[Any]]: if as_series: return { name: PolarsSeries.from_native(col, context=self) for name, col in self.native.to_dict().items() } else: return self.native.to_dict(as_series=False) def group_by( self, keys: Sequence[str] | Sequence[PolarsExpr], *, drop_null_keys: bool ) -> PolarsGroupBy: from narwhals._polars.group_by import PolarsGroupBy return PolarsGroupBy(self, keys, drop_null_keys=drop_null_keys) def with_row_index(self, name: str) -> Self: if self._backend_version < (0, 20, 4): return self._with_native(self.native.with_row_count(name)) return self._with_native(self.native.with_row_index(name)) def drop(self, columns: Sequence[str], *, strict: bool) -> Self: to_drop = parse_columns_to_drop(self, columns, strict=strict) return self._with_native(self.native.drop(to_drop)) def unpivot( self, on: Sequence[str] | None, index: Sequence[str] | None, variable_name: str, value_name: str, ) -> Self: if self._backend_version < (1, 0, 0): return self._with_native( self.native.melt( id_vars=index, value_vars=on, variable_name=variable_name, value_name=value_name, ) ) return self._with_native( self.native.unpivot( on=on, index=index, variable_name=variable_name, value_name=value_name ) ) @requires.backend_version((1,)) def pivot( self, on: Sequence[str], *, index: Sequence[str] | None, values: Sequence[str] | None, aggregate_function: PivotAgg | None, sort_columns: bool, separator: str, ) -> Self: try: result = self.native.pivot( on, index=index, values=values, aggregate_function=aggregate_function, sort_columns=sort_columns, separator=separator, ) except Exception as e: # noqa: BLE001 raise catch_polars_exception(e, self._backend_version) from None return self._from_native_object(result) def to_polars(self) -> pl.DataFrame: return self.native def join( self, other: Self, *, how: JoinStrategy, left_on: Sequence[str] | None, right_on: Sequence[str] | None, suffix: str, ) -> Self: how_native = ( "outer" if (self._backend_version < (0, 20, 29) and how == "full") else how ) try: return self._with_native( self.native.join( other=other.native, how=how_native, # type: ignore[arg-type] left_on=left_on, right_on=right_on, suffix=suffix, ) ) except Exception as e: # noqa: BLE001 raise catch_polars_exception(e, self._backend_version) from None def _check_columns_exist(self, subset: Sequence[str]) -> ColumnNotFoundError | None: return check_columns_exist(subset, available=self.columns) class PolarsLazyFrame: drop_nulls: Method[Self] explode: Method[Self] filter: Method[Self] gather_every: Method[Self] head: Method[Self] join_asof: Method[Self] rename: Method[Self] select: Method[Self] sort: Method[Self] tail: Method[Self] unique: Method[Self] with_columns: Method[Self] # CompliantLazyFrame _evaluate_expr: Any _evaluate_window_expr: Any _evaluate_aliases: Any def __init__( self, df: pl.LazyFrame, *, backend_version: tuple[int, ...], version: Version ) -> None: self._native_frame = df self._backend_version = backend_version self._implementation = Implementation.POLARS self._version = version validate_backend_version(self._implementation, self._backend_version) @staticmethod def _is_native(obj: pl.LazyFrame | Any) -> TypeIs[pl.LazyFrame]: return isinstance(obj, pl.LazyFrame) @classmethod def from_native(cls, data: pl.LazyFrame, /, *, context: _FullContext) -> Self: return cls( data, backend_version=context._backend_version, version=context._version ) def to_narwhals(self) -> LazyFrame[pl.LazyFrame]: return self._version.lazyframe(self, level="lazy") def __repr__(self) -> str: # pragma: no cover return "PolarsLazyFrame" def __narwhals_lazyframe__(self) -> Self: return self def __narwhals_namespace__(self) -> PolarsNamespace: return PolarsNamespace( backend_version=self._backend_version, version=self._version ) def __native_namespace__(self) -> ModuleType: if self._implementation is Implementation.POLARS: return self._implementation.to_native_namespace() msg = f"Expected polars, got: {type(self._implementation)}" # pragma: no cover raise AssertionError(msg) def _with_native(self, df: pl.LazyFrame) -> Self: return self.__class__( df, backend_version=self._backend_version, version=self._version ) def _with_version(self, version: Version) -> Self: return self.__class__( self.native, backend_version=self._backend_version, version=version ) def __getattr__(self, attr: str) -> Any: if attr not in INHERITED_METHODS: # pragma: no cover msg = f"{self.__class__.__name__} has not attribute '{attr}'." raise AttributeError(msg) def func(*args: Any, **kwargs: Any) -> Any: pos, kwds = extract_args_kwargs(args, kwargs) try: return self._with_native(getattr(self.native, attr)(*pos, **kwds)) except pl.exceptions.ColumnNotFoundError as e: # pragma: no cover raise ColumnNotFoundError(str(e)) from e return func def _iter_columns(self) -> Iterator[PolarsSeries]: # pragma: no cover yield from self.collect(self._implementation).iter_columns() @property def native(self) -> pl.LazyFrame: return self._native_frame @property def columns(self) -> list[str]: return self.native.columns @property def schema(self) -> dict[str, DType]: schema = self.native.schema return { name: native_to_narwhals_dtype(dtype, self._version, self._backend_version) for name, dtype in schema.items() } def collect_schema(self) -> dict[str, DType]: if self._backend_version < (1,): return { name: native_to_narwhals_dtype( dtype, self._version, self._backend_version ) for name, dtype in self.native.schema.items() } else: try: collected_schema = self.native.collect_schema() except Exception as e: # noqa: BLE001 raise catch_polars_exception(e, self._backend_version) from None return { name: native_to_narwhals_dtype( dtype, self._version, self._backend_version ) for name, dtype in collected_schema.items() } def collect( self, backend: Implementation | None, **kwargs: Any ) -> CompliantDataFrameAny: try: result = self.native.collect(**kwargs) except Exception as e: # noqa: BLE001 raise catch_polars_exception(e, self._backend_version) from None if backend is None or backend is Implementation.POLARS: return PolarsDataFrame.from_native(result, context=self) if backend is Implementation.PANDAS: import pandas as pd # ignore-banned-import from narwhals._pandas_like.dataframe import PandasLikeDataFrame return PandasLikeDataFrame( result.to_pandas(), implementation=Implementation.PANDAS, backend_version=parse_version(pd), version=self._version, validate_column_names=False, ) if backend is Implementation.PYARROW: import pyarrow as pa # ignore-banned-import from narwhals._arrow.dataframe import ArrowDataFrame return ArrowDataFrame( result.to_arrow(), backend_version=parse_version(pa), version=self._version, validate_column_names=False, ) msg = f"Unsupported `backend` value: {backend}" # pragma: no cover raise ValueError(msg) # pragma: no cover def group_by( self, keys: Sequence[str] | Sequence[PolarsExpr], *, drop_null_keys: bool ) -> PolarsLazyGroupBy: from narwhals._polars.group_by import PolarsLazyGroupBy return PolarsLazyGroupBy(self, keys, drop_null_keys=drop_null_keys) def with_row_index(self, name: str) -> Self: if self._backend_version < (0, 20, 4): return self._with_native(self.native.with_row_count(name)) return self._with_native(self.native.with_row_index(name)) def drop(self, columns: Sequence[str], *, strict: bool) -> Self: if self._backend_version < (1, 0, 0): return self._with_native(self.native.drop(columns)) return self._with_native(self.native.drop(columns, strict=strict)) def unpivot( self, on: Sequence[str] | None, index: Sequence[str] | None, variable_name: str, value_name: str, ) -> Self: if self._backend_version < (1, 0, 0): return self._with_native( self.native.melt( id_vars=index, value_vars=on, variable_name=variable_name, value_name=value_name, ) ) return self._with_native( self.native.unpivot( on=on, index=index, variable_name=variable_name, value_name=value_name ) ) def simple_select(self, *column_names: str) -> Self: return self._with_native(self.native.select(*column_names)) def aggregate(self, *exprs: Any) -> Self: return self.select(*exprs) def join( self, other: Self, *, how: JoinStrategy, left_on: Sequence[str] | None, right_on: Sequence[str] | None, suffix: str, ) -> Self: how_native = ( "outer" if (self._backend_version < (0, 20, 29) and how == "full") else how ) return self._with_native( self.native.join( other=other.native, how=how_native, # type: ignore[arg-type] left_on=left_on, right_on=right_on, suffix=suffix, ) ) def _check_columns_exist(self, subset: Sequence[str]) -> ColumnNotFoundError | None: return check_columns_exist( # pragma: no cover subset, available=self.columns )