chatExcel

from __future__ import annotations from typing import TYPE_CHECKING, Any, Iterator, Mapping, Sequence import dask.dataframe as dd import pandas as pd from narwhals._dask.utils import add_row_index, evaluate_exprs from narwhals._pandas_like.utils import native_to_narwhals_dtype, select_columns_by_name from narwhals._utils import ( Implementation, _remap_full_join_keys, check_column_names_are_unique, generate_temporary_column_name, not_implemented, parse_columns_to_drop, parse_version, validate_backend_version, ) from narwhals.typing import CompliantLazyFrame if TYPE_CHECKING: from types import ModuleType import dask.dataframe.dask_expr as dx from typing_extensions import Self, TypeIs from narwhals._compliant.typing import CompliantDataFrameAny from narwhals._dask.expr import DaskExpr from narwhals._dask.group_by import DaskLazyGroupBy from narwhals._dask.namespace import DaskNamespace from narwhals._utils import Version, _FullContext from narwhals.dataframe import LazyFrame from narwhals.dtypes import DType from narwhals.typing import AsofJoinStrategy, JoinStrategy, LazyUniqueKeepStrategy class DaskLazyFrame( CompliantLazyFrame["DaskExpr", "dd.DataFrame", "LazyFrame[dd.DataFrame]"] ): def __init__( self, native_dataframe: dd.DataFrame, *, backend_version: tuple[int, ...], version: Version, ) -> None: self._native_frame: dd.DataFrame = native_dataframe self._backend_version = backend_version self._implementation = Implementation.DASK self._version = version self._cached_schema: dict[str, DType] | None = None self._cached_columns: list[str] | None = None validate_backend_version(self._implementation, self._backend_version) @staticmethod def _is_native(obj: dd.DataFrame | Any) -> TypeIs[dd.DataFrame]: return isinstance(obj, dd.DataFrame) @classmethod def from_native(cls, data: dd.DataFrame, /, *, context: _FullContext) -> Self: return cls( data, backend_version=context._backend_version, version=context._version ) def to_narwhals(self) -> LazyFrame[dd.DataFrame]: return self._version.lazyframe(self, level="lazy") def __native_namespace__(self) -> ModuleType: if self._implementation is Implementation.DASK: return self._implementation.to_native_namespace() msg = f"Expected dask, got: {type(self._implementation)}" # pragma: no cover raise AssertionError(msg) def __narwhals_namespace__(self) -> DaskNamespace: from narwhals._dask.namespace import DaskNamespace return DaskNamespace(backend_version=self._backend_version, version=self._version) def __narwhals_lazyframe__(self) -> Self: return self def _with_version(self, version: Version) -> Self: return self.__class__( self.native, backend_version=self._backend_version, version=version ) def _with_native(self, df: Any) -> Self: return self.__class__( df, backend_version=self._backend_version, version=self._version ) def _iter_columns(self) -> Iterator[dx.Series]: for _col, ser in self.native.items(): # noqa: PERF102 yield ser def with_columns(self, *exprs: DaskExpr) -> Self: new_series = evaluate_exprs(self, *exprs) return self._with_native(self.native.assign(**dict(new_series))) def collect( self, backend: Implementation | None, **kwargs: Any ) -> CompliantDataFrameAny: result = self.native.compute(**kwargs) if backend is None or backend is Implementation.PANDAS: from narwhals._pandas_like.dataframe import PandasLikeDataFrame return PandasLikeDataFrame( result, implementation=Implementation.PANDAS, backend_version=parse_version(pd), version=self._version, validate_column_names=True, ) if backend is Implementation.POLARS: import polars as pl # ignore-banned-import from narwhals._polars.dataframe import PolarsDataFrame return PolarsDataFrame( pl.from_pandas(result), backend_version=parse_version(pl), version=self._version, ) if backend is Implementation.PYARROW: import pyarrow as pa # ignore-banned-import from narwhals._arrow.dataframe import ArrowDataFrame return ArrowDataFrame( pa.Table.from_pandas(result), backend_version=parse_version(pa), version=self._version, validate_column_names=True, ) msg = f"Unsupported `backend` value: {backend}" # pragma: no cover raise ValueError(msg) # pragma: no cover @property def columns(self) -> list[str]: if self._cached_columns is None: self._cached_columns = ( list(self.schema) if self._cached_schema is not None else self.native.columns.tolist() ) return self._cached_columns def filter(self, predicate: DaskExpr) -> Self: # `[0]` is safe as the predicate's expression only returns a single column mask = predicate(self)[0] return self._with_native(self.native.loc[mask]) def simple_select(self, *column_names: str) -> Self: native = select_columns_by_name( self.native, list(column_names), self._backend_version, self._implementation ) return self._with_native(native) def aggregate(self, *exprs: DaskExpr) -> Self: new_series = evaluate_exprs(self, *exprs) df = dd.concat([val.rename(name) for name, val in new_series], axis=1) return self._with_native(df) def select(self, *exprs: DaskExpr) -> Self: new_series = evaluate_exprs(self, *exprs) df = select_columns_by_name( self.native.assign(**dict(new_series)), [s[0] for s in new_series], self._backend_version, self._implementation, ) return self._with_native(df) def drop_nulls(self, subset: Sequence[str] | None) -> Self: if subset is None: return self._with_native(self.native.dropna()) plx = self.__narwhals_namespace__() return self.filter(~plx.any_horizontal(plx.col(*subset).is_null())) @property def schema(self) -> dict[str, DType]: if self._cached_schema is None: native_dtypes = self.native.dtypes self._cached_schema = { col: native_to_narwhals_dtype( native_dtypes[col], self._version, self._implementation ) for col in self.native.columns } return self._cached_schema def collect_schema(self) -> dict[str, DType]: return self.schema 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(columns=to_drop)) def with_row_index(self, name: str) -> Self: # Implementation is based on the following StackOverflow reply: # https://stackoverflow.com/questions/60831518/in-dask-how-does-one-add-a-range-of-integersauto-increment-to-a-new-column/60852409#60852409 return self._with_native( add_row_index(self.native, name, self._backend_version, self._implementation) ) def rename(self, mapping: Mapping[str, str]) -> Self: return self._with_native(self.native.rename(columns=mapping)) def head(self, n: int) -> Self: return self._with_native(self.native.head(n=n, compute=False, npartitions=-1)) def unique( self, subset: Sequence[str] | None, *, keep: LazyUniqueKeepStrategy ) -> Self: if subset and (error := self._check_columns_exist(subset)): raise error if keep == "none": subset = subset or self.columns token = generate_temporary_column_name(n_bytes=8, columns=subset) ser = self.native.groupby(subset).size().rename(token) ser = ser[ser == 1] unique = ser.reset_index().drop(columns=token) result = self.native.merge(unique, on=subset, how="inner") else: mapped_keep = {"any": "first"}.get(keep, keep) result = self.native.drop_duplicates(subset=subset, keep=mapped_keep) return self._with_native(result) def sort(self, *by: str, descending: bool | Sequence[bool], nulls_last: bool) -> Self: if isinstance(descending, bool): ascending: bool | list[bool] = not descending else: ascending = [not d for d in descending] position = "last" if nulls_last else "first" return self._with_native( self.native.sort_values(list(by), ascending=ascending, na_position=position) ) def join( # noqa: C901 self, other: Self, *, how: JoinStrategy, left_on: Sequence[str] | None, right_on: Sequence[str] | None, suffix: str, ) -> Self: if how == "cross": key_token = generate_temporary_column_name( n_bytes=8, columns=[*self.columns, *other.columns] ) return self._with_native( self.native.assign(**{key_token: 0}) .merge( other.native.assign(**{key_token: 0}), how="inner", left_on=key_token, right_on=key_token, suffixes=("", suffix), ) .drop(columns=key_token) ) if how == "anti": indicator_token = generate_temporary_column_name( n_bytes=8, columns=[*self.columns, *other.columns] ) if right_on is None: # pragma: no cover msg = "`right_on` cannot be `None` in anti-join" raise TypeError(msg) other_native = ( select_columns_by_name( other.native, list(right_on), self._backend_version, self._implementation, ) .rename( # rename to avoid creating extra columns in join columns=dict(zip(right_on, left_on)) # type: ignore[arg-type] ) .drop_duplicates() ) df = self.native.merge( other_native, how="outer", indicator=indicator_token, # pyright: ignore[reportArgumentType] left_on=left_on, right_on=left_on, ) return self._with_native( df[df[indicator_token] == "left_only"].drop(columns=[indicator_token]) ) if how == "semi": if right_on is None: # pragma: no cover msg = "`right_on` cannot be `None` in semi-join" raise TypeError(msg) other_native = ( select_columns_by_name( other.native, list(right_on), self._backend_version, self._implementation, ) .rename( # rename to avoid creating extra columns in join columns=dict(zip(right_on, left_on)) # type: ignore[arg-type] ) .drop_duplicates() # avoids potential rows duplication from inner join ) return self._with_native( self.native.merge( other_native, how="inner", left_on=left_on, right_on=left_on ) ) if how == "left": result_native = self.native.merge( other.native, how="left", left_on=left_on, right_on=right_on, suffixes=("", suffix), ) extra = [] for left_key, right_key in zip(left_on, right_on): # type: ignore[arg-type] if right_key != left_key and right_key not in self.columns: extra.append(right_key) elif right_key != left_key: extra.append(f"{right_key}_right") return self._with_native(result_native.drop(columns=extra)) if how == "full": # dask does not retain keys post-join # we must append the suffix to each key before-hand # help mypy assert left_on is not None # noqa: S101 assert right_on is not None # noqa: S101 right_on_mapper = _remap_full_join_keys(left_on, right_on, suffix) other_native = other.native.rename(columns=right_on_mapper) check_column_names_are_unique(other_native.columns) right_on = list(right_on_mapper.values()) # we now have the suffixed keys return self._with_native( self.native.merge( other_native, left_on=left_on, right_on=right_on, how="outer", suffixes=("", suffix), ) ) return self._with_native( self.native.merge( other.native, left_on=left_on, right_on=right_on, how=how, suffixes=("", suffix), ) ) def join_asof( self, other: Self, *, left_on: str, right_on: str, by_left: Sequence[str] | None, by_right: Sequence[str] | None, strategy: AsofJoinStrategy, suffix: str, ) -> Self: plx = self.__native_namespace__() return self._with_native( plx.merge_asof( self.native, other.native, left_on=left_on, right_on=right_on, left_by=by_left, right_by=by_right, direction=strategy, suffixes=("", suffix), ) ) def group_by( self, keys: Sequence[str] | Sequence[DaskExpr], *, drop_null_keys: bool ) -> DaskLazyGroupBy: from narwhals._dask.group_by import DaskLazyGroupBy return DaskLazyGroupBy(self, keys, drop_null_keys=drop_null_keys) def tail(self, n: int) -> Self: # pragma: no cover native_frame = self.native n_partitions = native_frame.npartitions if n_partitions == 1: return self._with_native(self.native.tail(n=n, compute=False)) else: msg = "`LazyFrame.tail` is not supported for Dask backend with multiple partitions." raise NotImplementedError(msg) def gather_every(self, n: int, offset: int) -> Self: row_index_token = generate_temporary_column_name(n_bytes=8, columns=self.columns) plx = self.__narwhals_namespace__() return ( self.with_row_index(row_index_token) .filter( (plx.col(row_index_token) >= offset) & ((plx.col(row_index_token) - offset) % n == 0) ) .drop([row_index_token], strict=False) ) def unpivot( self, on: Sequence[str] | None, index: Sequence[str] | None, variable_name: str, value_name: str, ) -> Self: return self._with_native( self.native.melt( id_vars=index, value_vars=on, var_name=variable_name, value_name=value_name, ) ) explode = not_implemented()