bbqddt2/Antigravity
0
1from __future__ import annotations2 3import functools4import operator5import re6from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast7 8import numpy as np9import pandas as pd10 11from narwhals._compliant import EagerSeriesNamespace12from narwhals._constants import (13 MS_PER_SECOND,14 NS_PER_MICROSECOND,15 NS_PER_MILLISECOND,16 NS_PER_SECOND,17 SECONDS_PER_DAY,18 US_PER_SECOND,19)20from narwhals._exceptions import issue_warning21from narwhals._utils import (22 Implementation,23 Version,24 _DeferredIterable,25 check_columns_exist,26 isinstance_or_issubclass,27 parse_version,28 requires,29)30from narwhals.exceptions import ShapeError31 32if TYPE_CHECKING:33 from collections.abc import Callable, Iterable, Iterator, Mapping34 from types import ModuleType35 from typing import TypeAlias36 37 import pyarrow as pa38 from pandas._typing import Dtype as PandasDtype39 from pandas.core.dtypes.dtypes import BaseMaskedDtype40 from typing_extensions import TypeIs41 42 from narwhals._duration import IntervalUnit43 from narwhals._pandas_like.expr import PandasLikeExpr44 from narwhals._pandas_like.series import PandasLikeSeries45 from narwhals._pandas_like.typing import (46 NativeDataFrameT,47 NativeNDFrameT,48 NativeSeriesT,49 )50 from narwhals.dtypes import DType51 from narwhals.typing import DTypeBackend, IntoDType, TimeUnit, _1DArray52 53 ExprT = TypeVar("ExprT", bound=PandasLikeExpr)54 UnitCurrent: TypeAlias = TimeUnit55 UnitTarget: TypeAlias = TimeUnit56 BinOpBroadcast: TypeAlias = Callable[[Any, int], Any]57 IntoRhs: TypeAlias = int58 59 60PANDAS_LIKE_IMPLEMENTATION = {61 Implementation.PANDAS,62 Implementation.CUDF,63 Implementation.MODIN,64}65PD_DATETIME_RGX = r"""^66 datetime64\[67 (?P<time_unit>s|ms|us|ns) # Match time unit: s, ms, us, or ns68 (?:, # Begin non-capturing group for optional timezone69 \s* # Optional whitespace after comma70 (?P<time_zone> # Start named group for timezone71 [a-zA-Z\/]+ # Match timezone name, e.g., UTC, America/New_York72 (?:[+-]\d{2}:\d{2})? # Optional offset in format +HH:MM or -HH:MM73 | # OR74 pytz\.FixedOffset\(\d+\) # Match pytz.FixedOffset with integer offset in parentheses75 ) # End time_zone group76 )? # End optional timezone group77 \] # Closing bracket for datetime6478$"""79PATTERN_PD_DATETIME = re.compile(PD_DATETIME_RGX, re.VERBOSE)80PD_DURATION_RGX = r"""^81 timedelta64\[82 (?P<time_unit>s|ms|us|ns) # Match time unit: s, ms, us, or ns83 \] # Closing bracket for timedelta6484$"""85PATTERN_PD_DURATION = re.compile(PD_DURATION_RGX, re.VERBOSE)86 87NativeIntervalUnit: TypeAlias = Literal[88 "year",89 "quarter",90 "month",91 "week",92 "day",93 "hour",94 "minute",95 "second",96 "millisecond",97 "microsecond",98 "nanosecond",99]100ALIAS_DICT = {"d": "D", "m": "min"}101UNITS_DICT: Mapping[IntervalUnit, NativeIntervalUnit] = {102 "y": "year",103 "q": "quarter",104 "mo": "month",105 "d": "day",106 "h": "hour",107 "m": "minute",108 "s": "second",109 "ms": "millisecond",110 "us": "microsecond",111 "ns": "nanosecond",112}113 114PANDAS_VERSION = Implementation.PANDAS._backend_version()115"""Static backend version for `pandas`.116 117Always available if we reached here, due to a module-level import.118"""119 120NUMPY_VERSION = parse_version(np)121"""Static version for `numpy`.122 123Always available if we reached here, as imported in both _pandas_like/dataframe.py and124_pandas_like/series.py.125"""126 127 128def is_pandas_or_modin(implementation: Implementation) -> bool:129 return implementation in {Implementation.PANDAS, Implementation.MODIN}130 131 132def align_and_extract_native(133 lhs: PandasLikeSeries, rhs: PandasLikeSeries | object134) -> tuple[pd.Series[Any], pd.Series[Any] | object]:135 """Validate RHS of binary operation.136 137 If the comparison isn't supported, return `NotImplemented` so that the138 "right-hand-side" operation (e.g. `__radd__`) can be tried.139 """140 from narwhals._pandas_like.series import PandasLikeSeries141 142 lhs_index = lhs.native.index143 144 if lhs._broadcast and isinstance(rhs, PandasLikeSeries) and not rhs._broadcast:145 return lhs.native.iloc[0], rhs.native146 147 if isinstance(rhs, PandasLikeSeries):148 if rhs._broadcast:149 return (lhs.native, rhs.native.iloc[0])150 if rhs.native.index is not lhs_index:151 return (152 lhs.native,153 set_index(rhs.native, lhs_index, implementation=rhs._implementation),154 )155 return (lhs.native, rhs.native)156 157 if isinstance(rhs, list):158 msg = "Expected Series or scalar, got list."159 raise TypeError(msg)160 161 # `rhs` must be scalar, so just leave it as-is162 return lhs.native, rhs163 164 165def set_index(166 obj: NativeNDFrameT, index: Any, *, implementation: Implementation167) -> NativeNDFrameT:168 """Wrapper around pandas' set_axis to set object index.169 170 We can set `copy` / `inplace` based on implementation/version.171 """172 if isinstance(index, implementation.to_native_namespace().Index) and (173 expected_len := len(index)174 ) != (actual_len := len(obj)):175 msg = f"Expected object of length {expected_len}, got length: {actual_len}"176 raise ShapeError(msg)177 if implementation is Implementation.CUDF:178 obj = obj.copy(deep=False)179 obj.index = index180 return obj181 if implementation is Implementation.PANDAS and (182 (1, 5) <= implementation._backend_version() < (3,)183 ): # pragma: no cover184 return obj.set_axis(index, axis=0, copy=False)185 return obj.set_axis(index, axis=0) # pragma: no cover186 187 188def rename(189 obj: NativeNDFrameT, *args: Any, implementation: Implementation, **kwargs: Any190) -> NativeNDFrameT:191 """Wrapper around pandas' rename so that we can set `copy` based on implementation/version."""192 if implementation is Implementation.PANDAS and (193 implementation._backend_version() < (3,)194 ): # pragma: no cover195 result = obj.rename(*args, **kwargs, copy=False, inplace=False)196 else:197 result = obj.rename(*args, **kwargs)198 return cast("NativeNDFrameT", result)199 200 201@functools.lru_cache(maxsize=16)202def is_dtype_non_pyarrow_string(native_dtype: Any) -> bool:203 """*There is no problem which can't be solved by adding an extra string type* pandas."""204 # TODO @dangotbanned: Investigate how we could handle `cudf` without `str(native_dtype)`205 # https://github.com/rapidsai/cudf/blob/a32b8cf62c9b086b645b0825b78b99f065b1887f/python/cudf/cudf/utils/dtypes.py#L646-L670206 return isinstance(native_dtype, pd.StringDtype) or str(native_dtype) in {207 "string",208 "string[python]",209 "string[pyarrow_numpy]",210 "<StringDtype(na_value=nan)>", # why? why? why?211 "str",212 }213 214 215@functools.lru_cache(maxsize=16)216def non_object_native_to_narwhals_dtype(native_dtype: Any, version: Version) -> DType: # noqa: C901, PLR0912217 dtype = str(native_dtype)218 219 dtypes = version.dtypes220 if dtype in {"int64", "Int64"}:221 return dtypes.Int64()222 if dtype in {"int32", "Int32"}:223 return dtypes.Int32()224 if dtype in {"int16", "Int16"}:225 return dtypes.Int16()226 if dtype in {"int8", "Int8"}:227 return dtypes.Int8()228 if dtype in {"uint64", "UInt64"}:229 return dtypes.UInt64()230 if dtype in {"uint32", "UInt32"}:231 return dtypes.UInt32()232 if dtype in {"uint16", "UInt16"}:233 return dtypes.UInt16()234 if dtype in {"uint8", "UInt8"}:235 return dtypes.UInt8()236 if dtype in {"float64", "Float64"}:237 return dtypes.Float64()238 if dtype in {"float32", "Float32"}:239 return dtypes.Float32()240 if dtype == "float16":241 return dtypes.Float16()242 if is_dtype_non_pyarrow_string(native_dtype):243 return dtypes.String()244 if dtype in {"bool", "boolean"}:245 return dtypes.Boolean()246 if match_ := PATTERN_PD_DATETIME.match(dtype):247 dt_time_unit: TimeUnit = match_.group("time_unit") # type: ignore[assignment]248 dt_time_zone: str | None = match_.group("time_zone")249 return dtypes.Datetime(dt_time_unit, dt_time_zone)250 if match_ := PATTERN_PD_DURATION.match(dtype):251 du_time_unit: TimeUnit = match_.group("time_unit") # type: ignore[assignment]252 return dtypes.Duration(du_time_unit)253 return dtypes.Unknown() # pragma: no cover254 255 256def object_native_to_narwhals_dtype(257 series: PandasLikeSeries | None, version: Version, implementation: Implementation258) -> DType:259 dtypes = version.dtypes260 if implementation is Implementation.CUDF:261 # Per conversations with their maintainers, they don't support arbitrary262 # objects, so we can just return String.263 return dtypes.String()264 265 infer = pd.api.types.infer_dtype266 # Arbitrary limit of 100 elements to use to sniff dtype.267 inferred_dtype = "empty" if series is None else infer(series.head(100), skipna=True)268 if inferred_dtype == "string":269 return dtypes.String()270 if inferred_dtype == "empty" and version is not Version.V1:271 # Default to String for empty Series.272 return dtypes.String()273 if inferred_dtype == "empty":274 # But preserve returning Object in V1.275 return dtypes.Object()276 return dtypes.Object()277 278 279def native_categorical_to_narwhals_dtype(280 native_dtype: pd.CategoricalDtype, version: Version, implementation: Implementation281) -> DType:282 dtypes = version.dtypes283 if version is Version.V1:284 return dtypes.Categorical()285 if native_dtype.ordered:286 into_iter = (287 _cudf_categorical_to_list(native_dtype)288 if implementation is Implementation.CUDF289 else native_dtype.categories.to_list290 )291 return dtypes.Enum(_DeferredIterable(into_iter))292 return dtypes.Categorical()293 294 295def _cudf_categorical_to_list(296 native_dtype: Any,297) -> Callable[[], list[Any]]: # pragma: no cover298 # NOTE: https://docs.rapids.ai/api/cudf/stable/user_guide/api_docs/api/cudf.core.dtypes.categoricaldtype/#cudf.core.dtypes.CategoricalDtype299 # https://github.com/rapidsai/cudf/issues/18536300 # https://github.com/rapidsai/cudf/issues/14027301 def fn() -> list[Any]:302 return native_dtype.categories.to_arrow().to_pylist()303 304 return fn305 306 307CUDF_BASE_DTYPE_PREFIX = ("list", "struct", "decimal")308 309 310def native_to_narwhals_dtype(311 native_dtype: Any,312 version: Version,313 implementation: Implementation,314 *,315 allow_object: bool = False,316) -> DType:317 str_dtype = str(native_dtype)318 319 if is_dtype_pyarrow(native_dtype) or str_dtype.startswith(CUDF_BASE_DTYPE_PREFIX):320 from narwhals._arrow.utils import (321 native_to_narwhals_dtype as arrow_native_to_narwhals_dtype,322 )323 324 if hasattr(native_dtype, "to_arrow"): # pragma: no cover325 pa_dtype: pa.DataType = native_dtype.to_arrow() # pyright: ignore[reportAttributeAccessIssue]326 else:327 pa_dtype = native_dtype.pyarrow_dtype328 return arrow_native_to_narwhals_dtype(pa_dtype, version)329 if str_dtype == "category":330 return native_categorical_to_narwhals_dtype(native_dtype, version, implementation)331 if str_dtype != "object":332 return non_object_native_to_narwhals_dtype(native_dtype, version)333 if implementation is Implementation.DASK:334 # Per conversations with their maintainers, they don't support arbitrary335 # objects, so we can just return String.336 return version.dtypes.String()337 if allow_object: # pragma: no cover338 return object_native_to_narwhals_dtype(None, version, implementation)339 msg = (340 "Unreachable code, object dtype should be handled separately" # pragma: no cover341 )342 raise AssertionError(msg)343 344 345def is_dtype_numpy_nullable(dtype: Any) -> TypeIs[BaseMaskedDtype]:346 """Return `True` if `dtype` is `"numpy_nullable"`."""347 # NOTE: We need a sentinel as the positive case is `BaseMaskedDtype.base = None`348 # See https://github.com/narwhals-dev/narwhals/pull/2740#discussion_r2171667055349 sentinel = object()350 return (351 isinstance(dtype, pd.api.extensions.ExtensionDtype)352 and getattr(dtype, "base", sentinel) is None353 )354 355 356def get_dtype_backend(dtype: Any, implementation: Implementation) -> DTypeBackend:357 """Get dtype backend for pandas type.358 359 Matches pandas' `dtype_backend` argument in `convert_dtypes`.360 """361 if implementation is Implementation.CUDF:362 return None363 if is_dtype_pyarrow(dtype):364 return "pyarrow"365 return "numpy_nullable" if is_dtype_numpy_nullable(dtype) else None366 367 368# NOTE: Use this to avoid annotating inline369def iter_dtype_backends(370 dtypes: Iterable[Any], implementation: Implementation371) -> Iterator[DTypeBackend]:372 """Yield a `DTypeBackend` per-dtype.373 374 Matches pandas' `dtype_backend` argument in `convert_dtypes`.375 """376 return (get_dtype_backend(dtype, implementation) for dtype in dtypes)377 378 379@functools.lru_cache(maxsize=16)380def is_dtype_pyarrow(dtype: Any) -> TypeIs[pd.ArrowDtype]:381 return hasattr(pd, "ArrowDtype") and isinstance(dtype, pd.ArrowDtype)382 383 384dtypes = Version.MAIN.dtypes385NW_TO_PD_DTYPES_INVARIANT: Mapping[type[DType], str] = {386 # TODO(Unassigned): is there no pyarrow-backed categorical?387 # or at least, convert_dtypes(dtype_backend='pyarrow') doesn't388 # convert to it?389 dtypes.Categorical: "category",390 dtypes.Object: "object",391}392NW_TO_PD_DTYPES_BACKEND: Mapping[type[DType], Mapping[DTypeBackend, str | type[Any]]] = {393 dtypes.Float64: {394 "pyarrow": "Float64[pyarrow]",395 "numpy_nullable": "Float64",396 None: "float64",397 },398 dtypes.Float32: {399 "pyarrow": "Float32[pyarrow]",400 "numpy_nullable": "Float32",401 None: "float32",402 },403 dtypes.Float16: {404 # pandas has no nullable (extension) `Float16`, so `numpy_nullable` falls405 # back to the NumPy `float16` (nulls become `NaN`).406 "pyarrow": "Float16[pyarrow]",407 "numpy_nullable": "float16",408 None: "float16",409 },410 dtypes.Int64: {"pyarrow": "Int64[pyarrow]", "numpy_nullable": "Int64", None: "int64"},411 dtypes.Int32: {"pyarrow": "Int32[pyarrow]", "numpy_nullable": "Int32", None: "int32"},412 dtypes.Int16: {"pyarrow": "Int16[pyarrow]", "numpy_nullable": "Int16", None: "int16"},413 dtypes.Int8: {"pyarrow": "Int8[pyarrow]", "numpy_nullable": "Int8", None: "int8"},414 dtypes.UInt64: {415 "pyarrow": "UInt64[pyarrow]",416 "numpy_nullable": "UInt64",417 None: "uint64",418 },419 dtypes.UInt32: {420 "pyarrow": "UInt32[pyarrow]",421 "numpy_nullable": "UInt32",422 None: "uint32",423 },424 dtypes.UInt16: {425 "pyarrow": "UInt16[pyarrow]",426 "numpy_nullable": "UInt16",427 None: "uint16",428 },429 dtypes.UInt8: {"pyarrow": "UInt8[pyarrow]", "numpy_nullable": "UInt8", None: "uint8"},430 dtypes.Boolean: {431 "pyarrow": "boolean[pyarrow]",432 "numpy_nullable": "boolean",433 None: "bool",434 },435}436 437 438def narwhals_to_native_dtype( # noqa: C901, PLR0912439 dtype: IntoDType,440 dtype_backend: DTypeBackend,441 implementation: Implementation,442 version: Version,443) -> str | PandasDtype:444 if dtype_backend not in {None, "pyarrow", "numpy_nullable"}:445 msg = f"Expected one of {{None, 'pyarrow', 'numpy_nullable'}}, got: '{dtype_backend}'"446 raise ValueError(msg)447 dtypes = version.dtypes448 base_type = dtype.base_type()449 if pd_type := NW_TO_PD_DTYPES_INVARIANT.get(base_type):450 return pd_type451 if into_pd_type := NW_TO_PD_DTYPES_BACKEND.get(base_type):452 return into_pd_type[dtype_backend]453 if issubclass(base_type, dtypes.String):454 if dtype_backend == "pyarrow":455 import pyarrow as pa # ignore-banned-import456 457 # Note: this is different from `string[pyarrow]`, even though the repr458 # looks the same.459 # >>> pd.DataFrame({'a':['foo']}, dtype='string[pyarrow]')['a'].str.len()460 # 0 3461 # Name: a, dtype: Int64462 # >>> pd.DataFrame({'a':['foo']}, dtype=pd.ArrowDtype(pa.string()))['a'].str.len()463 # 0 3464 # Name: a, dtype: int32[pyarrow]465 #466 # `ArrowDType(pa.string())` is what `.convert_dtypes(dtype_backend='pyarrow')` converts to,467 # so we use that here.468 return pd.ArrowDtype(pa.string())469 if dtype_backend == "numpy_nullable":470 return "string"471 return str472 if isinstance_or_issubclass(dtype, dtypes.Datetime):473 if is_pandas_or_modin(implementation) and PANDAS_VERSION < (474 2,475 ): # pragma: no cover476 if isinstance(dtype, dtypes.Datetime) and dtype.time_unit != "ns":477 found = requires._unparse_version(PANDAS_VERSION)478 available = f"available in 'pandas>=2.0', found version {found!r}."479 changelog_url = "https://pandas.pydata.org/docs/dev/whatsnew/v2.0.0.html#construction-with-datetime64-or-timedelta64-dtype-with-unsupported-resolution"480 msg = (481 f"`nw.Datetime(time_unit={dtype.time_unit!r})` is only {available}\n"482 "Narwhals has fallen back to using `time_unit='ns'` to avoid an error.\n\n"483 "Hint: to avoid this warning, consider either:\n"484 f"- Upgrading pandas: {changelog_url}\n"485 f"- Using a bare `nw.Datetime`, if this precision is not important"486 )487 issue_warning(msg, UserWarning)488 dt_time_unit = "ns"489 else:490 dt_time_unit = dtype.time_unit491 492 if dtype_backend == "pyarrow":493 tz_part = f", tz={tz}" if (tz := dtype.time_zone) else ""494 return f"timestamp[{dt_time_unit}{tz_part}][pyarrow]"495 tz_part = f", {tz}" if (tz := dtype.time_zone) else ""496 return f"datetime64[{dt_time_unit}{tz_part}]"497 if isinstance_or_issubclass(dtype, dtypes.Duration):498 if is_pandas_or_modin(implementation) and PANDAS_VERSION < (499 2,500 ): # pragma: no cover501 du_time_unit = "ns"502 else:503 du_time_unit = dtype.time_unit504 return (505 f"duration[{du_time_unit}][pyarrow]"506 if dtype_backend == "pyarrow"507 else f"timedelta64[{du_time_unit}]"508 )509 if isinstance_or_issubclass(dtype, dtypes.Date):510 try:511 import pyarrow as pa # ignore-banned-import512 except ModuleNotFoundError as exc:513 # BUG: Never re-raised?514 msg = "'pyarrow>=13.0.0' is required for `Date` dtype."515 raise ModuleNotFoundError(msg) from exc516 return "date32[pyarrow]"517 if isinstance_or_issubclass(dtype, dtypes.Enum):518 if version is Version.V1:519 msg = "Converting to Enum is not supported in narwhals.stable.v1"520 raise NotImplementedError(msg)521 if isinstance(dtype, dtypes.Enum):522 ns = implementation.to_native_namespace()523 return ns.CategoricalDtype(dtype.categories, ordered=True)524 msg = "Can not cast / initialize Enum without categories present"525 raise ValueError(msg)526 if issubclass(527 base_type,528 (529 dtypes.Struct,530 dtypes.Array,531 dtypes.List,532 dtypes.Time,533 dtypes.Binary,534 dtypes.Decimal,535 ),536 ):537 return narwhals_to_native_arrow_dtype(dtype, implementation, version)538 msg = f"Unknown dtype: {dtype}" # pragma: no cover539 raise AssertionError(msg)540 541 542def narwhals_to_native_arrow_dtype(543 dtype: IntoDType, implementation: Implementation, version: Version544) -> pd.ArrowDtype:545 if is_pandas_or_modin(implementation) and PANDAS_VERSION >= (2, 2):546 try:547 import pyarrow as pa # ignore-banned-import # noqa: F401548 except ImportError as exc: # pragma: no cover549 msg = (550 f"Unable to convert to {dtype} due to the following exception: {exc.msg}"551 )552 raise ImportError(msg) from exc553 from narwhals._arrow.utils import narwhals_to_native_dtype as _to_arrow_dtype554 555 return pd.ArrowDtype(_to_arrow_dtype(dtype, version))556 msg = ( # pragma: no cover557 f"Converting to {dtype} dtype is not supported for implementation "558 f"{implementation} and version {version}."559 )560 raise NotImplementedError(msg)561 562 563def int_dtype_mapper(dtype: Any) -> str:564 if "pyarrow" in str(dtype):565 return "Int64[pyarrow]"566 if str(dtype).lower() != str(dtype): # pragma: no cover567 return "Int64"568 return "int64"569 570 571_TIMESTAMP_DATETIME_OP_FACTOR: Mapping[572 tuple[UnitCurrent, UnitTarget], tuple[BinOpBroadcast, IntoRhs]573] = {574 ("ns", "us"): (operator.floordiv, 1_000),575 ("ns", "ms"): (operator.floordiv, 1_000_000),576 ("us", "ns"): (operator.mul, NS_PER_MICROSECOND),577 ("us", "ms"): (operator.floordiv, 1_000),578 ("ms", "ns"): (operator.mul, NS_PER_MILLISECOND),579 ("ms", "us"): (operator.mul, 1_000),580 ("s", "ns"): (operator.mul, NS_PER_SECOND),581 ("s", "us"): (operator.mul, US_PER_SECOND),582 ("s", "ms"): (operator.mul, MS_PER_SECOND),583}584 585 586def calculate_timestamp_datetime(587 s: NativeSeriesT, current: TimeUnit, time_unit: TimeUnit588) -> NativeSeriesT:589 if current == time_unit:590 return s591 if item := _TIMESTAMP_DATETIME_OP_FACTOR.get((current, time_unit)):592 fn, factor = item593 return fn(s, factor)594 msg = ( # pragma: no cover595 f"unexpected time unit {current}, please report an issue at "596 "https://github.com/narwhals-dev/narwhals"597 )598 raise AssertionError(msg)599 600 601_TIMESTAMP_DATE_FACTOR: Mapping[TimeUnit, int] = {602 "ns": NS_PER_SECOND,603 "us": US_PER_SECOND,604 "ms": MS_PER_SECOND,605 "s": 1,606}607 608 609def calculate_timestamp_date(s: NativeSeriesT, time_unit: TimeUnit) -> NativeSeriesT:610 return s * SECONDS_PER_DAY * _TIMESTAMP_DATE_FACTOR[time_unit]611 612 613def select_columns_by_name(614 df: NativeDataFrameT,615 column_names: list[str] | _1DArray, # NOTE: Cannot be a tuple!616 implementation: Implementation,617) -> NativeDataFrameT | Any:618 """Select columns by name.619 620 Prefer this over `df.loc[:, column_names]` as it's621 generally more performant.622 """623 if len(column_names) == df.shape[1] and (df.columns == column_names).all():624 return df625 if (df.columns.dtype.kind == "b") or (626 implementation is Implementation.PANDAS627 and implementation._backend_version() < (1, 5)628 ):629 # See https://github.com/narwhals-dev/narwhals/issues/1349#issuecomment-2470118122630 # for why we need this631 if error := check_columns_exist(column_names, available=df.columns.tolist()):632 raise error633 return df.loc[:, column_names]634 try:635 return df[column_names]636 except KeyError as e:637 if error := check_columns_exist(column_names, available=df.columns.tolist()):638 raise error from e639 raise640 641 642def is_non_nullable_boolean(s: PandasLikeSeries) -> bool:643 # cuDF booleans are nullable but the native dtype is still 'bool'.644 return (645 s._implementation646 in {Implementation.PANDAS, Implementation.MODIN, Implementation.DASK}647 and s.native.dtype == "bool"648 )649 650 651def import_array_module(implementation: Implementation, /) -> ModuleType:652 """Returns numpy or cupy module depending on the given implementation."""653 if implementation in {Implementation.PANDAS, Implementation.MODIN}:654 import numpy as np655 656 return np657 if implementation is Implementation.CUDF:658 import cupy as cp # ignore-banned-import # cuDF dependency.659 660 return cp661 msg = f"Expected pandas/modin/cudf, got: {implementation}" # pragma: no cover662 raise AssertionError(msg)663 664 665class PandasLikeSeriesNamespace(EagerSeriesNamespace["PandasLikeSeries", Any]): ...666 667 668def make_group_by_kwargs(*, drop_null_keys: bool) -> dict[str, bool]:669 return {"sort": False, "as_index": True, "dropna": drop_null_keys, "observed": True}670 671 672def broadcast_series_to_index(673 native: pd.Series[Any],674 index: Any,675 *,676 is_nested: bool,677 series_class: type[pd.Series[Any]],678) -> pd.Series[Any]:679 """Broadcast a scalar value from a (one element) Series to match a target index.680 681 For nested (arrow-backed) types, we rely on682 [`pandas.array`](https://pandas.pydata.org/docs/reference/api/pandas.array.html).683 684 Arguments:685 native: The native pandas-like Series containing the scalar value to broadcast.686 index: The target index to broadcast to.687 is_nested: Whether the Series has a nested (arrow-backed) dtype.688 series_class: Series class to use for constructing the result.689 690 Returns:691 A new Series with the scalar value broadcast to match the target index.692 """693 value = native.iloc[0]694 if is_nested:695 from narwhals._arrow.utils import repeat696 697 # NOTE: Ignore typing because `pandas-stubs` are wrong698 # TODO(FBruzzesi): Should we pass the `copy=False` flag?699 pa_array = pd.array(repeat(value, len(index)), dtype=native.dtype) # type: ignore[arg-type]700 701 return series_class(pa_array, index=index, name=native.name)702 703 return series_class(value, index=index, dtype=native.dtype, name=native.name)704 705 706def binary_string_sum_fallback( # pragma: no cover707 left: pd.Series, right: Any, pdx: Any708) -> pd.Series:709 # Workaround some upstream issues:710 # - https://github.com/pandas-dev/pandas/issues/64393711 # - https://github.com/pandas-dev/pandas/issues/65220712 left_dtype = left.dtype713 left_dtype_str = str(left_dtype)714 if left_dtype_str == "large_string[pyarrow]" and isinstance(right, str):715 import pyarrow as pa # ignore-banned-import716 717 return left + pa.scalar(right, type=pa.large_string())718 if isinstance(right, pdx.Series):719 right_dtype = right.dtype720 if left_dtype_str == "object":721 # Only for pandas pre 3.0. Anything is better than `object`, so take RHS.722 return left.astype(right_dtype) + right723 if hasattr(left.values, "__arrow_array__") and hasattr(724 right.values, "__arrow_array__"725 ):726 import pyarrow as pa # ignore-banned-import727 728 left_arrow = left.values.__arrow_array__().type # noqa: PD011 # type: ignore[attr-defined]729 right_arrow = right.values.__arrow_array__().type # noqa: PD011 # type: ignore[attr-defined]730 if pa.types.is_string(left_arrow) and pa.types.is_large_string(right_arrow):731 # https://github.com/pandas-dev/pandas/blob/b00d4f6710ff6c1c80319196657c31c2cf6c70ff/pandas/core/arrays/arrow/array.py#L1064-L1068732 pd_pa_large_string = pd.ArrowDtype(pa.large_string())733 return left.astype(pd_pa_large_string) + right.astype(pd_pa_large_string)734 else:735 pass736 # Give precedence to the left-hand-side dtype.737 return left + right.astype(left_dtype)738 return left + right739