bbqddt2/Antigravity
0
1from __future__ import annotations2 3import warnings4from typing import TYPE_CHECKING, Any, cast5 6from narwhals._compliant import EagerExpr7from narwhals._expression_parsing import evaluate_nodes, evaluate_output_names_and_aliases8from narwhals._pandas_like.group_by import _REMAP_ORDERED_INDEX, PandasLikeGroupBy9from narwhals._pandas_like.series import PandasLikeSeries10from narwhals._pandas_like.utils import make_group_by_kwargs11from narwhals._utils import generate_temporary_column_name12 13if TYPE_CHECKING:14 from collections.abc import Sequence15 16 from typing_extensions import Self17 18 from narwhals._compliant.typing import (19 AliasNames,20 EvalNames,21 EvalSeries,22 NarwhalsAggregation,23 )24 from narwhals._pandas_like.dataframe import PandasLikeDataFrame25 from narwhals._pandas_like.namespace import PandasLikeNamespace26 from narwhals._utils import Implementation, Version, _LimitedContext27 from narwhals.typing import PythonLiteral28 29WINDOW_FUNCTIONS_TO_PANDAS_EQUIVALENT = {30 "cum_sum": "cumsum",31 "cum_min": "cummin",32 "cum_max": "cummax",33 "cum_prod": "cumprod",34 # Pandas cumcount starts counting from 0 while Polars starts from 135 # Pandas cumcount counts nulls while Polars does not36 # So, instead of using "cumcount" we use "cumsum" on notna() to get the same result37 "cum_count": "cumsum",38 "rolling_sum": "sum",39 "rolling_mean": "mean",40 "rolling_std": "std",41 "rolling_var": "var",42 "shift": "shift",43 "rank": "rank",44 "diff": "diff",45 "fill_null": "fillna",46 "quantile": "quantile",47 "ewm_mean": "mean",48}49 50 51def window_kwargs_to_pandas_equivalent( # noqa: C90152 function_name: str, kwargs: dict[str, Any]53) -> dict[str, PythonLiteral]:54 if function_name == "shift":55 assert "n" in kwargs # noqa: S10156 pandas_kwargs: dict[str, PythonLiteral] = {"periods": kwargs["n"]}57 elif function_name == "rank":58 assert "method" in kwargs # noqa: S10159 assert "descending" in kwargs # noqa: S10160 _method = kwargs["method"]61 pandas_kwargs = {62 "method": "first" if _method == "ordinal" else _method,63 "ascending": not kwargs["descending"],64 "na_option": "keep",65 "pct": False,66 }67 elif function_name.startswith("cum_"): # Cumulative operation68 pandas_kwargs = {"skipna": True}69 elif function_name == "n_unique":70 pandas_kwargs = {"dropna": False}71 elif function_name.startswith("rolling_"): # Rolling operation72 assert "min_samples" in kwargs # noqa: S10173 assert "window_size" in kwargs # noqa: S10174 assert "center" in kwargs # noqa: S10175 pandas_kwargs = {76 "min_periods": kwargs["min_samples"],77 "window": kwargs["window_size"],78 "center": kwargs["center"],79 }80 elif function_name in {"std", "var"}:81 assert "ddof" in kwargs # noqa: S10182 pandas_kwargs = {"ddof": kwargs["ddof"]}83 elif function_name == "fill_null":84 assert "strategy" in kwargs # noqa: S10185 assert "limit" in kwargs # noqa: S10186 pandas_kwargs = {"strategy": kwargs["strategy"], "limit": kwargs["limit"]}87 elif function_name == "quantile":88 assert "quantile" in kwargs # noqa: S10189 assert "interpolation" in kwargs # noqa: S10190 pandas_kwargs = {91 "q": kwargs["quantile"],92 "interpolation": kwargs["interpolation"],93 }94 elif function_name.startswith("ewm_"):95 assert "com" in kwargs # noqa: S10196 assert "span" in kwargs # noqa: S10197 assert "half_life" in kwargs # noqa: S10198 assert "alpha" in kwargs # noqa: S10199 assert "adjust" in kwargs # noqa: S101100 assert "min_samples" in kwargs # noqa: S101101 assert "ignore_nulls" in kwargs # noqa: S101102 103 pandas_kwargs = {104 "com": kwargs["com"],105 "span": kwargs["span"],106 "halflife": kwargs["half_life"],107 "alpha": kwargs["alpha"],108 "adjust": kwargs["adjust"],109 "min_periods": kwargs["min_samples"],110 "ignore_na": kwargs["ignore_nulls"],111 }112 elif function_name in {"first", "last", "any_value"}:113 if kwargs.get("ignore_nulls"):114 msg = (115 "`Expr.any_value(ignore_nulls=True)` is not supported in a `over` "116 "context for pandas-like backend."117 )118 raise NotImplementedError(msg)119 pandas_kwargs = {120 "n": _REMAP_ORDERED_INDEX[cast("NarwhalsAggregation", function_name)]121 }122 else: # sum, len, ...123 pandas_kwargs = {}124 return pandas_kwargs125 126 127class PandasLikeExpr(EagerExpr["PandasLikeDataFrame", PandasLikeSeries]):128 def __init__(129 self,130 call: EvalSeries[PandasLikeDataFrame, PandasLikeSeries],131 *,132 evaluate_output_names: EvalNames[PandasLikeDataFrame],133 alias_output_names: AliasNames | None,134 implementation: Implementation,135 version: Version,136 ) -> None:137 self._call = call138 self._evaluate_output_names = evaluate_output_names139 self._alias_output_names = alias_output_names140 self._implementation = implementation141 self._version = version142 143 def __narwhals_namespace__(self) -> PandasLikeNamespace:144 from narwhals._pandas_like.namespace import PandasLikeNamespace145 146 return PandasLikeNamespace(self._implementation, version=self._version)147 148 @classmethod149 def from_column_names(150 cls: type[Self],151 evaluate_column_names: EvalNames[PandasLikeDataFrame],152 /,153 *,154 context: _LimitedContext,155 ) -> Self:156 def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:157 try:158 return [159 PandasLikeSeries(160 df._native_frame[column_name],161 implementation=df._implementation,162 version=df._version,163 )164 for column_name in evaluate_column_names(df)165 ]166 except KeyError as e:167 if error := df._check_columns_exist(evaluate_column_names(df)):168 raise error from e169 raise170 171 return cls(172 func,173 evaluate_output_names=evaluate_column_names,174 alias_output_names=None,175 implementation=context._implementation,176 version=context._version,177 )178 179 @classmethod180 def from_column_indices(cls, *column_indices: int, context: _LimitedContext) -> Self:181 def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:182 native = df.native183 return [184 PandasLikeSeries.from_native(native.iloc[:, i], context=df)185 for i in column_indices186 ]187 188 return cls(189 func,190 evaluate_output_names=cls._eval_names_indices(column_indices),191 alias_output_names=None,192 implementation=context._implementation,193 version=context._version,194 )195 196 def ewm_mean(197 self,198 *,199 com: float | None,200 span: float | None,201 half_life: float | None,202 alpha: float | None,203 adjust: bool,204 min_samples: int,205 ignore_nulls: bool,206 ) -> Self:207 return self._reuse_series(208 "ewm_mean",209 com=com,210 span=span,211 half_life=half_life,212 alpha=alpha,213 adjust=adjust,214 min_samples=min_samples,215 ignore_nulls=ignore_nulls,216 )217 218 def _over_without_partition_by(self, order_by: Sequence[str]) -> Self:219 # e.g. `nw.col('a').cum_sum().order_by(key)`220 # We can always easily support this as it doesn't require grouping.221 222 def func(df: PandasLikeDataFrame) -> Sequence[PandasLikeSeries]:223 token = generate_temporary_column_name(8, df.columns)224 df = df.with_row_index(token, order_by=None).sort(225 *order_by, descending=False, nulls_last=False226 )227 results = self(df.drop([token], strict=True))228 meta = self._metadata229 if meta is not None and meta.is_scalar_like:230 # We need to broadcast the result to the original size, since231 # `over` is a length-preserving operation.232 index = df.native.index233 ns = self._implementation.to_native_namespace()234 return [235 s._with_native(ns.Series(s.item(), index=index, name=s.name))236 for s in results237 ]238 239 sorting_indices = df.get_column(token)240 for s in results:241 s.scatter(sorting_indices, s, in_place=True)242 return results243 244 return self.__class__(245 func,246 evaluate_output_names=self._evaluate_output_names,247 alias_output_names=self._alias_output_names,248 implementation=self._implementation,249 version=self._version,250 )251 252 def over( # noqa: C901, PLR0915253 self, partition_by: Sequence[str], order_by: Sequence[str]254 ) -> Self:255 if not partition_by:256 assert order_by # noqa: S101257 return self._over_without_partition_by(order_by)258 259 # We have something like prev.leaf().over(...) (e.g. `nw.col('a').sum().over('b')`), where:260 # - `prev` must be elementwise (in the example: `nw.col('a')`)261 # - `leaf` must be a "simple" function, i.e. one that pandas supports in `transform`262 # (in the example: `sum`)263 #264 # We first evaluate `prev` as-is, and then evaluate `leaf().over(...)`` by using `transform`265 # or other DataFrameGroupBy methods.266 meta = self._metadata267 if partition_by and (meta.prev is not None and not meta.prev.is_elementwise):268 msg = (269 "Only elementary expressions are supported for `.over` in pandas-like backends "270 "when `partition_by` is specified.\n\n"271 "Please see: "272 "https://narwhals-dev.github.io/narwhals/concepts/improve_group_by_operation/"273 )274 raise NotImplementedError(msg)275 nodes = list(reversed(list(self._metadata.iter_nodes_reversed())))276 277 leaf_node = nodes[-1]278 function_name = leaf_node.name279 pandas_agg = PandasLikeGroupBy._REMAP_AGGS.get(280 cast("NarwhalsAggregation", function_name)281 )282 pandas_function_name = WINDOW_FUNCTIONS_TO_PANDAS_EQUIVALENT.get(283 function_name, pandas_agg284 )285 if pandas_function_name is None:286 msg = (287 f"Unsupported function: {function_name} in `over` context.\n\n"288 f"Supported functions are {', '.join(WINDOW_FUNCTIONS_TO_PANDAS_EQUIVALENT)}\n"289 f"and {', '.join(PandasLikeGroupBy._REMAP_AGGS)}."290 )291 raise NotImplementedError(msg)292 scalar_kwargs = leaf_node.kwargs293 pandas_kwargs = window_kwargs_to_pandas_equivalent(function_name, scalar_kwargs)294 295 def func(df: PandasLikeDataFrame) -> Sequence[PandasLikeSeries]: # noqa: C901, PLR0912, PLR0914, PLR0915296 assert pandas_function_name is not None # help mypy # noqa: S101297 plx = self.__narwhals_namespace__()298 if meta.prev is not None:299 df = df.with_columns(300 cast("PandasLikeExpr", evaluate_nodes(nodes[:-1], plx))301 )302 _, aliases = evaluate_output_names_and_aliases(self, df, [])303 if function_name == "cum_count":304 df = df.with_columns(~plx.col(*aliases).is_null())305 306 if function_name.startswith("cum_"):307 assert "reverse" in scalar_kwargs # noqa: S101308 reverse = scalar_kwargs["reverse"]309 else:310 assert "reverse" not in scalar_kwargs # noqa: S101311 reverse = False312 313 if order_by:314 columns = list(set(partition_by).union(aliases).union(order_by))315 token = generate_temporary_column_name(8, columns)316 df = (317 df.simple_select(*columns)318 .with_row_index(token, order_by=None)319 .sort(320 *partition_by, *order_by, descending=reverse, nulls_last=reverse321 )322 )323 sorting_indices = df.get_column(token)324 elif reverse:325 columns = list(set(partition_by).union(aliases))326 df = df.simple_select(*columns)._gather_slice(slice(None, None, -1))327 group_by_kwargs = make_group_by_kwargs(drop_null_keys=False)328 grouped = df._native_frame.groupby(partition_by, **group_by_kwargs)329 if function_name.startswith("rolling"):330 rolling = grouped[list(aliases)].rolling(**pandas_kwargs)331 if pandas_function_name in {"std", "var"}:332 assert "ddof" in scalar_kwargs # noqa: S101333 res_native = getattr(rolling, pandas_function_name)(334 ddof=scalar_kwargs["ddof"]335 )336 else:337 res_native = getattr(rolling, pandas_function_name)()338 elif function_name.startswith("ewm"):339 ewm = grouped[list(aliases)].ewm(**pandas_kwargs)340 assert pandas_function_name is not None # help mypy # noqa: S101341 res_native = getattr(ewm, pandas_function_name)()342 elif function_name == "fill_null":343 assert "strategy" in scalar_kwargs # noqa: S101344 assert "limit" in scalar_kwargs # noqa: S101345 df_grouped = grouped[list(aliases)]346 if scalar_kwargs["strategy"] == "forward":347 res_native = df_grouped.ffill(limit=scalar_kwargs["limit"])348 elif scalar_kwargs["strategy"] == "backward":349 res_native = df_grouped.bfill(limit=scalar_kwargs["limit"])350 else: # pragma: no cover351 # This is deprecated in pandas. Indeed, `nw.col('a').fill_null(3).over('b')`352 # does not seem very useful, and DuckDB doesn't support it either.353 msg = "`fill_null` with `over` without `strategy` specified is not supported."354 raise NotImplementedError(msg)355 elif function_name == "len":356 if len(aliases) != 1: # pragma: no cover357 msg = "Safety check failed, please report a bug."358 raise AssertionError(msg)359 res_native = grouped.transform("size").to_frame(aliases[0])360 elif function_name in {"first", "last", "any_value"}:361 with warnings.catch_warnings():362 # Ignore settingwithcopy warnings/errors, they're false-positives here.363 warnings.filterwarnings("ignore", message="\n.*copy of a slice")364 _agg = getattr(365 grouped[[*partition_by, *aliases]], pandas_function_name366 )(**pandas_kwargs)367 impl = self._implementation368 backend_version = impl._backend_version()369 if impl.is_pandas() and backend_version < (3, 0): # pragma: no cover370 # NOTE: Keep `inplace=True` to avoid making a redundant copy.371 _agg.reset_index(drop=True, inplace=True)372 else:373 _agg = _agg.reset_index(drop=True)374 375 keys = list(partition_by)376 res_native = df.native[keys].merge(_agg, on=keys)[list(aliases)]377 else:378 res_native = grouped[list(aliases)].transform(379 pandas_function_name, **pandas_kwargs380 )381 result_frame = df._with_native(res_native)382 results = [result_frame.get_column(name) for name in aliases]383 if order_by:384 with warnings.catch_warnings():385 # Ignore settingwithcopy warnings/errors, they're false-positives here.386 warnings.filterwarnings("ignore", message="\n.*copy of a slice")387 for s in results:388 s.scatter(sorting_indices, s, in_place=True)389 return results390 if reverse:391 return [s._gather_slice(slice(None, None, -1)) for s in results]392 return results393 394 return self.__class__(395 func,396 evaluate_output_names=self._evaluate_output_names,397 alias_output_names=self._alias_output_names,398 implementation=self._implementation,399 version=self._version,400 )401