bbqddt2/Antigravity
0
1from __future__ import annotations2 3from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence4from itertools import chain, product5from typing import TYPE_CHECKING, Any, Literal, cast, overload6 7import numpy as np8 9from narwhals._compliant import EagerDataFrame10from narwhals._pandas_like.series import PANDAS_TO_NUMPY_DTYPE_MISSING, PandasLikeSeries11from narwhals._pandas_like.utils import (12 align_and_extract_native,13 broadcast_series_to_index,14 get_dtype_backend,15 import_array_module,16 iter_dtype_backends,17 narwhals_to_native_dtype,18 native_to_narwhals_dtype,19 object_native_to_narwhals_dtype,20 rename,21 select_columns_by_name,22 set_index,23)24from narwhals._typing_compat import assert_never25from narwhals._utils import (26 Implementation,27 _into_arrow_table,28 _remap_full_join_keys,29 check_column_names_are_unique,30 exclude_column_names,31 generate_temporary_column_name,32 parse_columns_to_drop,33 scale_bytes,34)35from narwhals.dependencies import is_pandas_like_dataframe36from narwhals.exceptions import InvalidOperationError, ShapeError37from narwhals.functions import col as nw_col38 39if TYPE_CHECKING:40 from io import BytesIO41 from pathlib import Path42 from types import ModuleType43 from typing import TypeAlias44 45 import pandas as pd46 import polars as pl47 from typing_extensions import Self, TypeIs48 49 from narwhals._compliant.typing import CompliantDataFrameAny, CompliantLazyFrameAny50 from narwhals._pandas_like.expr import PandasLikeExpr51 from narwhals._pandas_like.group_by import PandasLikeGroupBy52 from narwhals._pandas_like.namespace import PandasLikeNamespace53 from narwhals._spark_like.utils import SparkSession54 from narwhals._translate import IntoArrowTable55 from narwhals._typing import _EagerAllowedImpl, _LazyAllowedImpl56 from narwhals._utils import Version, _LimitedContext57 from narwhals.dtypes import DType58 from narwhals.typing import (59 AsofJoinStrategy,60 DTypeBackend,61 IntoSchema,62 JoinStrategy,63 PivotAgg,64 SizedMultiIndexSelector,65 SizedMultiNameSelector,66 SizeUnit,67 UniqueKeepStrategy,68 _2DArray,69 _SliceIndex,70 _SliceName,71 )72 73 Constructor: TypeAlias = Callable[..., pd.DataFrame]74 75 76CLASSICAL_NUMPY_DTYPES: frozenset[np.dtype[Any]] = frozenset(77 [78 np.dtype("float64"),79 np.dtype("float32"),80 np.dtype("int64"),81 np.dtype("int32"),82 np.dtype("int16"),83 np.dtype("int8"),84 np.dtype("uint64"),85 np.dtype("uint32"),86 np.dtype("uint16"),87 np.dtype("uint8"),88 np.dtype("bool"),89 np.dtype("datetime64[s]"),90 np.dtype("datetime64[ms]"),91 np.dtype("datetime64[us]"),92 np.dtype("datetime64[ns]"),93 np.dtype("timedelta64[s]"),94 np.dtype("timedelta64[ms]"),95 np.dtype("timedelta64[us]"),96 np.dtype("timedelta64[ns]"),97 np.dtype("object"),98 ]99)100 101 102class PandasLikeDataFrame(103 EagerDataFrame["PandasLikeSeries", "PandasLikeExpr", "Any", "pd.Series[Any]"]104):105 def __init__(106 self,107 native_dataframe: Any,108 *,109 implementation: Implementation,110 version: Version,111 validate_column_names: bool,112 validate_backend_version: bool = False,113 ) -> None:114 self._native_frame = native_dataframe115 self._implementation = implementation116 self._version = version117 if validate_column_names:118 check_column_names_are_unique(native_dataframe.columns)119 if validate_backend_version:120 self._validate_backend_version()121 122 @classmethod123 def from_arrow(cls, data: IntoArrowTable, /, *, context: _LimitedContext) -> Self:124 implementation = context._implementation125 tbl = _into_arrow_table(data, context)126 if implementation.is_pandas():127 native = tbl.to_pandas()128 elif implementation.is_modin():129 # NOTE: Function moved + deprecated (0.26.0), then old path removed (0.31.0)130 # https://github.com/modin-project/modin/pull/6806131 # https://github.com/modin-project/modin/pull/7274132 if implementation._backend_version() >= (0, 26, 0):133 from modin.pandas.io import from_arrow as mpd_from_arrow134 else: # pragma: no cover135 from modin.pandas.utils import (136 from_arrow as mpd_from_arrow, # pyright: ignore[reportAttributeAccessIssue]137 )138 native = mpd_from_arrow(tbl)139 elif implementation.is_cudf(): # pragma: no cover140 native = implementation.to_native_namespace().DataFrame.from_arrow(tbl)141 else: # pragma: no cover142 msg = "congratulations, you entered unreachable code - please report a bug"143 raise AssertionError(msg)144 return cls.from_native(native, context=context)145 146 @classmethod147 def from_dict(148 cls,149 data: Mapping[str, Any],150 /,151 *,152 context: _LimitedContext,153 schema: IntoSchema | Mapping[str, DType | None] | None,154 ) -> Self:155 implementation = context._implementation156 pdx = implementation.to_native_namespace()157 Series = cast("type[pd.Series[Any]]", pdx.Series)158 DataFrame = cast("type[pd.DataFrame]", pdx.DataFrame)159 aligned_data: dict[str, pd.Series[Any] | Any] = {}160 left_most: PandasLikeSeries | None = None161 for name, series in data.items():162 if isinstance(series, Series):163 compliant = PandasLikeSeries.from_native(series, context=context)164 if left_most is None:165 left_most = compliant166 aligned_data[name] = series167 else:168 aligned_data[name] = align_and_extract_native(left_most, compliant)[1]169 else:170 aligned_data[name] = series171 if aligned_data or not schema:172 native = DataFrame.from_dict(aligned_data)173 else:174 native = DataFrame.from_dict({col: [] for col in schema})175 if schema:176 backends: Iterable[DTypeBackend]177 if aligned_data:178 backends = iter_dtype_backends(native.dtypes, implementation)179 else:180 backends = (None for _ in range(len(schema)))181 native_schema = {182 key: narwhals_to_native_dtype(183 dtype,184 backend,185 implementation=context._implementation,186 version=context._version,187 )188 for ((key, dtype), backend) in zip(schema.items(), backends, strict=False)189 if dtype is not None190 }191 native = native.astype(native_schema)192 return cls.from_native(native, context=context)193 194 @classmethod195 def from_dicts(196 cls,197 data: Sequence[Mapping[str, Any]],198 /,199 *,200 context: _LimitedContext,201 schema: IntoSchema | Mapping[str, DType | None] | None,202 ) -> Self:203 implementation = context._implementation204 ns = implementation.to_native_namespace()205 DataFrame = cast("type[pd.DataFrame]", ns.DataFrame)206 if data or not schema:207 native = DataFrame.from_records(data)208 else:209 native = DataFrame.from_dict({col: [] for col in schema})210 if schema:211 backends: Iterable[DTypeBackend]212 if data:213 backends = iter_dtype_backends(native.dtypes, implementation)214 else:215 backends = (None for _ in range(len(schema)))216 native_schema = {217 key: narwhals_to_native_dtype(218 dtype,219 backend,220 implementation=context._implementation,221 version=context._version,222 )223 for ((key, dtype), backend) in zip(schema.items(), backends, strict=False)224 if dtype is not None225 }226 native = native.astype(native_schema)227 return cls.from_native(native, context=context)228 229 @staticmethod230 def _is_native(obj: Any) -> TypeIs[Any]:231 return is_pandas_like_dataframe(obj) # pragma: no cover232 233 @classmethod234 def from_native(cls, data: Any, /, *, context: _LimitedContext) -> Self:235 return cls(236 data,237 implementation=context._implementation,238 version=context._version,239 validate_column_names=True,240 )241 242 @classmethod243 def from_numpy(244 cls,245 data: _2DArray,246 /,247 *,248 context: _LimitedContext,249 schema: IntoSchema | Sequence[str] | None,250 ) -> Self:251 from narwhals.schema import Schema252 253 implementation = context._implementation254 DataFrame: Constructor = implementation.to_native_namespace().DataFrame255 if isinstance(schema, (Mapping, Schema)):256 it: Iterable[DTypeBackend] = (257 get_dtype_backend(native_type, implementation)258 for native_type in schema.values()259 )260 native = DataFrame(data, columns=schema.keys()).astype(261 Schema(schema).to_pandas(it)262 )263 else:264 native = DataFrame(data, columns=cls._numpy_column_names(data, schema))265 return cls.from_native(native, context=context)266 267 def __narwhals_dataframe__(self) -> Self:268 return self269 270 def __narwhals_lazyframe__(self) -> Self:271 return self272 273 def __narwhals_namespace__(self) -> PandasLikeNamespace:274 from narwhals._pandas_like.namespace import PandasLikeNamespace275 276 return PandasLikeNamespace(self._implementation, version=self._version)277 278 def __native_namespace__(self) -> ModuleType:279 if self._implementation in {280 Implementation.PANDAS,281 Implementation.MODIN,282 Implementation.CUDF,283 }:284 return self._implementation.to_native_namespace()285 286 msg = f"Expected pandas/modin/cudf, got: {type(self._implementation)}" # pragma: no cover287 raise AssertionError(msg)288 289 def __len__(self) -> int:290 return len(self.native)291 292 def _with_version(self, version: Version) -> Self:293 return self.__class__(294 self.native,295 implementation=self._implementation,296 version=version,297 validate_column_names=False,298 )299 300 def _with_native(self, df: Any, *, validate_column_names: bool = True) -> Self:301 return self.__class__(302 df,303 implementation=self._implementation,304 version=self._version,305 validate_column_names=validate_column_names,306 )307 308 def _extract_comparand(self, other: PandasLikeSeries) -> pd.Series[Any]:309 index = self.native.index310 if other._broadcast:311 native = other.native312 is_nested = other.dtype.is_nested()313 return broadcast_series_to_index(314 native, index, is_nested=is_nested, series_class=type(native)315 )316 317 if (len_other := len(other)) != (len_idx := len(index)):318 msg = f"Expected object of length {len_idx}, got: {len_other}."319 raise ShapeError(msg)320 if other.native.index is not index:321 return set_index(other.native, index, implementation=other._implementation)322 return other.native # pragma: no cover323 324 @property325 def _array_funcs(self): # type: ignore[no-untyped-def] # noqa: ANN202326 if TYPE_CHECKING:327 import numpy as np328 329 return np330 return import_array_module(self._implementation)331 332 def get_column(self, name: str) -> PandasLikeSeries:333 return PandasLikeSeries.from_native(self.native[name], context=self)334 335 def __array__(self, dtype: Any = None, *, copy: bool | None = None) -> _2DArray:336 return self.to_numpy(dtype=dtype, copy=copy)337 338 def _gather(self, rows: SizedMultiIndexSelector[pd.Series[Any]]) -> Self:339 items = list(rows) if isinstance(rows, tuple) else rows340 return self._with_native(self.native.iloc[items, :])341 342 def _gather_slice(self, rows: _SliceIndex | range) -> Self:343 return self._with_native(344 self.native.iloc[slice(rows.start, rows.stop, rows.step), :],345 validate_column_names=False,346 )347 348 def _select_slice_name(self, columns: _SliceName) -> Self:349 start = (350 self.native.columns.get_loc(columns.start)351 if columns.start is not None352 else None353 )354 stop = (355 self.native.columns.get_loc(columns.stop) + 1356 if columns.stop is not None357 else None358 )359 selector = slice(start, stop, columns.step)360 return self._with_native(361 self.native.iloc[:, selector], validate_column_names=False362 )363 364 def _select_slice_index(self, columns: _SliceIndex | range) -> Self:365 return self._with_native(366 self.native.iloc[:, columns], validate_column_names=False367 )368 369 def _select_multi_index(370 self, columns: SizedMultiIndexSelector[pd.Series[Any]]371 ) -> Self:372 columns = list(columns) if isinstance(columns, tuple) else columns373 return self._with_native(374 self.native.iloc[:, columns], validate_column_names=False375 )376 377 def _select_multi_name(self, columns: SizedMultiNameSelector[pd.Series[Any]]) -> Self:378 return self._with_native(self.native.loc[:, columns])379 380 # --- properties ---381 @property382 def columns(self) -> list[str]:383 return self.native.columns.tolist()384 385 @overload386 def rows(self, *, named: Literal[True]) -> list[dict[str, Any]]: ...387 388 @overload389 def rows(self, *, named: Literal[False]) -> list[tuple[Any, ...]]: ...390 391 @overload392 def rows(self, *, named: bool) -> list[tuple[Any, ...]] | list[dict[str, Any]]: ...393 394 def rows(self, *, named: bool) -> list[tuple[Any, ...]] | list[dict[str, Any]]:395 if not named:396 # cuDF does not support itertuples. But it does support to_dict!397 if self._implementation is Implementation.CUDF:398 # Extract the row values from the named rows399 return [tuple(row.values()) for row in self.rows(named=True)]400 401 return list(self.native.itertuples(index=False, name=None))402 403 return self.native.to_dict(orient="records")404 405 def iter_columns(self) -> Iterator[PandasLikeSeries]:406 for _name, series in self.native.items(): # noqa: PERF102407 yield PandasLikeSeries.from_native(series, context=self)408 409 _iter_columns = iter_columns410 411 def iter_rows(412 self, *, named: bool, buffer_size: int413 ) -> Iterator[tuple[Any, ...]] | Iterator[dict[str, Any]]:414 # The param ``buffer_size`` is only here for compatibility with the Polars API415 # and has no effect on the output.416 if not named:417 yield from self.native.itertuples(index=False, name=None)418 else:419 col_names = self.native.columns420 for row in self.native.itertuples(index=False):421 yield dict(zip(col_names, row, strict=False))422 423 @property424 def schema(self) -> dict[str, DType]:425 native_dtypes = self.native.dtypes426 return {427 col: native_to_narwhals_dtype(428 native_dtypes[col], self._version, self._implementation429 )430 if native_dtypes[col] != "object"431 else object_native_to_narwhals_dtype(432 self.native[col], self._version, self._implementation433 )434 for col in self.native.columns435 }436 437 def collect_schema(self) -> dict[str, DType]:438 return self.schema439 440 # --- reshape ---441 def simple_select(self, *column_names: str) -> Self:442 return self._with_native(443 select_columns_by_name(self.native, list(column_names), self._implementation),444 validate_column_names=False,445 )446 447 def select(self, *exprs: PandasLikeExpr) -> Self:448 new_series = self._evaluate_exprs(*exprs)449 if not new_series:450 # return empty dataframe, like Polars does451 return self._with_native(type(self.native)(), validate_column_names=False)452 new_series = new_series[0]._align_full_broadcast(*new_series)453 namespace = self.__narwhals_namespace__()454 df = namespace._concat_horizontal([s.native for s in new_series])455 # `concat` creates a new object, so fine to modify `.columns.name` inplace.456 df.columns.name = self.native.columns.name457 return self._with_native(df, validate_column_names=True)458 459 def drop_nulls(self, subset: Sequence[str] | None) -> Self:460 if subset is None:461 return self._with_native(462 self.native.dropna(axis=0), validate_column_names=False463 )464 plx = self.__narwhals_namespace__()465 mask = ~plx.any_horizontal(plx.col(*subset).is_null(), ignore_nulls=True)466 return self.filter(mask)467 468 def estimated_size(self, unit: SizeUnit) -> int | float:469 sz = self.native.memory_usage(deep=True).sum()470 return scale_bytes(sz, unit=unit)471 472 def with_row_index(self, name: str, order_by: Sequence[str] | None) -> Self:473 plx = self.__narwhals_namespace__()474 if order_by is None:475 data = self._array_funcs.arange(len(self))476 row_index = plx._expr._from_series(477 plx._series.from_iterable(478 data, context=self, index=self.native.index, name=name479 )480 )481 else:482 rank = cast(483 "PandasLikeExpr",484 nw_col(order_by[0]).rank(method="ordinal")._to_compliant_expr(plx),485 )486 row_index = (487 rank.over(partition_by=[], order_by=order_by)488 - plx.lit(1, None).broadcast()489 ).alias(name)490 return self.select(row_index, plx.all())491 492 def row(self, index: int) -> tuple[Any, ...]:493 return tuple(x for x in self.native.iloc[index])494 495 def filter(self, predicate: PandasLikeExpr) -> Self:496 mask = self._evaluate_single_output_expr(predicate)497 mask_native = self._extract_comparand(mask)498 return self._with_native(499 self.native.loc[mask_native], validate_column_names=False500 )501 502 def with_columns(self, *exprs: PandasLikeExpr) -> Self:503 columns = self._evaluate_exprs(*exprs)504 if not columns and len(self) == 0:505 return self506 name_columns: dict[str, PandasLikeSeries] = {s.name: s for s in columns}507 to_concat = []508 # Make sure to preserve column order509 for name in self.native.columns:510 if name in name_columns:511 series = self._extract_comparand(name_columns.pop(name))512 else:513 series = self.native[name]514 to_concat.append(series)515 to_concat.extend(self._extract_comparand(s) for s in name_columns.values())516 namespace = self.__narwhals_namespace__()517 df = namespace._concat_horizontal(to_concat)518 # `concat` creates a new object, so fine to modify `.columns.name` inplace.519 df.columns.name = self.native.columns.name520 return self._with_native(df, validate_column_names=False)521 522 def rename(self, mapping: Mapping[str, str]) -> Self:523 return self._with_native(524 rename(self.native, columns=mapping, implementation=self._implementation)525 )526 527 def drop(self, columns: Sequence[str], *, strict: bool) -> Self:528 to_drop = parse_columns_to_drop(self, columns, strict=strict)529 return self._with_native(530 self.native.drop(columns=to_drop), validate_column_names=False531 )532 533 # --- transform ---534 def sort(self, *by: str, descending: bool | Sequence[bool], nulls_last: bool) -> Self:535 df = self.native536 if isinstance(descending, bool):537 ascending: bool | list[bool] = not descending538 else:539 ascending = [not d for d in descending]540 na_position = "last" if nulls_last else "first"541 return self._with_native(542 df.sort_values(list(by), ascending=ascending, na_position=na_position),543 validate_column_names=False,544 )545 546 def top_k(self, k: int, *, by: Iterable[str], reverse: bool | Sequence[bool]) -> Self:547 df = self.native548 schema = self.schema549 if isinstance(reverse, bool) and all(schema[x].is_numeric() for x in by):550 if reverse:551 return self._with_native(df.nsmallest(k, by))552 return self._with_native(df.nlargest(k, by))553 return self._with_native(554 df.sort_values(list(by), ascending=reverse).head(k),555 validate_column_names=False,556 )557 558 # --- convert ---559 def collect(560 self, backend: _EagerAllowedImpl | None, **kwargs: Any561 ) -> CompliantDataFrameAny:562 if backend is None:563 return PandasLikeDataFrame(564 self.native,565 implementation=self._implementation,566 version=self._version,567 validate_column_names=False,568 )569 570 if backend is Implementation.PANDAS:571 kwds: dict[str, Any] = {572 "implementation": Implementation.PANDAS,573 "version": self._version,574 "validate_column_names": False,575 }576 if backend is not self._implementation: # pragma: no cover577 kwds.update(validate_backend_version=True)578 return PandasLikeDataFrame(self.to_pandas(), **kwds)579 580 if backend is Implementation.PYARROW:581 from narwhals._arrow.dataframe import ArrowDataFrame582 583 return ArrowDataFrame(584 native_dataframe=self.to_arrow(),585 validate_backend_version=True,586 version=self._version,587 validate_column_names=False,588 )589 590 if backend is Implementation.POLARS:591 from narwhals._polars.dataframe import PolarsDataFrame592 593 return PolarsDataFrame(594 df=self.to_polars(), validate_backend_version=True, version=self._version595 )596 597 msg = f"Unsupported `backend` value: {backend}" # pragma: no cover598 raise ValueError(msg) # pragma: no cover599 600 # --- actions ---601 def group_by(602 self, keys: Sequence[str] | Sequence[PandasLikeExpr], *, drop_null_keys: bool603 ) -> PandasLikeGroupBy:604 from narwhals._pandas_like.group_by import PandasLikeGroupBy605 606 return PandasLikeGroupBy(self, keys, drop_null_keys=drop_null_keys)607 608 def _join_inner(609 self, other: Self, *, left_on: Sequence[str], right_on: Sequence[str], suffix: str610 ) -> pd.DataFrame:611 return self.native.dropna(subset=left_on, how="any").merge(612 other.native,613 left_on=left_on,614 right_on=right_on,615 how="inner",616 suffixes=("", suffix),617 )618 619 def _join_left(620 self, other: Self, *, left_on: Sequence[str], right_on: Sequence[str], suffix: str621 ) -> pd.DataFrame:622 result_native = self.native.merge(623 other.native.dropna(subset=right_on, how="any"),624 how="left",625 left_on=left_on,626 right_on=right_on,627 suffixes=("", suffix),628 )629 extra = [630 right_key if right_key not in self.columns else f"{right_key}{suffix}"631 for left_key, right_key in zip(left_on, right_on, strict=True)632 if right_key != left_key633 ]634 impl = self._implementation635 if impl.is_pandas() and impl._backend_version() < (3, 0): # pragma: no cover636 # NOTE: Keep `inplace=True` to avoid making a redundant copy.637 result_native.drop(columns=extra, inplace=True) # noqa: PD002638 return result_native639 640 return result_native.drop(columns=extra)641 642 def _join_full(643 self, other: Self, *, left_on: Sequence[str], right_on: Sequence[str], suffix: str644 ) -> pd.DataFrame:645 # Pandas coalesces keys in full joins unless there's no collision646 ns = self.__narwhals_namespace__()647 self_native = self.native648 right_on_mapper = _remap_full_join_keys(left_on, right_on, suffix)649 other_native = other.native.rename(columns=right_on_mapper)650 check_column_names_are_unique(other_native.columns)651 right_suffixed = list(right_on_mapper.values())652 653 left_null_mask = self_native[list(left_on)].isna().any(axis=1)654 right_null_mask = other_native[right_suffixed].isna().any(axis=1)655 656 # We need to add suffix to `other` columns overlapping in `self` if not in keys657 to_rename = set(other.columns).intersection(self.columns).difference(right_on)658 right_null_rows = other_native[right_null_mask].rename(659 columns={col: f"{col}{suffix}" for col in to_rename}660 )661 662 join_result = self_native[~left_null_mask].merge(663 other_native[~right_null_mask],664 left_on=left_on,665 right_on=right_suffixed,666 how="outer",667 suffixes=("", suffix),668 )669 670 return ns._concat_diagonal(671 [join_result, self_native[left_null_mask], right_null_rows]672 )673 674 def _join_cross(self, other: Self, *, suffix: str) -> pd.DataFrame:675 impl = self._implementation676 backend_version = impl._backend_version()677 if (impl.is_modin() or impl.is_cudf()) or (678 impl.is_pandas() and backend_version < (1, 4)679 ): # pragma: no cover680 key_token = generate_temporary_column_name(681 n_bytes=8, columns=(*self.columns, *other.columns)682 )683 result_native = self.native.assign(**{key_token: 0}).merge(684 other.native.assign(**{key_token: 0}),685 how="inner",686 left_on=key_token,687 right_on=key_token,688 suffixes=("", suffix),689 )690 if impl.is_pandas() and backend_version < (3, 0): # pragma: no cover691 # NOTE: Keep `inplace=True` to avoid making a redundant copy.692 result_native.drop(columns=key_token, inplace=True) # noqa: PD002693 return result_native694 return result_native.drop(columns=key_token)695 696 return self.native.merge(other.native, how="cross", suffixes=("", suffix))697 698 def _join_semi(699 self, other: Self, *, left_on: Sequence[str], right_on: Sequence[str]700 ) -> pd.DataFrame:701 other_native = self._join_filter_rename(702 other=other,703 columns_to_select=list(right_on),704 columns_mapping=dict(zip(right_on, left_on, strict=False)),705 )706 return self.native.dropna(subset=left_on, how="any").merge(707 other_native, how="inner", left_on=left_on, right_on=left_on708 )709 710 def _join_anti(711 self, other: Self, *, left_on: Sequence[str], right_on: Sequence[str]712 ) -> pd.DataFrame:713 impl = self._implementation714 715 if impl.is_cudf():716 return self.native.merge(717 other.native.dropna(subset=left_on, how="any"),718 how="leftanti",719 left_on=left_on,720 right_on=right_on,721 )722 723 indicator_token = generate_temporary_column_name(724 n_bytes=8, columns=(*self.columns, *other.columns)725 )726 727 other_native = self._join_filter_rename(728 other=other,729 columns_to_select=list(right_on),730 columns_mapping=dict(zip(right_on, left_on, strict=True)),731 )732 result_native = self.native.merge(733 other_native.dropna(subset=left_on, how="any"),734 # TODO(FBruzzesi): See https://github.com/modin-project/modin/issues/7384735 how="left" if impl.is_pandas() else "outer",736 indicator=indicator_token,737 left_on=left_on,738 right_on=left_on,739 ).loc[lambda t: t[indicator_token] == "left_only"]740 741 if impl.is_pandas() and impl._backend_version() < (3, 0): # pragma: no cover742 # NOTE: Keep `inplace=True` to avoid making a redundant copy.743 result_native.drop(columns=indicator_token, inplace=True) # noqa: PD002744 return result_native745 746 return result_native.drop(columns=indicator_token)747 748 def _join_filter_rename(749 self, other: Self, columns_to_select: list[str], columns_mapping: dict[str, str]750 ) -> pd.DataFrame:751 """Helper function to avoid creating extra columns and row duplication.752 753 Used in `"anti"` and `"semi`" join's.754 755 Notice that a native object is returned.756 """757 implementation = self._implementation758 return rename(759 select_columns_by_name(760 other.native,761 column_names=columns_to_select,762 implementation=implementation,763 ),764 columns=columns_mapping,765 implementation=implementation,766 ).drop_duplicates()767 768 def join(769 self,770 other: Self,771 *,772 how: JoinStrategy,773 left_on: Sequence[str] | None,774 right_on: Sequence[str] | None,775 suffix: str,776 ) -> Self:777 if how == "cross":778 result = self._join_cross(other=other, suffix=suffix)779 780 elif left_on is None or right_on is None: # pragma: no cover781 raise ValueError(left_on, right_on)782 783 elif how == "inner":784 result = self._join_inner(785 other=other, left_on=left_on, right_on=right_on, suffix=suffix786 )787 elif how == "anti":788 result = self._join_anti(other=other, left_on=left_on, right_on=right_on)789 elif how == "semi":790 result = self._join_semi(other=other, left_on=left_on, right_on=right_on)791 elif how == "left":792 result = self._join_left(793 other=other, left_on=left_on, right_on=right_on, suffix=suffix794 )795 elif how == "full":796 result = self._join_full(797 other=other, left_on=left_on, right_on=right_on, suffix=suffix798 )799 else:800 assert_never(how)801 802 return self._with_native(result)803 804 def join_asof(805 self,806 other: Self,807 *,808 left_on: str,809 right_on: str,810 by_left: Sequence[str] | None,811 by_right: Sequence[str] | None,812 strategy: AsofJoinStrategy,813 suffix: str,814 ) -> Self:815 plx = self.__native_namespace__()816 return self._with_native(817 plx.merge_asof(818 self.native,819 other.native,820 left_on=left_on,821 right_on=right_on,822 left_by=by_left,823 right_by=by_right,824 direction=strategy,825 suffixes=("", suffix),826 )827 )828 829 # --- partial reduction ---830 831 def head(self, n: int) -> Self:832 return self._with_native(self.native.head(n), validate_column_names=False)833 834 def tail(self, n: int) -> Self:835 return self._with_native(self.native.tail(n), validate_column_names=False)836 837 def unique(838 self,839 subset: Sequence[str] | None,840 *,841 keep: UniqueKeepStrategy,842 maintain_order: bool | None = None,843 order_by: Sequence[str] | None,844 ) -> Self:845 # The param `maintain_order` is only here for compatibility with the Polars API846 # and has no effect on the output.847 mapped_keep = {"none": False, "any": "first"}.get(keep, keep)848 if subset and (error := self._check_columns_exist(subset)):849 raise error850 if order_by and maintain_order:851 token = generate_temporary_column_name(8, self.columns)852 res = (853 self.with_row_index(token, order_by=None)854 .sort(*order_by, nulls_last=False, descending=False)855 .native.drop_duplicates(subset or self.columns, keep=mapped_keep)856 .sort_values(token)857 )858 if (859 self._implementation.is_pandas()860 and self._implementation._backend_version() < (3, 0)861 ):862 res.drop(columns=token, inplace=True) # noqa: PD002 # pragma: no cover863 else:864 res = res.drop(columns=token)865 elif order_by:866 res = self.sort(867 *order_by, nulls_last=False, descending=False868 ).native.drop_duplicates(subset or self.columns, keep=mapped_keep)869 else:870 res = self.native.drop_duplicates(subset or self.columns, keep=mapped_keep)871 return self._with_native(res, validate_column_names=False)872 873 # --- lazy-only ---874 def lazy(875 self,876 backend: _LazyAllowedImpl | None = None,877 *,878 session: SparkSession | None = None,879 ) -> CompliantLazyFrameAny:880 pandas_df = self.to_pandas()881 if backend is None:882 return self883 if backend is Implementation.DUCKDB:884 import duckdb # ignore-banned-import885 886 from narwhals._duckdb.dataframe import DuckDBLazyFrame887 888 return DuckDBLazyFrame(889 df=duckdb.table("pandas_df"),890 validate_backend_version=True,891 version=self._version,892 )893 if backend is Implementation.POLARS:894 import polars as pl # ignore-banned-import895 896 from narwhals._polars.dataframe import PolarsLazyFrame897 898 return PolarsLazyFrame(899 df=pl.from_pandas(pandas_df).lazy(),900 validate_backend_version=True,901 version=self._version,902 )903 if backend is Implementation.DASK:904 import dask.dataframe as dd # ignore-banned-import905 906 from narwhals._dask.dataframe import DaskLazyFrame907 908 return DaskLazyFrame(909 native_dataframe=dd.from_pandas(pandas_df),910 validate_backend_version=True,911 version=self._version,912 )913 if backend is Implementation.IBIS:914 import ibis # ignore-banned-import915 916 from narwhals._ibis.dataframe import IbisLazyFrame917 918 return IbisLazyFrame(919 ibis.memtable(pandas_df, columns=self.columns),920 validate_backend_version=True,921 version=self._version,922 )923 924 if backend.is_spark_like():925 from narwhals._spark_like.dataframe import SparkLikeLazyFrame926 927 if session is None:928 msg = "Spark like backends require `session` to be not None."929 raise ValueError(msg)930 931 return SparkLikeLazyFrame(932 session.createDataFrame(pandas_df),933 version=self._version,934 implementation=backend,935 validate_backend_version=True,936 )937 938 raise AssertionError # pragma: no cover939 940 @property941 def shape(self) -> tuple[int, int]:942 return self.native.shape943 944 def to_dict(self, *, as_series: bool) -> dict[str, Any]:945 if as_series:946 return {947 col: PandasLikeSeries.from_native(self.native[col], context=self)948 for col in self.columns949 }950 return self.native.to_dict(orient="list")951 952 def to_numpy(self, dtype: Any = None, *, copy: bool | None = None) -> _2DArray:953 native_dtypes = self.native.dtypes954 955 if copy is None:956 # pandas default differs from Polars, but cuDF default is True957 copy = self._implementation is Implementation.CUDF958 959 if native_dtypes.isin(CLASSICAL_NUMPY_DTYPES).all():960 # Fast path, no conversions necessary.961 if dtype is not None:962 return self.native.to_numpy(dtype=dtype, copy=copy)963 return self.native.to_numpy(copy=copy)964 965 dtype_datetime = self._version.dtypes.Datetime966 to_convert = [967 key968 for key, val in self.schema.items()969 if isinstance(val, dtype_datetime) and val.time_zone is not None970 ]971 if to_convert:972 df = self.with_columns(973 self.__narwhals_namespace__()974 .col(*to_convert)975 .dt.convert_time_zone("UTC")976 .dt.replace_time_zone(None)977 ).native978 else:979 df = self.native980 981 if dtype is not None:982 return df.to_numpy(dtype=dtype, copy=copy)983 984 # pandas return `object` dtype for nullable dtypes if dtype=None,985 # so we cast each Series to numpy and let numpy find a common dtype.986 # If there aren't any dtypes where `to_numpy()` is "broken" (i.e. it987 # returns Object) then we just call `to_numpy()` on the DataFrame.988 for col_dtype in native_dtypes:989 if str(col_dtype) in PANDAS_TO_NUMPY_DTYPE_MISSING:990 arr: Any = np.hstack(991 [992 self.get_column(col).to_numpy(copy=copy, dtype=None)[:, None]993 for col in self.columns994 ]995 )996 return arr997 return df.to_numpy(copy=copy)998 999 def to_pandas(self) -> pd.DataFrame:1000 if self._implementation is Implementation.PANDAS:1001 return self.native1002 if self._implementation is Implementation.CUDF:1003 return self.native.to_pandas()1004 if self._implementation is Implementation.MODIN:1005 return self.native._to_pandas()1006 msg = f"Unknown implementation: {self._implementation}" # pragma: no cover1007 raise AssertionError(msg)1008 1009 def to_polars(self) -> pl.DataFrame:1010 import polars as pl # ignore-banned-import1011 1012 return pl.from_pandas(self.to_pandas())1013 1014 def write_parquet(self, file: str | Path | BytesIO) -> None:1015 self.native.to_parquet(file)1016 1017 @overload1018 def write_csv(self, file: None) -> str: ...1019 1020 @overload1021 def write_csv(self, file: str | Path | BytesIO) -> None: ...1022 1023 def write_csv(self, file: str | Path | BytesIO | None) -> str | None:1024 return self.native.to_csv(file, index=False)1025 1026 # --- descriptive ---1027 def is_unique(self) -> PandasLikeSeries:1028 return PandasLikeSeries.from_native(1029 ~self.native.duplicated(keep=False), context=self1030 )1031 1032 def item(self, row: int | None, column: int | str | None) -> Any:1033 if row is None and column is None:1034 if (shape := self.shape) != (1, 1):1035 msg = (1036 'can only call `.item()` without "row" or "column" values if the '1037 f"DataFrame has a single element; shape={shape!r}"1038 )1039 raise ValueError(msg)1040 return self.native.iloc[0, 0]1041 1042 if row is None or column is None:1043 msg = "cannot call `.item()` with only one of `row` or `column`"1044 raise ValueError(msg)1045 1046 _col = self.columns.index(column) if isinstance(column, str) else column1047 return self.native.iloc[row, _col]1048 1049 def clone(self) -> Self:1050 return self._with_native(self.native.copy(), validate_column_names=False)1051 1052 def gather_every(self, n: int, offset: int) -> Self:1053 return self._with_native(self.native.iloc[offset::n], validate_column_names=False)1054 1055 def _pivot_into_index_values(1056 self,1057 on: Sequence[str],1058 index: Sequence[str] | None,1059 values: Sequence[str] | None,1060 /,1061 ) -> tuple[Sequence[str], Sequence[str]]:1062 index = index or (1063 exclude_column_names(self, {*on, *values})1064 if values1065 else exclude_column_names(self, on)1066 )1067 values = values or exclude_column_names(self, {*on, *index})1068 return index, values1069 1070 @staticmethod1071 def _pivot_multi_on_name(unique_values: tuple[str, ...], /) -> str:1072 LB, RB, Q = "{", "}", '"' # noqa: N8061073 body = '","'.join(unique_values)1074 return f"{LB}{Q}{body}{Q}{RB}"1075 1076 @staticmethod1077 def _pivot_single_on_names(1078 column_names: Iterable[str], n_values: int, separator: str, /1079 ) -> list[str]:1080 if n_values > 1:1081 return [separator.join(col).strip() for col in column_names]1082 return [col[-1] for col in column_names]1083 1084 def _pivot_multi_on_names(1085 self,1086 column_names: Iterable[tuple[str, ...]],1087 n_on: int,1088 n_values: int,1089 separator: str,1090 /,1091 ) -> Iterator[str]:1092 if n_values > 1:1093 for col in column_names:1094 names = col[-n_on:]1095 prefix = col[0]1096 yield separator.join((prefix, self._pivot_multi_on_name(names)))1097 else:1098 for col in column_names:1099 yield self._pivot_multi_on_name(col[-n_on:])1100 1101 def _pivot_remap_column_names(1102 self, column_names: Iterable[Any], *, n_on: int, n_values: int, separator: str1103 ) -> list[str]:1104 """Reformat output column names from a native pivot operation, to match `polars`.1105 1106 Note:1107 `column_names` is a `pd.MultiIndex`, but not in the stubs.1108 """1109 if n_on == 1:1110 return self._pivot_single_on_names(column_names, n_values, separator)1111 return list(self._pivot_multi_on_names(column_names, n_on, n_values, separator))1112 1113 def _pivot_table(1114 self,1115 on: Sequence[str],1116 index: Sequence[str],1117 values: Sequence[str],1118 aggregate_function: Literal[1119 "min", "max", "first", "last", "sum", "mean", "median"1120 ],1121 /,1122 ) -> Any:1123 kwds: dict[Any, Any] = (1124 {} if self._implementation is Implementation.CUDF else {"observed": True}1125 )1126 return self.native.pivot_table(1127 values=values,1128 index=index,1129 columns=on,1130 aggfunc=aggregate_function,1131 margins=False,1132 **kwds,1133 )1134 1135 def _pivot(1136 self,1137 on: Sequence[str],1138 index: Sequence[str],1139 values: Sequence[str],1140 aggregate_function: PivotAgg | None,1141 /,1142 ) -> pd.DataFrame:1143 if aggregate_function is None:1144 return self.native.pivot(columns=on, index=index, values=values)1145 if aggregate_function == "len":1146 return (1147 self.native.groupby([*on, *index], as_index=False)1148 .agg(dict.fromkeys(values, "size"))1149 .pivot(columns=on, index=index, values=values)1150 )1151 return self._pivot_table(on, index, values, aggregate_function)1152 1153 def pivot(1154 self,1155 on: Sequence[str],1156 *,1157 index: Sequence[str] | None,1158 values: Sequence[str] | None,1159 aggregate_function: PivotAgg | None,1160 sort_columns: bool,1161 separator: str,1162 ) -> Self:1163 implementation = self._implementation1164 if implementation.is_modin():1165 msg = "pivot is not supported for Modin backend due to https://github.com/modin-project/modin/issues/7409."1166 raise NotImplementedError(msg)1167 1168 index, values = self._pivot_into_index_values(on, index, values)1169 result = self._pivot(on, index, values, aggregate_function)1170 1171 # Select the columns in the right order1172 uniques = (1173 (1174 self.get_column(col)1175 .unique()1176 .sort(descending=False, nulls_last=False)1177 .to_list()1178 for col in on1179 )1180 if sort_columns1181 else (self.get_column(col).unique().to_list() for col in on)1182 )1183 ordered_cols = list(product(values, *chain(uniques)))1184 result = result.loc[:, ordered_cols]1185 columns = result.columns1186 remapped = self._pivot_remap_column_names(1187 columns, n_on=len(on), n_values=len(values), separator=separator1188 )1189 result.columns = remapped1190 result.columns.names = [""]1191 return self._with_native(result.reset_index())1192 1193 def to_arrow(self) -> Any:1194 if self._implementation is Implementation.CUDF:1195 return self.native.to_arrow(preserve_index=False)1196 1197 import pyarrow as pa # ignore-banned-import()1198 1199 return pa.Table.from_pandas(self.native)1200