narwhals-2.11.0/narwhals/__init__.py0000644000000000000000000000630213615410400014263 0ustar00from __future__ import annotations import typing as _t from narwhals import dependencies, dtypes, exceptions, selectors from narwhals._utils import ( Implementation, generate_temporary_column_name, is_ordered_categorical, maybe_align_index, maybe_convert_dtypes, maybe_get_index, maybe_reset_index, maybe_set_index, ) from narwhals.dataframe import DataFrame, LazyFrame from narwhals.dtypes import ( Array, Binary, Boolean, Categorical, Date, Datetime, Decimal, Duration, Enum, Field, Float32, Float64, Int8, Int16, Int32, Int64, Int128, List, Object, String, Struct, Time, UInt8, UInt16, UInt32, UInt64, UInt128, Unknown, ) from narwhals.expr import Expr from narwhals.functions import ( all_ as all, all_horizontal, any_horizontal, coalesce, col, concat, concat_str, exclude, format, from_arrow, from_dict, from_dicts, from_numpy, len_ as len, lit, max, max_horizontal, mean, mean_horizontal, median, min, min_horizontal, new_series, nth, read_csv, read_parquet, scan_csv, scan_parquet, show_versions, sum, sum_horizontal, when, ) from narwhals.schema import Schema from narwhals.series import Series from narwhals.translate import ( from_native, get_native_namespace, narwhalify, to_native, to_py_scalar, ) __version__: str __all__ = [ "Array", "Binary", "Boolean", "Categorical", "DataFrame", "Date", "Datetime", "Decimal", "Duration", "Enum", "Expr", "Field", "Float32", "Float64", "Implementation", "Int8", "Int16", "Int32", "Int64", "Int128", "LazyFrame", "List", "Object", "Schema", "Series", "String", "Struct", "Time", "UInt8", "UInt16", "UInt32", "UInt64", "UInt128", "Unknown", "all", "all_horizontal", "any_horizontal", "coalesce", "col", "concat", "concat_str", "dependencies", "dtypes", "exceptions", "exclude", "format", "from_arrow", "from_dict", "from_dicts", "from_native", "from_numpy", "generate_temporary_column_name", "get_native_namespace", "is_ordered_categorical", "len", "lit", "max", "max_horizontal", "maybe_align_index", "maybe_convert_dtypes", "maybe_get_index", "maybe_reset_index", "maybe_set_index", "mean", "mean_horizontal", "median", "min", "min_horizontal", "narwhalify", "new_series", "nth", "read_csv", "read_parquet", "scan_csv", "scan_parquet", "selectors", "show_versions", "sum", "sum_horizontal", "to_native", "to_py_scalar", "when", ] def __getattr__(name: _t.Literal["__version__"]) -> str: # type: ignore[misc] if name == "__version__": global __version__ # noqa: PLW0603 from importlib import metadata __version__ = metadata.version(__name__) return __version__ msg = f"module {__name__!r} has no attribute {name!r}" raise AttributeError(msg) narwhals-2.11.0/narwhals/_constants.py0000644000000000000000000000210613615410400014675 0ustar00from __future__ import annotations import datetime as dt # Temporal (from `polars._utils.constants`) SECONDS_PER_DAY = 86_400 SECONDS_PER_MINUTE = 60 NS_PER_MINUTE = 60_000_000_000 """Nanoseconds (`[ns]`) per minute.""" US_PER_MINUTE = 60_000_000 """Microseconds (`[Ξs]`) per minute.""" MS_PER_MINUTE = 60_000 """Milliseconds (`[ms]`) per minute.""" NS_PER_SECOND = 1_000_000_000 """Nanoseconds (`[ns]`) per second (`[s]`).""" US_PER_SECOND = 1_000_000 """Microseconds (`[Ξs]`) per second (`[s]`).""" MS_PER_SECOND = 1_000 """Milliseconds (`[ms]`) per second (`[s]`).""" NS_PER_MICROSECOND = 1_000 """Nanoseconds (`[ns]`) per microsecond (`[Ξs]`).""" NS_PER_MILLISECOND = 1_000_000 """Nanoseconds (`[ns]`) per millisecond (`[ms]`). From [polars](https://github.com/pola-rs/polars/blob/2c7a3e77f0faa37c86a3745db4ef7707ae50c72e/crates/polars-time/src/chunkedarray/duration.rs#L7). """ EPOCH_YEAR = 1970 """See [Unix time](https://en.wikipedia.org/wiki/Unix_time).""" EPOCH = dt.datetime(EPOCH_YEAR, 1, 1).replace(tzinfo=None) """See [Unix time](https://en.wikipedia.org/wiki/Unix_time).""" narwhals-2.11.0/narwhals/_duration.py0000644000000000000000000000610313615410400014507 0ustar00"""Tools for working with the Polars duration string language.""" from __future__ import annotations import datetime as dt import re from typing import TYPE_CHECKING, Literal, cast, get_args if TYPE_CHECKING: from collections.abc import Container, Mapping from typing_extensions import TypeAlias __all__ = ["IntervalUnit"] IntervalUnit: TypeAlias = Literal["ns", "us", "ms", "s", "m", "h", "d", "mo", "q", "y"] """A Polars duration string interval unit. - 'ns': nanosecond. - 'us': microsecond. - 'ms': millisecond. - 's': second. - 'm': minute. - 'h': hour. - 'd': day. - 'mo': month. - 'q': quarter. - 'y': year. """ TimedeltaKwd: TypeAlias = Literal[ "days", "hours", "minutes", "seconds", "milliseconds", "microseconds" ] PATTERN_INTERVAL: re.Pattern[str] = re.compile( r"^(?P-?\d+)(?Pns|us|ms|mo|m|s|h|d|q|y)\Z" ) MONTH_MULTIPLES = frozenset([1, 2, 3, 4, 6, 12]) QUARTER_MULTIPLES = frozenset([1, 2, 4]) UNIT_TO_TIMEDELTA: Mapping[IntervalUnit, TimedeltaKwd] = { "d": "days", "h": "hours", "m": "minutes", "s": "seconds", "ms": "milliseconds", "us": "microseconds", } class Interval: def __init__(self, multiple: int, unit: IntervalUnit, /) -> None: self.multiple: int = multiple self.unit: IntervalUnit = unit def to_timedelta( self, *, unsupported: Container[IntervalUnit] = frozenset(("ns", "mo", "q", "y")) ) -> dt.timedelta: if self.unit in unsupported: # pragma: no cover msg = f"Creating timedelta with {self.unit} unit is not supported." raise NotImplementedError(msg) kwd = UNIT_TO_TIMEDELTA[self.unit] # error: Keywords must be strings (bad mypy) return dt.timedelta(**{kwd: self.multiple}) # type: ignore[misc] @classmethod def parse(cls, every: str) -> Interval: multiple, unit = cls._parse(every) if unit == "mo" and multiple not in MONTH_MULTIPLES: msg = f"Only the following multiples are supported for 'mo' unit: {MONTH_MULTIPLES}.\nGot: {multiple}." raise ValueError(msg) if unit == "q" and multiple not in QUARTER_MULTIPLES: msg = f"Only the following multiples are supported for 'q' unit: {QUARTER_MULTIPLES}.\nGot: {multiple}." raise ValueError(msg) if unit == "y" and multiple != 1: msg = ( f"Only multiple 1 is currently supported for 'y' unit.\nGot: {multiple}." ) raise ValueError(msg) return cls(multiple, unit) @classmethod def parse_no_constraints(cls, every: str) -> Interval: return cls(*cls._parse(every)) @staticmethod def _parse(every: str) -> tuple[int, IntervalUnit]: if match := PATTERN_INTERVAL.match(every): multiple = int(match["multiple"]) unit = cast("IntervalUnit", match["unit"]) return multiple, unit msg = ( f"Invalid `every` string: {every}. Expected string of kind , " f"where 'unit' is one of: {get_args(IntervalUnit)}." ) raise ValueError(msg) narwhals-2.11.0/narwhals/_enum.py0000644000000000000000000000225013615410400013625 0ustar00from __future__ import annotations # ruff: noqa: ARG004 from enum import Enum from typing import Any class NoAutoEnum(Enum): """Enum base class that prohibits the use of enum.auto() for value assignment. This behavior is achieved by overriding the value generation mechanism. Examples: >>> from enum import auto >>> from narwhals._enum import NoAutoEnum >>> >>> class Colors(NoAutoEnum): ... RED = 1 ... GREEN = 2 >>> Colors.RED >>> class ColorsWithAuto(NoAutoEnum): ... RED = 1 ... GREEN = auto() Traceback (most recent call last): ... ValueError: Creating values with `auto()` is not allowed. Please provide a value manually instead. Raises: ValueError: If `auto()` is attempted to be used for any enum member value. """ @staticmethod def _generate_next_value_( name: str, start: int, count: int, last_values: list[Any] ) -> Any: msg = "Creating values with `auto()` is not allowed. Please provide a value manually instead." raise ValueError(msg) __all__ = ["NoAutoEnum"] narwhals-2.11.0/narwhals/_exceptions.py0000644000000000000000000000354613615410400015053 0ustar00from __future__ import annotations from warnings import warn def find_stacklevel() -> int: """Find the first place in the stack that is not inside narwhals. Taken from: https://github.com/pandas-dev/pandas/blob/ab89c53f48df67709a533b6a95ce3d911871a0a8/pandas/util/_exceptions.py#L30-L51 """ import inspect from pathlib import Path import narwhals as nw pkg_dir = str(Path(nw.__file__).parent) # https://stackoverflow.com/questions/17407119/python-inspect-stack-is-slow frame = inspect.currentframe() n = 0 try: while frame: fname = inspect.getfile(frame) if fname.startswith(pkg_dir) or ( (qualname := getattr(frame.f_code, "co_qualname", None)) # ignore @singledispatch wrappers and qualname.startswith("singledispatch.") ): frame = frame.f_back n += 1 else: # pragma: no cover break else: # pragma: no cover pass finally: # https://docs.python.org/3/library/inspect.html # > Though the cycle detector will catch these, destruction of the frames # > (and local variables) can be made deterministic by removing the cycle # > in a finally clause. del frame return n def issue_deprecation_warning(message: str, _version: str) -> None: # pragma: no cover """Issue a deprecation warning. Arguments: message: The message associated with the warning. _version: Narwhals version when the warning was introduced. Just used for internal bookkeeping. """ warn(message=message, category=DeprecationWarning, stacklevel=find_stacklevel()) def issue_warning(message: str, category: type[Warning]) -> None: warn(message=message, category=category, stacklevel=find_stacklevel()) narwhals-2.11.0/narwhals/_expression_parsing.py0000644000000000000000000007763113615410400016622 0ustar00# Utilities for expression parsing # Useful for backends which don't have any concept of expressions, such # and pandas or PyArrow. # ! Any change to this module will trigger the pyspark and pyspark-connect tests in CI from __future__ import annotations from enum import Enum, auto from typing import TYPE_CHECKING, Any, Callable, Literal, cast from narwhals._utils import zip_strict from narwhals.dependencies import is_numpy_array_1d from narwhals.exceptions import ( InvalidIntoExprError, InvalidOperationError, MultiOutputExpressionError, ) if TYPE_CHECKING: from collections.abc import Iterator, Sequence from typing_extensions import Never, TypeIs from narwhals._compliant import CompliantExpr, CompliantFrameT from narwhals._compliant.typing import ( AliasNames, CompliantExprAny, CompliantFrameAny, CompliantNamespaceAny, EvalNames, ) from narwhals.expr import Expr from narwhals.series import Series from narwhals.typing import IntoExpr, NonNestedLiteral, _1DArray def is_expr(obj: Any) -> TypeIs[Expr]: """Check whether `obj` is a Narwhals Expr.""" from narwhals.expr import Expr return isinstance(obj, Expr) def is_series(obj: Any) -> TypeIs[Series[Any]]: """Check whether `obj` is a Narwhals Expr.""" from narwhals.series import Series return isinstance(obj, Series) def combine_evaluate_output_names( *exprs: CompliantExpr[CompliantFrameT, Any], ) -> EvalNames[CompliantFrameT]: # Follow left-hand-rule for naming. E.g. `nw.sum_horizontal(expr1, expr2)` takes the # first name of `expr1`. def evaluate_output_names(df: CompliantFrameT) -> Sequence[str]: return exprs[0]._evaluate_output_names(df)[:1] return evaluate_output_names def combine_alias_output_names(*exprs: CompliantExprAny) -> AliasNames | None: # Follow left-hand-rule for naming. E.g. `nw.sum_horizontal(expr1.alias(alias), expr2)` takes the # aliasing function of `expr1` and apply it to the first output name of `expr1`. if exprs[0]._alias_output_names is None: return None def alias_output_names(names: Sequence[str]) -> Sequence[str]: return exprs[0]._alias_output_names(names)[:1] # type: ignore[misc] return alias_output_names def evaluate_output_names_and_aliases( expr: CompliantExprAny, df: CompliantFrameAny, exclude: Sequence[str] ) -> tuple[Sequence[str], Sequence[str]]: output_names = expr._evaluate_output_names(df) aliases = ( output_names if expr._alias_output_names is None else expr._alias_output_names(output_names) ) if exclude and expr._metadata.expansion_kind.is_multi_unnamed(): output_names, aliases = zip_strict( *[ (x, alias) for x, alias in zip_strict(output_names, aliases) if x not in exclude ] ) return output_names, aliases class ExprKind(Enum): """Describe which kind of expression we are dealing with.""" LITERAL = auto() """e.g. `nw.lit(1)`""" AGGREGATION = auto() """Reduces to a single value, not affected by row order, e.g. `nw.col('a').mean()`""" ORDERABLE_AGGREGATION = auto() """Reduces to a single value, affected by row order, e.g. `nw.col('a').arg_max()`""" ELEMENTWISE = auto() """Preserves length, can operate without context for surrounding rows, e.g. `nw.col('a').abs()`.""" ORDERABLE_WINDOW = auto() """Depends on the rows around it and on their order, e.g. `diff`.""" WINDOW = auto() """Depends on the rows around it and possibly their order, e.g. `rank`.""" FILTRATION = auto() """Changes length, not affected by row order, e.g. `drop_nulls`.""" ORDERABLE_FILTRATION = auto() """Changes length, affected by row order, e.g. `tail`.""" OVER = auto() """Results from calling `.over` on expression.""" COL = auto() """Results from calling `nw.col`.""" NTH = auto() """Results from calling `nw.nth`.""" EXCLUDE = auto() """Results from calling `nw.exclude`.""" ALL = auto() """Results from calling `nw.all`.""" SELECTOR = auto() """Results from creating an expression with a selector.""" WHEN_THEN = auto() """Results from `when/then expression`, possibly followed by `otherwise`.""" SERIES = auto() """Results from converting a Series to Expr.""" @property def is_orderable(self) -> bool: # Any operation which may be affected by `order_by`, such as `cum_sum`, # `diff`, `rank`, `arg_max`, ... return self in { ExprKind.ORDERABLE_WINDOW, ExprKind.WINDOW, ExprKind.ORDERABLE_AGGREGATION, ExprKind.ORDERABLE_FILTRATION, } @property def is_elementwise(self) -> bool: # Any operation which can operate on each row independently # of the rows around it, e.g. `abs(), __add__, sum_horizontal, ...` return self in { ExprKind.ALL, ExprKind.COL, ExprKind.ELEMENTWISE, ExprKind.EXCLUDE, ExprKind.LITERAL, ExprKind.NTH, ExprKind.SELECTOR, ExprKind.SERIES, ExprKind.WHEN_THEN, } def is_scalar_like(obj: CompliantExprAny) -> bool: return obj._metadata.is_scalar_like class ExpansionKind(Enum): """Describe what kind of expansion the expression performs.""" SINGLE = auto() """e.g. `nw.col('a'), nw.sum_horizontal(nw.all())`""" MULTI_NAMED = auto() """e.g. `nw.col('a', 'b')`""" MULTI_UNNAMED = auto() """e.g. `nw.all()`, nw.nth(0, 1)""" def is_multi_unnamed(self) -> bool: return self is ExpansionKind.MULTI_UNNAMED def is_multi_output(self) -> bool: return self in {ExpansionKind.MULTI_NAMED, ExpansionKind.MULTI_UNNAMED} def __and__(self, other: ExpansionKind) -> Literal[ExpansionKind.MULTI_UNNAMED]: if self is ExpansionKind.MULTI_UNNAMED and other is ExpansionKind.MULTI_UNNAMED: # e.g. nw.selectors.all() - nw.selectors.numeric(). return ExpansionKind.MULTI_UNNAMED # Don't attempt anything more complex, keep it simple and raise in the face of ambiguity. msg = f"Unsupported ExpansionKind combination, got {self} and {other}, please report a bug." # pragma: no cover raise AssertionError(msg) # pragma: no cover class ExprNode: """An operation to create or modify an expression. Parameters: kind: ExprKind of operation. name: Name of function, as defined in the compliant protocols. exprs: Expressifiable arguments to function. str_as_lit: Whether to interpret strings as literals when they are present in `exprs`. allow_multi_output: Whether to allow any of `exprs` to be multi-output. kwargs: Other (non-expressifiable) arguments to function. """ def __init__( self, kind: ExprKind, name: str, /, *exprs: IntoExpr | NonNestedLiteral, str_as_lit: bool = False, allow_multi_output: bool = False, **kwargs: Any, ) -> None: self.kind: ExprKind = kind self.name: str = name self.exprs: Sequence[IntoExpr | NonNestedLiteral] = exprs self.kwargs: dict[str, Any] = kwargs self.str_as_lit: bool = str_as_lit self.allow_multi_output: bool = allow_multi_output # Cached methods. self._is_orderable_cached: bool | None = None self._is_elementwise_cached: bool | None = None def __repr__(self) -> str: if self.name == "col": names = ", ".join(str(x) for x in self.kwargs["names"]) return f"col({names})" arg_str = [] expr_repr = ", ".join(str(x) for x in self.exprs) kwargs_repr = ", ".join(f"{key}={value}" for key, value in self.kwargs.items()) if self.exprs: arg_str.append(expr_repr) if self.kwargs: arg_str.append(kwargs_repr) return f"{self.name}({', '.join(arg_str)})" def as_dict(self) -> dict[str, Any]: # pragma: no cover # Just for debugging. return { "kind": self.kind, "name": self.name, "exprs": self.exprs, "kwargs": self.kwargs, "str_as_lit": self.str_as_lit, "allow_multi_output": self.allow_multi_output, } def _with_kwargs(self, **kwargs: Any) -> ExprNode: return self.__class__( self.kind, self.name, *self.exprs, str_as_lit=self.str_as_lit, **kwargs ) def _push_down_over_node_in_place( self, over_node: ExprNode, over_node_without_order_by: ExprNode ) -> None: exprs: list[IntoExpr | NonNestedLiteral] = [] # Note: please keep this as a for-loop (rather than a list-comprehension) # so that pytest-cov highlights any uncovered branches. over_node_order_by = over_node.kwargs["order_by"] over_node_partition_by = over_node.kwargs["partition_by"] for expr in self.exprs: if not is_expr(expr): exprs.append(expr) elif over_node_order_by and any( expr_node.is_orderable() for expr_node in expr._nodes ): exprs.append(expr._with_over_node(over_node)) elif over_node_partition_by and not all( expr_node.is_elementwise() for expr_node in expr._nodes ): exprs.append(expr._with_over_node(over_node_without_order_by)) else: # If there's no `partition_by`, then `over_node_without_order_by` is a no-op. exprs.append(expr) self.exprs = exprs def is_orderable(self) -> bool: if self._is_orderable_cached is None: # Note: don't combine these if/then statements so that pytest-cov shows if # anything is uncovered. if self.kind.is_orderable: # noqa: SIM114 self._is_orderable_cached = True elif any( any(node.is_orderable() for node in expr._nodes) for expr in self.exprs if is_expr(expr) ): self._is_orderable_cached = True else: self._is_orderable_cached = False return self._is_orderable_cached def is_elementwise(self) -> bool: if self._is_elementwise_cached is None: # Note: don't combine these if/then statements so that pytest-cov shows if # anything is uncovered. if not self.kind.is_elementwise: # noqa: SIM114 self._is_elementwise_cached = False elif any( any(not node.is_elementwise() for node in expr._nodes) for expr in self.exprs if is_expr(expr) ): self._is_elementwise_cached = False else: self._is_elementwise_cached = True return self._is_elementwise_cached class ExprMetadata: """Expression metadata. Parameters: expansion_kind: What kind of expansion the expression performs. has_windows: Whether it already contains window functions. is_elementwise: Whether it can operate row-by-row without context of the other rows around it. is_literal: Whether it is just a literal wrapped in an expression. is_scalar_like: Whether it is a literal or an aggregation. n_orderable_ops: The number of order-dependent operations. In the lazy case, this number must be `0` by the time the expression is evaluated. preserves_length: Whether the expression preserves the input length. current_node: The current ExprNode in the linked list. prev: Reference to the previous ExprMetadata in the linked list (None for root). """ __slots__ = ( "current_node", "expansion_kind", "has_windows", "is_elementwise", "is_literal", "is_scalar_like", "n_orderable_ops", "preserves_length", "prev", ) def __init__( self, expansion_kind: ExpansionKind, *, has_windows: bool = False, n_orderable_ops: int = 0, preserves_length: bool = True, is_elementwise: bool = True, is_scalar_like: bool = False, is_literal: bool = False, current_node: ExprNode, prev: ExprMetadata | None = None, ) -> None: if is_literal: assert is_scalar_like # noqa: S101 # debug assertion if is_elementwise: assert preserves_length # noqa: S101 # debug assertion self.expansion_kind: ExpansionKind = expansion_kind self.has_windows: bool = has_windows self.n_orderable_ops: int = n_orderable_ops self.is_elementwise: bool = is_elementwise self.preserves_length: bool = preserves_length self.is_scalar_like: bool = is_scalar_like self.is_literal: bool = is_literal self.current_node: ExprNode = current_node self.prev: ExprMetadata | None = prev def __init_subclass__(cls, /, *args: Any, **kwds: Any) -> Never: # pragma: no cover msg = f"Cannot subclass {cls.__name__!r}" raise TypeError(msg) def __repr__(self) -> str: # pragma: no cover nodes = tuple(reversed(tuple(self.iter_nodes_reversed()))) return ( f"ExprMetadata(\n" f" expansion_kind: {self.expansion_kind},\n" f" has_windows: {self.has_windows},\n" f" n_orderable_ops: {self.n_orderable_ops},\n" f" is_elementwise: {self.is_elementwise},\n" f" preserves_length: {self.preserves_length},\n" f" is_scalar_like: {self.is_scalar_like},\n" f" is_literal: {self.is_literal},\n" f" nodes: {nodes},\n" ")" ) def iter_nodes_reversed(self) -> Iterator[ExprNode]: """Iterate through all nodes from current to root.""" current: ExprMetadata | None = self while current is not None: yield current.current_node current = current.prev @classmethod def from_node( cls, node: ExprNode, *compliant_exprs: CompliantExprAny ) -> ExprMetadata: return KIND_TO_METADATA_CONSTRUCTOR[node.kind](node, *compliant_exprs) def with_node( self, node: ExprNode, compliant_expr: CompliantExprAny, *compliant_expr_args: CompliantExprAny, ) -> ExprMetadata: return KIND_TO_METADATA_UPDATER[node.kind]( self, node, compliant_expr, *compliant_expr_args ) @classmethod def from_aggregation(cls, node: ExprNode) -> ExprMetadata: return cls( ExpansionKind.SINGLE, is_elementwise=False, preserves_length=False, is_scalar_like=True, current_node=node, prev=None, ) @classmethod def from_literal(cls, node: ExprNode) -> ExprMetadata: return cls( ExpansionKind.SINGLE, is_elementwise=False, preserves_length=False, is_literal=True, is_scalar_like=True, current_node=node, prev=None, ) @classmethod def from_series(cls, node: ExprNode) -> ExprMetadata: return cls(ExpansionKind.SINGLE, current_node=node, prev=None) @classmethod def from_col(cls, node: ExprNode) -> ExprMetadata: # e.g. `nw.col('a')`, `nw.nth(0)` return ( cls(ExpansionKind.SINGLE, current_node=node, prev=None) if len(node.kwargs["names"]) == 1 else cls.from_selector_multi_named(node) ) @classmethod def from_nth(cls, node: ExprNode) -> ExprMetadata: return ( cls(ExpansionKind.SINGLE, current_node=node, prev=None) if len(node.kwargs["indices"]) == 1 else cls.from_selector_multi_named(node) ) @classmethod def from_selector_multi_named(cls, node: ExprNode) -> ExprMetadata: # e.g. `nw.col('a', 'b')` return cls(ExpansionKind.MULTI_NAMED, current_node=node, prev=None) @classmethod def from_selector_multi_unnamed(cls, node: ExprNode) -> ExprMetadata: # e.g. `nw.all()` return cls(ExpansionKind.MULTI_UNNAMED, current_node=node, prev=None) @classmethod def from_elementwise( cls, node: ExprNode, *compliant_exprs: CompliantExprAny ) -> ExprMetadata: return combine_metadata( *compliant_exprs, to_single_output=True, current_node=node, prev=None ) @property def is_filtration(self) -> bool: return not self.preserves_length and not self.is_scalar_like def with_aggregation(self, node: ExprNode, _ce: CompliantExprAny) -> ExprMetadata: if self.is_scalar_like: msg = "Can't apply aggregations to scalar-like expressions." raise InvalidOperationError(msg) return ExprMetadata( self.expansion_kind, has_windows=self.has_windows, n_orderable_ops=self.n_orderable_ops, preserves_length=False, is_elementwise=False, is_scalar_like=True, is_literal=False, current_node=node, prev=self, ) def with_orderable_aggregation( self, node: ExprNode, _ce: CompliantExprAny ) -> ExprMetadata: # Deprecated, used only in stable.v1. if self.is_scalar_like: # pragma: no cover msg = "Can't apply aggregations to scalar-like expressions." raise InvalidOperationError(msg) return ExprMetadata( self.expansion_kind, has_windows=self.has_windows, n_orderable_ops=self.n_orderable_ops + 1, preserves_length=False, is_elementwise=False, is_scalar_like=True, is_literal=False, current_node=node, prev=self, ) def with_elementwise( self, node: ExprNode, compliant_expr: CompliantExprAny, *compliant_expr_args: CompliantExprAny, ) -> ExprMetadata: return combine_metadata( compliant_expr, *compliant_expr_args, to_single_output=False, current_node=node, prev=compliant_expr._metadata, ) def with_window(self, node: ExprNode, _ce: CompliantExprAny) -> ExprMetadata: # Window function which may (but doesn't have to) be used with `over(order_by=...)`. if self.is_scalar_like: msg = "Can't apply window (e.g. `rank`) to scalar-like expression." raise InvalidOperationError(msg) return ExprMetadata( self.expansion_kind, has_windows=self.has_windows, # The function isn't order-dependent (but, users can still use `order_by` if they wish!), # so we don't increment `n_orderable_ops`. n_orderable_ops=self.n_orderable_ops, preserves_length=self.preserves_length, is_elementwise=False, is_scalar_like=False, is_literal=False, current_node=node, prev=self, ) def with_orderable_window( self, node: ExprNode, _ce: CompliantExprAny ) -> ExprMetadata: # Window function which must be used with `over(order_by=...)`. if self.is_scalar_like: msg = "Can't apply orderable window (e.g. `diff`, `shift`) to scalar-like expression." raise InvalidOperationError(msg) return ExprMetadata( self.expansion_kind, has_windows=self.has_windows, n_orderable_ops=self.n_orderable_ops + 1, preserves_length=self.preserves_length, is_elementwise=False, is_scalar_like=False, is_literal=False, current_node=node, prev=self, ) def with_ordered_over(self, node: ExprNode, _ce: CompliantExprAny) -> ExprMetadata: if self.has_windows: msg = "Cannot nest `over` statements." raise InvalidOperationError(msg) if self.is_elementwise or self.is_filtration: msg = ( "Cannot use `over` on expressions which are elementwise\n" "(e.g. `abs`) or which change length (e.g. `drop_nulls`)." ) raise InvalidOperationError(msg) n_orderable_ops = self.n_orderable_ops if ( not n_orderable_ops and next(self.op_nodes_reversed()).kind is not ExprKind.WINDOW ): msg = ( "Cannot use `order_by` in `over` on expression which isn't orderable.\n" "If your expression is orderable, then make sure that `over(order_by=...)`\n" "comes immediately after the order-dependent expression.\n\n" "Hint: instead of\n" " - `(nw.col('price').diff() + 1).over(order_by='date')`\n" "write:\n" " + `nw.col('price').diff().over(order_by='date') + 1`\n" ) raise InvalidOperationError(msg) if next(self.op_nodes_reversed()).kind.is_orderable and n_orderable_ops > 0: n_orderable_ops -= 1 return ExprMetadata( self.expansion_kind, has_windows=True, n_orderable_ops=n_orderable_ops, preserves_length=True, is_elementwise=False, is_scalar_like=False, is_literal=False, current_node=node, prev=self, ) def with_partitioned_over( self, node: ExprNode, _ce: CompliantExprAny ) -> ExprMetadata: if self.has_windows: msg = "Cannot nest `over` statements." raise InvalidOperationError(msg) if self.is_elementwise or self.is_filtration: msg = ( "Cannot use `over` on expressions which are elementwise\n" "(e.g. `abs`) or which change length (e.g. `drop_nulls`)." ) raise InvalidOperationError(msg) return ExprMetadata( self.expansion_kind, has_windows=True, n_orderable_ops=self.n_orderable_ops, preserves_length=True, is_elementwise=False, is_scalar_like=False, is_literal=False, current_node=node, prev=self, ) def with_over(self, node: ExprNode, _ce: CompliantExprAny) -> ExprMetadata: if node.kwargs["order_by"]: return self.with_ordered_over(node, _ce) if not node.kwargs["partition_by"]: # pragma: no cover msg = "At least one of `partition_by` or `order_by` must be specified." raise InvalidOperationError(msg) return self.with_partitioned_over(node, _ce) def with_filtration( self, node: ExprNode, *compliant_exprs: CompliantExprAny ) -> ExprMetadata: if self.is_scalar_like: msg = "Can't apply filtration (e.g. `drop_nulls`) to scalar-like expression." raise InvalidOperationError(msg) result_has_windows = any(x._metadata.has_windows for x in compliant_exprs) result_n_orderable_ops = sum(x._metadata.n_orderable_ops for x in compliant_exprs) return ExprMetadata( self.expansion_kind, has_windows=result_has_windows, n_orderable_ops=result_n_orderable_ops, preserves_length=False, is_elementwise=False, is_scalar_like=False, is_literal=False, current_node=node, prev=self, ) def with_orderable_filtration( self, node: ExprNode, _ce: CompliantExprAny ) -> ExprMetadata: if self.is_scalar_like: msg = "Can't apply filtration (e.g. `drop_nulls`) to scalar-like expression." raise InvalidOperationError(msg) return ExprMetadata( self.expansion_kind, has_windows=self.has_windows, n_orderable_ops=self.n_orderable_ops + 1, preserves_length=False, is_elementwise=False, is_scalar_like=False, is_literal=False, current_node=node, prev=self, ) def op_nodes_reversed(self) -> Iterator[ExprNode]: for node in self.iter_nodes_reversed(): if node.name.startswith(("name.", "alias")): # Skip nodes which only do aliasing. continue yield node KIND_TO_METADATA_CONSTRUCTOR: dict[ExprKind, Callable[[ExprNode], ExprMetadata]] = { ExprKind.AGGREGATION: ExprMetadata.from_aggregation, ExprKind.ALL: ExprMetadata.from_selector_multi_unnamed, ExprKind.ELEMENTWISE: ExprMetadata.from_elementwise, ExprKind.EXCLUDE: ExprMetadata.from_selector_multi_unnamed, ExprKind.SERIES: ExprMetadata.from_series, ExprKind.COL: ExprMetadata.from_col, ExprKind.LITERAL: ExprMetadata.from_literal, ExprKind.NTH: ExprMetadata.from_nth, ExprKind.SELECTOR: ExprMetadata.from_selector_multi_unnamed, } KIND_TO_METADATA_UPDATER: dict[ExprKind, Callable[..., ExprMetadata]] = { ExprKind.AGGREGATION: ExprMetadata.with_aggregation, ExprKind.ELEMENTWISE: ExprMetadata.with_elementwise, ExprKind.FILTRATION: ExprMetadata.with_filtration, ExprKind.ORDERABLE_AGGREGATION: ExprMetadata.with_orderable_aggregation, ExprKind.ORDERABLE_FILTRATION: ExprMetadata.with_orderable_filtration, ExprKind.OVER: ExprMetadata.with_over, ExprKind.ORDERABLE_WINDOW: ExprMetadata.with_orderable_window, ExprKind.WINDOW: ExprMetadata.with_window, } def combine_metadata( *compliant_exprs: CompliantExprAny, to_single_output: bool, current_node: ExprNode, prev: ExprMetadata | None, ) -> ExprMetadata: """Combine metadata from `args`. Arguments: compliant_exprs: Expression arguments. to_single_output: Whether the result is always single-output, regardless of the inputs (e.g. `nw.sum_horizontal`). current_node: The current node being added. prev: ExprMetadata of previous node. """ n_filtrations = 0 result_expansion_kind = ExpansionKind.SINGLE result_has_windows = False result_n_orderable_ops = 0 # result preserves length if at least one input does result_preserves_length = False # result is elementwise if all inputs are elementwise result_is_elementwise = True # result is scalar-like if all inputs are scalar-like result_is_scalar_like = True # result is literal if all inputs are literal result_is_literal = True for i, ce in enumerate(compliant_exprs): metadata = ce._metadata assert metadata is not None # noqa: S101 if metadata.expansion_kind.is_multi_output(): expansion_kind = metadata.expansion_kind if not to_single_output: result_expansion_kind = ( result_expansion_kind & expansion_kind if i > 0 else expansion_kind ) result_has_windows |= metadata.has_windows result_n_orderable_ops += metadata.n_orderable_ops result_preserves_length |= metadata.preserves_length result_is_elementwise &= metadata.is_elementwise result_is_scalar_like &= metadata.is_scalar_like result_is_literal &= metadata.is_literal n_filtrations += int(metadata.is_filtration) if n_filtrations > 1: msg = "Length-changing expressions can only be used in isolation, or followed by an aggregation" raise InvalidOperationError(msg) if result_preserves_length and n_filtrations: msg = "Cannot combine length-changing expressions with length-preserving ones or aggregations" raise InvalidOperationError(msg) return ExprMetadata( result_expansion_kind, has_windows=result_has_windows, n_orderable_ops=result_n_orderable_ops, preserves_length=result_preserves_length, is_elementwise=result_is_elementwise, is_scalar_like=result_is_scalar_like, is_literal=result_is_literal, current_node=current_node, prev=prev, ) def check_expressions_preserve_length( *args: CompliantExprAny, function_name: str ) -> None: # Raise if any argument in `args` isn't length-preserving. # For Series input, we don't raise (yet), we let such checks happen later, # as this function works lazily and so can't evaluate lengths. if not all(x._metadata.preserves_length for x in args): msg = f"Expressions which aggregate or change length cannot be passed to '{function_name}'." raise InvalidOperationError(msg) def _parse_into_expr( arg: IntoExpr | NonNestedLiteral | _1DArray, *, str_as_lit: bool = False, backend: Any = None, allow_literal: bool = True, ) -> Expr: from narwhals.functions import col, lit, new_series if isinstance(arg, str) and not str_as_lit: return col(arg) if is_numpy_array_1d(arg): return new_series("", arg, backend=backend)._to_expr() if is_series(arg): return arg._to_expr() if is_expr(arg): return arg if not allow_literal: raise InvalidIntoExprError.from_invalid_type(type(arg)) return lit(arg) def evaluate_into_exprs( *exprs: IntoExpr | NonNestedLiteral | _1DArray, ns: CompliantNamespaceAny, str_as_lit: bool, allow_multi_output: bool, ) -> Iterator[CompliantExprAny]: for expr in exprs: ret = _parse_into_expr( expr, str_as_lit=str_as_lit, backend=ns._implementation )._to_compliant_expr(ns) if not allow_multi_output and ret._metadata.expansion_kind.is_multi_output(): msg = "Multi-output expressions are not allowed in this context." raise MultiOutputExpressionError(msg) yield ret def maybe_broadcast_ces(*compliant_exprs: CompliantExprAny) -> list[CompliantExprAny]: broadcast = any(not is_scalar_like(ce) for ce in compliant_exprs) results: list[CompliantExprAny] = [] for compliant_expr in compliant_exprs: if broadcast and is_scalar_like(compliant_expr): _compliant_expr: CompliantExprAny = compliant_expr.broadcast() # Make sure to preserve metadata. _compliant_expr._opt_metadata = compliant_expr._metadata results.append(_compliant_expr) else: results.append(compliant_expr) return results def evaluate_root_node(node: ExprNode, ns: CompliantNamespaceAny) -> CompliantExprAny: if node.name in {"col", "exclude"}: # There's too much potential for Sequence[str] vs str bugs, so we pass down # `names` positionally rather than as a sequence of strings. ce = getattr(ns, node.name)(*node.kwargs["names"]) ces = [] else: if "." in node.name: module, method = node.name.split(".") func = getattr(getattr(ns, module), method) else: func = getattr(ns, node.name) ces = maybe_broadcast_ces( *evaluate_into_exprs( *node.exprs, ns=ns, str_as_lit=node.str_as_lit, allow_multi_output=node.allow_multi_output, ) ) ce = cast("CompliantExprAny", func(*ces, **node.kwargs)) md = ExprMetadata.from_node(node, *ces) ce._opt_metadata = md return ce def evaluate_node( compliant_expr: CompliantExprAny, node: ExprNode, ns: CompliantNamespaceAny ) -> CompliantExprAny: md: ExprMetadata = compliant_expr._metadata compliant_expr, *compliant_expr_args = maybe_broadcast_ces( compliant_expr, *evaluate_into_exprs( *node.exprs, ns=ns, str_as_lit=node.str_as_lit, allow_multi_output=node.allow_multi_output, ), ) md = md.with_node(node, compliant_expr, *compliant_expr_args) if "." in node.name: accessor, method = node.name.split(".") func = getattr(getattr(compliant_expr, accessor), method) else: func = getattr(compliant_expr, node.name) ret = cast("CompliantExprAny", func(*compliant_expr_args, **node.kwargs)) ret._opt_metadata = md return ret narwhals-2.11.0/narwhals/_namespace.py0000644000000000000000000002061513615410400014622 0ustar00"""Narwhals-level equivalent of `CompliantNamespace`.""" from __future__ import annotations from typing import TYPE_CHECKING, Any, ClassVar, Generic, overload from narwhals._compliant.typing import CompliantNamespaceAny, CompliantNamespaceT_co from narwhals._native import ( NativeAny, NativeArrow, NativeCuDF, NativeDask, NativeDuckDB, NativeIbis, NativeModin, NativePandas, NativePandasLike, NativePolars, NativeSparkLike, NativeUnknown, _CuDFDataFrame, _CuDFSeries, _ModinDataFrame, _ModinSeries, is_native_arrow, is_native_cudf, is_native_dask, is_native_duckdb, is_native_ibis, is_native_modin, is_native_pandas, is_native_polars, is_native_pyspark_connect, is_native_spark_like, is_native_sqlframe, ) from narwhals._utils import Implementation, Version if TYPE_CHECKING: import pandas as pd from typing_extensions import TypeAlias from narwhals._arrow.namespace import ArrowNamespace from narwhals._dask.namespace import DaskNamespace from narwhals._duckdb.namespace import DuckDBNamespace from narwhals._ibis.namespace import IbisNamespace from narwhals._pandas_like.namespace import PandasLikeNamespace from narwhals._polars.namespace import PolarsNamespace from narwhals._spark_like.namespace import SparkLikeNamespace from narwhals._typing import ( Arrow, Backend, Dask, DuckDB, EagerAllowed, Ibis, IntoBackend, PandasLike, Polars, SparkLike, ) EagerAllowedNamespace: TypeAlias = "Namespace[PandasLikeNamespace] | Namespace[ArrowNamespace] | Namespace[PolarsNamespace]" __all__ = ["Namespace"] class Namespace(Generic[CompliantNamespaceT_co]): _compliant_namespace: CompliantNamespaceT_co _version: ClassVar[Version] = Version.MAIN def __init__(self, namespace: CompliantNamespaceT_co, /) -> None: self._compliant_namespace = namespace def __init_subclass__(cls, *args: Any, version: Version, **kwds: Any) -> None: super().__init_subclass__(*args, **kwds) if isinstance(version, Version): cls._version = version else: msg = f"Expected {Version} but got {type(version).__name__!r}" raise TypeError(msg) def __repr__(self) -> str: return f"Namespace[{type(self.compliant).__name__}]" @property def compliant(self) -> CompliantNamespaceT_co: return self._compliant_namespace @property def implementation(self) -> Implementation: return self.compliant._implementation @property def version(self) -> Version: return self._version @overload @classmethod def from_backend(cls, backend: PandasLike, /) -> Namespace[PandasLikeNamespace]: ... @overload @classmethod def from_backend(cls, backend: Polars, /) -> Namespace[PolarsNamespace]: ... @overload @classmethod def from_backend(cls, backend: Arrow, /) -> Namespace[ArrowNamespace]: ... @overload @classmethod def from_backend(cls, backend: SparkLike, /) -> Namespace[SparkLikeNamespace]: ... @overload @classmethod def from_backend(cls, backend: DuckDB, /) -> Namespace[DuckDBNamespace]: ... @overload @classmethod def from_backend(cls, backend: Dask, /) -> Namespace[DaskNamespace]: ... @overload @classmethod def from_backend(cls, backend: Ibis, /) -> Namespace[IbisNamespace]: ... @overload @classmethod def from_backend(cls, backend: EagerAllowed, /) -> EagerAllowedNamespace: ... @overload @classmethod def from_backend( cls, backend: IntoBackend[Backend], / ) -> Namespace[CompliantNamespaceAny]: ... @classmethod def from_backend( cls: type[Namespace[Any]], backend: IntoBackend[Backend], / ) -> Namespace[Any]: """Instantiate from native namespace module, string, or Implementation. Arguments: backend: native namespace module, string, or Implementation. Examples: >>> from narwhals._namespace import Namespace >>> Namespace.from_backend("polars") Namespace[PolarsNamespace] """ impl = Implementation.from_backend(backend) backend_version = impl._backend_version() # noqa: F841 version = cls._version ns: CompliantNamespaceAny if impl.is_pandas_like(): from narwhals._pandas_like.namespace import PandasLikeNamespace ns = PandasLikeNamespace(implementation=impl, version=version) elif impl.is_polars(): from narwhals._polars.namespace import PolarsNamespace ns = PolarsNamespace(version=version) elif impl.is_pyarrow(): from narwhals._arrow.namespace import ArrowNamespace ns = ArrowNamespace(version=version) elif impl.is_spark_like(): from narwhals._spark_like.namespace import SparkLikeNamespace ns = SparkLikeNamespace(implementation=impl, version=version) elif impl.is_duckdb(): from narwhals._duckdb.namespace import DuckDBNamespace ns = DuckDBNamespace(version=version) elif impl.is_dask(): from narwhals._dask.namespace import DaskNamespace ns = DaskNamespace(version=version) elif impl.is_ibis(): from narwhals._ibis.namespace import IbisNamespace ns = IbisNamespace(version=version) else: msg = "Not supported Implementation" # pragma: no cover raise AssertionError(msg) return cls(ns) @overload @classmethod def from_native_object( cls, native: NativePolars, / ) -> Namespace[PolarsNamespace]: ... @overload @classmethod def from_native_object( cls, native: NativePandas, / ) -> Namespace[PandasLikeNamespace[pd.DataFrame, pd.Series[Any]]]: ... @overload @classmethod def from_native_object(cls, native: NativeArrow, /) -> Namespace[ArrowNamespace]: ... @overload @classmethod def from_native_object( cls, native: NativeSparkLike, / ) -> Namespace[SparkLikeNamespace]: ... @overload @classmethod def from_native_object( cls, native: NativeDuckDB, / ) -> Namespace[DuckDBNamespace]: ... @overload @classmethod def from_native_object(cls, native: NativeDask, /) -> Namespace[DaskNamespace]: ... @overload @classmethod def from_native_object(cls, native: NativeIbis, /) -> Namespace[IbisNamespace]: ... @overload @classmethod def from_native_object( cls, native: NativeModin, / ) -> Namespace[PandasLikeNamespace[_ModinDataFrame, _ModinSeries]]: ... @overload @classmethod def from_native_object( cls, native: NativeCuDF, / ) -> Namespace[PandasLikeNamespace[_CuDFDataFrame, _CuDFSeries]]: ... @overload @classmethod def from_native_object( cls, native: NativePandasLike, / ) -> Namespace[PandasLikeNamespace[Any, Any]]: ... @overload @classmethod def from_native_object( cls, native: NativeUnknown, / ) -> Namespace[CompliantNamespaceAny]: ... @classmethod def from_native_object( cls: type[Namespace[Any]], native: NativeAny, / ) -> Namespace[Any]: impl: Backend if is_native_polars(native): impl = Implementation.POLARS elif is_native_pandas(native): impl = Implementation.PANDAS elif is_native_arrow(native): impl = Implementation.PYARROW elif is_native_spark_like(native): impl = ( Implementation.SQLFRAME if is_native_sqlframe(native) else Implementation.PYSPARK_CONNECT if is_native_pyspark_connect(native) else Implementation.PYSPARK ) elif is_native_dask(native): # pragma: no cover impl = Implementation.DASK elif is_native_duckdb(native): impl = Implementation.DUCKDB elif is_native_cudf(native): # pragma: no cover impl = Implementation.CUDF elif is_native_modin(native): # pragma: no cover impl = Implementation.MODIN elif is_native_ibis(native): impl = Implementation.IBIS else: msg = f"Unsupported type: {type(native).__qualname__!r}" raise TypeError(msg) return cls.from_backend(impl) narwhals-2.11.0/narwhals/_native.py0000644000000000000000000003515013615410400014154 0ustar00"""The home for *mostly* [structural] counterparts to [nominal] native types. If you find yourself being yelled at by a typechecker and ended up here - **do not fear!** We have 5 funky flavors, which tackle two different problem spaces. How do we describe [Native types] when ... - ... **wrapping in** a [Narwhals type]? - ... **matching to** an [`Implementation`]? ## Wrapping in a Narwhals type [//]: # (TODO @dangotbanned: Replace `Thing` with a better name) The following examples use the placeholder type `Thing` which represents one of: - `DataFrame`: (Eager) 2D data structure representing data as a table with rows and columns. - `LazyFrame`: (Lazy) Computation graph/query against a DataFrame/database. - `Series`: 1D data structure representing a single column. Our goal is to **wrap** a *partially-unknown* native object **in** a [generic class]: def wrapping_in_df(native: IntoDataFrameT) -> DataFrame[IntoDataFrameT]: ... def wrapping_in_lf(native: IntoLazyFrameT) -> LazyFrame[IntoLazyFrameT]: ... def wrapping_in_ser(native: IntoSeriesT) -> Series[IntoSeriesT]: ... ### (1) `Native` Minimal [`Protocol`]s that are [assignable to] *almost any* supported native type of that group: class NativeThing(Protocol): def something_common(self, *args: Any, **kwargs: Any) -> Any: ... Note: This group is primarily a building block for more useful types. ### (2) `Into` *Publicly* exported [`TypeAlias`]s of **(1)**: IntoThing: TypeAlias = NativeThing **But**, occasionally, there'll be an edge-case which we can spell like: IntoThing: TypeAlias = Union[, NativeThing] Tip: Reach for these when there **isn't a need to preserve** the original native type. ### (3) `IntoT` *Publicly* exported [`TypeVar`]s, bound to **(2)**: IntoThingT = TypeVar("IntoThingT", bound=IntoThing) Important: In most situations, you'll want to use these as they **do preserve** the original native type. Putting it all together, we can now add a *narwhals-level* wrapper: class Thing(Generic[IntoThingT]): def to_native(self) -> IntoThingT: ... ## Matching to an `Implementation` This problem differs as we need to *create* a relationship between *otherwise-unrelated* types. Comparing the problems side-by-side, we can more clearly see this difference: def wrapping_in_df(native: IntoDataFrameT) -> DataFrame[IntoDataFrameT]: ... def matching_to_polars(native: pl.DataFrame) -> Literal[Implementation.POLARS]: ... ### (4) `Native` If we want to describe a set of specific types and **match** them in [`@overload`s], then these the tools we need. For common and easily-installed backends, [`TypeAlias`]s are composed of the native type(s): NativePolars: TypeAlias = pl.DataFrame | pl.LazyFrame | pl.Series Otherwise, we need to define a [`Protocol`] which the native type(s) can **match** against *when* installed: class NativeDask(NativeLazyFrame, Protocol): _partition_type: type[pd.DataFrame] Tip: The goal is to be as minimal as possible, while still being *specific-enough* to **not match** something else. Important: See [ibis#9276 comment] for a more *in-depth* example that doesn't fit here 😄 ### (5) `is_native_` [Type guards] for **(4)**, *similar* to those found in `nw.dependencies`. They differ by checking **all** native types/protocols in a single-call and using ``Native`` aliases. [structural]: https://typing.python.org/en/latest/spec/glossary.html#term-structural [nominal]: https://typing.python.org/en/latest/spec/glossary.html#term-nominal [Native types]: https://narwhals-dev.github.io/narwhals/how_it_works/#polars-and-other-implementations [Narwhals type]: https://narwhals-dev.github.io/narwhals/api-reference/dataframe/ [`Implementation`]: https://narwhals-dev.github.io/narwhals/api-reference/implementation/ [`Protocol`]: https://typing.python.org/en/latest/spec/protocol.html [assignable to]: https://typing.python.org/en/latest/spec/glossary.html#term-assignable [`TypeAlias`]: https://mypy.readthedocs.io/en/stable/kinds_of_types.html#type-aliases [`TypeVar`]: https://mypy.readthedocs.io/en/stable/generics.html#type-variables-with-upper-bounds [generic class]: https://docs.python.org/3/library/typing.html#user-defined-generic-types [`@overload`s]: https://typing.python.org/en/latest/spec/overload.html [ibis#9276 comment]: https://github.com/ibis-project/ibis/issues/9276#issuecomment-3292016818 [Type guards]: https://typing.python.org/en/latest/spec/narrowing.html """ from __future__ import annotations from collections.abc import Callable, Collection, Iterable, Sized from typing import TYPE_CHECKING, Any, Protocol, TypeVar, Union, cast from narwhals.dependencies import ( get_cudf, get_modin, get_pandas, get_polars, get_pyarrow, is_dask_dataframe, is_duckdb_relation, is_ibis_table, is_pyspark_connect_dataframe, is_pyspark_dataframe, is_sqlframe_dataframe, ) if TYPE_CHECKING: import duckdb import pandas as pd import polars as pl import pyarrow as pa from sqlframe.base.dataframe import BaseDataFrame as _BaseDataFrame from typing_extensions import Self, TypeAlias, TypeIs SQLFrameDataFrame = _BaseDataFrame[Any, Any, Any, Any, Any] T = TypeVar("T") _Guard: TypeAlias = "Callable[[Any], TypeIs[T]]" Incomplete: TypeAlias = Any __all__ = [ "IntoDataFrame", "IntoDataFrameT", "IntoFrame", "IntoFrameT", "IntoLazyFrame", "IntoLazyFrameT", "IntoSeries", "IntoSeriesT", "NativeAny", "NativeArrow", "NativeCuDF", "NativeDask", "NativeDataFrame", "NativeDuckDB", "NativeFrame", "NativeIbis", "NativeKnown", "NativeLazyFrame", "NativeModin", "NativePandas", "NativePandasLike", "NativePandasLikeDataFrame", "NativePandasLikeSeries", "NativePolars", "NativePySpark", "NativePySparkConnect", "NativeSQLFrame", "NativeSeries", "NativeSparkLike", "NativeUnknown", "is_native_arrow", "is_native_cudf", "is_native_dask", "is_native_duckdb", "is_native_ibis", "is_native_modin", "is_native_pandas", "is_native_pandas_like", "is_native_polars", "is_native_pyspark", "is_native_pyspark_connect", "is_native_spark_like", "is_native_sqlframe", ] # All dataframes supported by Narwhals have a # `columns` property. Their similarities don't extend # _that_ much further unfortunately... class NativeFrame(Protocol): @property def columns(self) -> Any: ... def join(self, *args: Any, **kwargs: Any) -> Any: ... class NativeDataFrame(Sized, NativeFrame, Protocol): ... class NativeLazyFrame(NativeFrame, Protocol): def explain(self, *args: Any, **kwargs: Any) -> Any: ... class NativeSeries(Sized, Iterable[Any], Protocol): def filter(self, *args: Any, **kwargs: Any) -> Any: ... class _BasePandasLike(Sized, Protocol): index: Any """`mypy` doesn't like the asymmetric `property` setter in `pandas`.""" def __getitem__(self, key: Any, /) -> Any: ... def __mul__(self, other: float | Collection[float] | Self, /) -> Self: ... def __floordiv__(self, other: float | Collection[float] | Self, /) -> Self: ... @property def loc(self) -> Any: ... @property def shape(self) -> tuple[int, ...]: ... def set_axis(self, labels: Any, *, axis: Any = ..., copy: bool = ...) -> Self: ... def copy(self, deep: bool = ...) -> Self: ... # noqa: FBT001 def rename(self, *args: Any, **kwds: Any) -> Self | Incomplete: """`mypy` & `pyright` disagree on overloads. `Incomplete` used to fix [more important issue](https://github.com/narwhals-dev/narwhals/pull/3016#discussion_r2296139744). """ class _BasePandasLikeFrame(NativeDataFrame, _BasePandasLike, Protocol): ... class _BasePandasLikeSeries(NativeSeries, _BasePandasLike, Protocol): def where(self, cond: Any, other: Any = ..., /) -> Self | Incomplete: ... class NativeDask(NativeLazyFrame, Protocol): _partition_type: type[pd.DataFrame] class _CuDFDataFrame(_BasePandasLikeFrame, Protocol): def to_pylibcudf(self, *args: Any, **kwds: Any) -> Any: ... class _CuDFSeries(_BasePandasLikeSeries, Protocol): def to_pylibcudf(self, *args: Any, **kwds: Any) -> Any: ... class NativeIbis(NativeFrame, Protocol): def sql(self, *args: Any, **kwds: Any) -> Any: ... def __pyarrow_result__(self, *args: Any, **kwds: Any) -> Any: ... def __pandas_result__(self, *args: Any, **kwds: Any) -> Any: ... def __polars_result__(self, *args: Any, **kwds: Any) -> Any: ... class _ModinDataFrame(_BasePandasLikeFrame, Protocol): _pandas_class: type[pd.DataFrame] class _ModinSeries(_BasePandasLikeSeries, Protocol): _pandas_class: type[pd.Series[Any]] class _PySparkDataFrame(NativeLazyFrame, Protocol): def dropDuplicatesWithinWatermark(self, *arg: Any, **kwargs: Any) -> Any: ... # noqa: N802 NativePolars: TypeAlias = "pl.DataFrame | pl.LazyFrame | pl.Series" NativeArrow: TypeAlias = "pa.Table | pa.ChunkedArray[Any]" NativeDuckDB: TypeAlias = "duckdb.DuckDBPyRelation" NativePandas: TypeAlias = "pd.DataFrame | pd.Series[Any]" NativeModin: TypeAlias = "_ModinDataFrame | _ModinSeries" NativeCuDF: TypeAlias = "_CuDFDataFrame | _CuDFSeries" NativePandasLikeSeries: TypeAlias = "pd.Series[Any] | _CuDFSeries | _ModinSeries" NativePandasLikeDataFrame: TypeAlias = "pd.DataFrame | _CuDFDataFrame | _ModinDataFrame" NativePandasLike: TypeAlias = "NativePandasLikeDataFrame | NativePandasLikeSeries" NativeSQLFrame: TypeAlias = "_BaseDataFrame[Any, Any, Any, Any, Any]" NativePySpark: TypeAlias = _PySparkDataFrame NativePySparkConnect: TypeAlias = _PySparkDataFrame NativeSparkLike: TypeAlias = "NativeSQLFrame | NativePySpark | NativePySparkConnect" NativeKnown: TypeAlias = "NativePolars | NativeArrow | NativePandasLike | NativeSparkLike | NativeDuckDB | NativeDask | NativeIbis" NativeUnknown: TypeAlias = "NativeDataFrame | NativeSeries | NativeLazyFrame" NativeAny: TypeAlias = "NativeKnown | NativeUnknown" IntoDataFrame: TypeAlias = NativeDataFrame """Anything which can be converted to a Narwhals DataFrame. Use this if your function accepts a narwhalifiable object but doesn't care about its backend. Examples: >>> import narwhals as nw >>> from narwhals.typing import IntoDataFrame >>> def agnostic_shape(df_native: IntoDataFrame) -> tuple[int, int]: ... df = nw.from_native(df_native, eager_only=True) ... return df.shape """ IntoLazyFrame: TypeAlias = Union[NativeLazyFrame, NativeIbis] IntoFrame: TypeAlias = Union[IntoDataFrame, IntoLazyFrame] """Anything which can be converted to a Narwhals DataFrame or LazyFrame. Use this if your function can accept an object which can be converted to either `nw.DataFrame` or `nw.LazyFrame` and it doesn't care about its backend. Examples: >>> import narwhals as nw >>> from narwhals.typing import IntoFrame >>> def agnostic_columns(df_native: IntoFrame) -> list[str]: ... df = nw.from_native(df_native) ... return df.collect_schema().names() """ IntoSeries: TypeAlias = NativeSeries """Anything which can be converted to a Narwhals Series. Use this if your function can accept an object which can be converted to `nw.Series` and it doesn't care about its backend. Examples: >>> from typing import Any >>> import narwhals as nw >>> from narwhals.typing import IntoSeries >>> def agnostic_to_list(s_native: IntoSeries) -> list[Any]: ... s = nw.from_native(s_native) ... return s.to_list() """ IntoFrameT = TypeVar("IntoFrameT", bound=IntoFrame) """TypeVar bound to object convertible to Narwhals DataFrame or Narwhals LazyFrame. Use this if your function accepts an object which is convertible to `nw.DataFrame` or `nw.LazyFrame` and returns an object of the same type. Examples: >>> import narwhals as nw >>> from narwhals.typing import IntoFrameT >>> def agnostic_func(df_native: IntoFrameT) -> IntoFrameT: ... df = nw.from_native(df_native) ... return df.with_columns(c=nw.col("a") + 1).to_native() """ IntoDataFrameT = TypeVar("IntoDataFrameT", bound=IntoDataFrame) """TypeVar bound to object convertible to Narwhals DataFrame. Use this if your function accepts an object which can be converted to `nw.DataFrame` and returns an object of the same class. Examples: >>> import narwhals as nw >>> from narwhals.typing import IntoDataFrameT >>> def agnostic_func(df_native: IntoDataFrameT) -> IntoDataFrameT: ... df = nw.from_native(df_native, eager_only=True) ... return df.with_columns(c=df["a"] + 1).to_native() """ IntoLazyFrameT = TypeVar("IntoLazyFrameT", bound=IntoLazyFrame) IntoSeriesT = TypeVar("IntoSeriesT", bound=IntoSeries) """TypeVar bound to object convertible to Narwhals Series. Use this if your function accepts an object which can be converted to `nw.Series` and returns an object of the same class. Examples: >>> import narwhals as nw >>> from narwhals.typing import IntoSeriesT >>> def agnostic_abs(s_native: IntoSeriesT) -> IntoSeriesT: ... s = nw.from_native(s_native, series_only=True) ... return s.abs().to_native() """ def is_native_polars(obj: Any) -> TypeIs[NativePolars]: return (pl := get_polars()) is not None and isinstance( obj, (pl.DataFrame, pl.Series, pl.LazyFrame) ) def is_native_arrow(obj: Any) -> TypeIs[NativeArrow]: return (pa := get_pyarrow()) is not None and isinstance( obj, (pa.Table, pa.ChunkedArray) ) is_native_dask = cast("_Guard[NativeDask]", is_dask_dataframe) is_native_duckdb: _Guard[NativeDuckDB] = is_duckdb_relation is_native_sqlframe: _Guard[NativeSQLFrame] = is_sqlframe_dataframe is_native_pyspark = cast("_Guard[NativePySpark]", is_pyspark_dataframe) is_native_pyspark_connect = cast( "_Guard[NativePySparkConnect]", is_pyspark_connect_dataframe ) is_native_ibis = cast("_Guard[NativeIbis]", is_ibis_table) def is_native_pandas(obj: Any) -> TypeIs[NativePandas]: return (pd := get_pandas()) is not None and isinstance(obj, (pd.DataFrame, pd.Series)) def is_native_modin(obj: Any) -> TypeIs[NativeModin]: return (mpd := get_modin()) is not None and isinstance( obj, (mpd.DataFrame, mpd.Series) ) def is_native_cudf(obj: Any) -> TypeIs[NativeCuDF]: return (cudf := get_cudf()) is not None and isinstance( obj, (cudf.DataFrame, cudf.Series) ) # pragma: no cover def is_native_pandas_like(obj: Any) -> TypeIs[NativePandasLike]: return is_native_pandas(obj) or is_native_cudf(obj) or is_native_modin(obj) def is_native_spark_like(obj: Any) -> TypeIs[NativeSparkLike]: return ( is_native_sqlframe(obj) or is_native_pyspark(obj) or is_native_pyspark_connect(obj) ) narwhals-2.11.0/narwhals/_translate.py0000644000000000000000000001374013615410400014664 0ustar00"""[Protocols] defining conversion methods between representations. These come in 3 flavors and are [generic] to promote reuse. The following examples use the placeholder types `Narwhal` and `Other`: - `Narwhal`: some class written in `narwhals`. - `Other`: any other class, could be native, compliant, or a builtin. ## `To` When we want to convert or unwrap a `Narwhal` into an `Other`, we provide an **instance** method: ToOtherT_co = TypeVar("ToOtherT_co", covariant=True) class ToOther(Protocol[ToOtherT_co]): def to_other(self, *args: Any, **kwds: Any) -> ToOtherT_co: ... - `*args`, `**kwds` are defined to be *permissive* and allow a wider set of signatures when implementing. - In most cases, they are unused. - But come in handy when adapting an [upstream signature]. - We use a **covariant** `TypeVar`. ## `From` But what if we have `Other` and want to do the reverse? Our `Narwhal` will need to provide a `@classmethod`: FromOtherT_contra = TypeVar("FromOtherT_contra", contravariant=True) class FromOther(Protocol[FromOtherT_contra]): @classmethod def from_other(cls, data: FromOtherT_contra, *args: Any, **kwds: Any) -> Self: ... - `*args`, `**kwds` serve a similar purpose as before, but are much more frequently used. - We've added a **required** [positional-only] parameter `data` which will always be passed `Other`. - This removes the name from the contract of the protocol. - Implementations are free to use something more descriptive for documentation purposes. - We use a **contravariant** `TypeVar`. ## `Convertible` Combining our `to_` and `from_` methods allows us to convert in both directions `Narwhal` <-> `Other`: class OtherConvertible( ToOther[ToOtherT_co], FromOther[FromOtherT_contra], Protocol[ToOtherT_co, FromOtherT_contra], ): ... ## See Also Variance of `TypeVar`(s) can be tricky to wrap your head around. To learn more see [moist], [dry], or [even drier] - depending on how deep you wanna go. [Protocols]: https://typing.python.org/en/latest/spec/protocol.html [generic]: https://typing.python.org/en/latest/spec/generics.html [upstream signature]: https://numpy.org/doc/stable/user/basics.interoperability.html#the-array-method [positional-only]: https://peps.python.org/pep-0570/ [moist]: https://mypy.readthedocs.io/en/stable/generics.html#variance-of-generic-types [dry]: https://typing.python.org/en/latest/spec/generics.html#variance [even drier]: https://en.wikipedia.org/wiki/Covariance_and_contravariance_%28computer_science%29 """ from __future__ import annotations from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Any, Protocol from narwhals._typing_compat import TypeVar if TYPE_CHECKING: import pyarrow as pa from typing_extensions import Self, TypeAlias, TypeIs class ArrowStreamExportable(Protocol): def __arrow_c_stream__(self, requested_schema: object | None = None) -> object: ... ToNumpyT_co = TypeVar("ToNumpyT_co", covariant=True) FromNumpyDT_contra = TypeVar( "FromNumpyDT_contra", contravariant=True, default=ToNumpyT_co ) FromNumpyT_contra = TypeVar("FromNumpyT_contra", contravariant=True) class ToNumpy(Protocol[ToNumpyT_co]): def to_numpy(self, *args: Any, **kwds: Any) -> ToNumpyT_co: ... class FromNumpy(Protocol[FromNumpyT_contra]): @classmethod def from_numpy(cls, data: FromNumpyT_contra, *args: Any, **kwds: Any) -> Self: ... class NumpyConvertible( ToNumpy[ToNumpyT_co], FromNumpy[FromNumpyDT_contra], Protocol[ToNumpyT_co, FromNumpyDT_contra], ): def to_numpy(self, dtype: Any, *, copy: bool | None) -> ToNumpyT_co: ... FromIterableT_contra = TypeVar("FromIterableT_contra", contravariant=True, default=Any) class FromIterable(Protocol[FromIterableT_contra]): @classmethod def from_iterable( cls, data: Iterable[FromIterableT_contra], *args: Any, **kwds: Any ) -> Self: ... ToDictDT_co = TypeVar( "ToDictDT_co", bound=Mapping[str, Any], covariant=True, default="dict[str, Any]" ) FromDictDT_contra = TypeVar( "FromDictDT_contra", bound=Mapping[str, Any], contravariant=True, default=Mapping[str, Any], ) class ToDict(Protocol[ToDictDT_co]): def to_dict(self, *args: Any, **kwds: Any) -> ToDictDT_co: ... class FromDict(Protocol[FromDictDT_contra]): @classmethod def from_dict(cls, data: FromDictDT_contra, *args: Any, **kwds: Any) -> Self: ... class DictConvertible( ToDict[ToDictDT_co], FromDict[FromDictDT_contra], Protocol[ToDictDT_co, FromDictDT_contra], ): ... IntoArrowTable: TypeAlias = "ArrowStreamExportable | pa.Table" """An object supporting the [Arrow PyCapsule Interface], or a native [`pyarrow.Table`]. [Arrow PyCapsule Interface]: https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html#arrowstream-export [`pyarrow.Table`]: https://arrow.apache.org/docs/python/generated/pyarrow.Table.html """ ToArrowT_co = TypeVar("ToArrowT_co", covariant=True) FromArrowDT_contra = TypeVar( "FromArrowDT_contra", contravariant=True, default=IntoArrowTable ) class ToArrow(Protocol[ToArrowT_co]): def to_arrow(self, *args: Any, **kwds: Any) -> ToArrowT_co: ... class FromArrow(Protocol[FromArrowDT_contra]): @classmethod def from_arrow(cls, data: FromArrowDT_contra, *args: Any, **kwds: Any) -> Self: ... class ArrowConvertible( ToArrow[ToArrowT_co], FromArrow[FromArrowDT_contra], Protocol[ToArrowT_co, FromArrowDT_contra], ): ... FromNativeT = TypeVar("FromNativeT") class FromNative(Protocol[FromNativeT]): @classmethod def from_native(cls, data: FromNativeT, *args: Any, **kwds: Any) -> Self: ... @staticmethod def _is_native(obj: FromNativeT | Any, /) -> TypeIs[FromNativeT]: """Return `True` if `obj` can be passed to `from_native`.""" ... ToNarwhalsT_co = TypeVar("ToNarwhalsT_co", covariant=True) class ToNarwhals(Protocol[ToNarwhalsT_co]): def to_narwhals(self) -> ToNarwhalsT_co: """Convert into public representation.""" ... narwhals-2.11.0/narwhals/_typing.py0000644000000000000000000001454113615410400014201 0ustar00from __future__ import annotations from types import ModuleType from typing import TYPE_CHECKING, Literal, Union from narwhals._typing_compat import TypeVar from narwhals._utils import Implementation, _NoDefault if TYPE_CHECKING: from typing_extensions import TypeAlias # `str` aliases _Polars: TypeAlias = Literal["polars"] _Arrow: TypeAlias = Literal["pyarrow"] _Dask: TypeAlias = Literal["dask"] _DuckDB: TypeAlias = Literal["duckdb"] _Pandas: TypeAlias = Literal["pandas"] _Modin: TypeAlias = Literal["modin"] _CuDF: TypeAlias = Literal["cudf"] _PySpark: TypeAlias = Literal["pyspark"] _SQLFrame: TypeAlias = Literal["sqlframe"] _PySparkConnect: TypeAlias = Literal["pyspark[connect]"] _Ibis: TypeAlias = Literal["ibis"] _PandasLike: TypeAlias = Literal[_Pandas, _CuDF, _Modin] _SparkLike: TypeAlias = Literal[_PySpark, _SQLFrame, _PySparkConnect] _EagerOnly: TypeAlias = Literal[_PandasLike, _Arrow] _EagerAllowed: TypeAlias = Literal[_Polars, _EagerOnly] _LazyOnly: TypeAlias = Literal[_SparkLike, _Dask, _DuckDB, _Ibis] _LazyAllowed: TypeAlias = Literal[_Polars, _LazyOnly] # `Implementation` aliases _PandasImpl: TypeAlias = Literal[Implementation.PANDAS] _ModinImpl: TypeAlias = Literal[Implementation.MODIN] _CuDFImpl: TypeAlias = Literal[Implementation.CUDF] _PySparkImpl: TypeAlias = Literal[Implementation.PYSPARK] _SQLFrameImpl: TypeAlias = Literal[Implementation.SQLFRAME] _PySparkConnectImpl: TypeAlias = Literal[Implementation.PYSPARK_CONNECT] _PolarsImpl: TypeAlias = Literal[Implementation.POLARS] _ArrowImpl: TypeAlias = Literal[Implementation.PYARROW] _DaskImpl: TypeAlias = Literal[Implementation.DASK] _DuckDBImpl: TypeAlias = Literal[Implementation.DUCKDB] _IbisImpl: TypeAlias = Literal[Implementation.IBIS] _PandasLikeImpl: TypeAlias = Literal[_PandasImpl, _CuDFImpl, _ModinImpl] _SparkLikeImpl: TypeAlias = Literal[_PySparkImpl, _SQLFrameImpl, _PySparkConnectImpl] _EagerOnlyImpl: TypeAlias = Literal[_PandasLikeImpl, _ArrowImpl] _EagerAllowedImpl: TypeAlias = Literal[_EagerOnlyImpl, _PolarsImpl] # noqa: PYI047 _LazyOnlyImpl: TypeAlias = Literal[_SparkLikeImpl, _DaskImpl, _DuckDBImpl, _IbisImpl] _LazyAllowedImpl: TypeAlias = Literal[_LazyOnlyImpl, _PolarsImpl] # noqa: PYI047 # NOTE: Temporary aliases for gaps in `LazyFrame.collect`, `DataFrame.lazy`, see: # - https://github.com/narwhals-dev/narwhals/pull/2971#discussion_r2277137003 # - https://github.com/narwhals-dev/narwhals/pull/3002#issuecomment-3194267667 _LazyFrameCollectImpl: TypeAlias = Literal[_PandasImpl, _PolarsImpl, _ArrowImpl] # noqa: PYI047 # `str | Implementation` aliases Pandas: TypeAlias = Literal[_Pandas, _PandasImpl] CuDF: TypeAlias = Literal[_CuDF, _CuDFImpl] Modin: TypeAlias = Literal[_Modin, _ModinImpl] PySpark: TypeAlias = Literal[_PySpark, _PySparkImpl] SQLFrame: TypeAlias = Literal[_SQLFrame, _SQLFrameImpl] PySparkConnect: TypeAlias = Literal[_PySparkConnect, _PySparkConnectImpl] Polars: TypeAlias = Literal[_Polars, _PolarsImpl] Arrow: TypeAlias = Literal[_Arrow, _ArrowImpl] Dask: TypeAlias = Literal[_Dask, _DaskImpl] DuckDB: TypeAlias = Literal[_DuckDB, _DuckDBImpl] Ibis: TypeAlias = Literal[_Ibis, _IbisImpl] PandasLike: TypeAlias = Literal[_PandasLike, _PandasLikeImpl] SparkLike: TypeAlias = Literal[_SparkLike, _SparkLikeImpl] EagerOnly: TypeAlias = Literal[PandasLike, Arrow] LazyOnly: TypeAlias = Literal[SparkLike, Dask, DuckDB, Ibis] EagerAllowed: TypeAlias = Literal[EagerOnly, Polars] """A string name or [`narwhals.Implementation`][] of an eager backend. - A string name, one of: `"cudf"`, `"modin"`, `"pandas"`, `"pyarrow"`, `"polars"`. - An Implementation, such as: `Implementation.CUDF`, `Implementation.MODIN`, ... """ LazyAllowed: TypeAlias = Literal[LazyOnly, Polars] """A string name or [`narwhals.Implementation`][] of a lazy backend. - A string name, such as: `"duckdb"`, `"ibis"`, `"dask"`, `"sqlframe"`, ... - An Implementation, such as: `Implementation.POLARS`, `Implementation.PYSPARK`, ... """ BackendName: TypeAlias = Literal[_EagerAllowed, _LazyAllowed] Backend: TypeAlias = Literal[EagerAllowed, LazyAllowed] """A string name or [`narwhals.Implementation`][] of a supported backend (either eager or lazy). - A string name, such as: `"duckdb"`, `"ibis"`, `"pandas"`, `"pyarrow"`, `"polars"`, ... - An Implementation, such as: `Implementation.DASK`, `Implementation.PYSPARK`, ... """ BackendT = TypeVar("BackendT", bound=Backend) IntoBackend: TypeAlias = Union[BackendT, ModuleType] """Anything that can be converted into a [`narwhals.Implementation`][]. `backend` can be specified in three ways. Examples: A string backend name, such as: `"pandas"`, `"pyarrow"`, `"modin"`, `"cudf"` >>> import pandas as pd >>> import narwhals as nw >>> >>> data = {"c": [5, 2], "d": [1, 4]} >>> nw.DataFrame.from_dict(data, backend="pandas") ┌──────────────────┐ |Narwhals DataFrame| |------------------| | c d | | 0 5 1 | | 1 2 4 | └──────────────────┘ An Implementation, such as: `Implementation.POLARS`, `Implementation.DUCKDB`, `Implementation.PYSPARK` >>> import narwhals as nw >>> nw.read_parquet("file.parquet", backend=nw.Implementation.PYARROW) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | pyarrow.Table | | a: int64 | | b: int64 | | ---- | | a: [[1,2]] | | b: [[4,5]] | └──────────────────┘ A python module, such as `dask`, `ibis`, `sqlframe` >>> import numpy as np >>> import polars as pl >>> import narwhals as nw >>> >>> arr = np.arange(5, 10) >>> nw.Series.from_numpy("arr", arr, dtype=nw.Int8, backend=pl) ┌──────────────────┐ | Narwhals Series | |------------------| |shape: (5,) | |Series: 'arr' [i8]| |[ | | 5 | | 6 | | 7 | | 8 | | 9 | |] | └──────────────────┘ """ IntoBackendAny: TypeAlias = IntoBackend[Backend] IntoBackendEager: TypeAlias = IntoBackend[EagerAllowed] IntoBackendLazy: TypeAlias = IntoBackend[LazyAllowed] NoDefault: TypeAlias = Literal[_NoDefault.no_default] narwhals-2.11.0/narwhals/_typing_compat.py0000644000000000000000000000472413615410400015546 0ustar00"""Backward compatibility for newer/less buggy typing features. ## Important Import from here to avoid introducing a runtime dependency on [`typing_extensions`] ## Notes - `TypeVar` defaults - https://typing.python.org/en/latest/spec/generics.html#type-parameter-defaults - https://peps.python.org/pep-0696/ - `@deprecated` - https://docs.python.org/3/library/warnings.html#warnings.deprecated - https://typing.python.org/en/latest/spec/directives.html#deprecated - https://peps.python.org/pep-0702/ [`typing_extensions`]: https://github.com/python/typing_extensions """ from __future__ import annotations # ruff: noqa: ARG001, ANN202, N802 import sys from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from typing import Callable if sys.version_info >= (3, 13): from typing import TypeVar from warnings import deprecated else: from typing_extensions import TypeVar, deprecated if sys.version_info >= (3, 11): from typing import Never, assert_never else: from typing_extensions import Never, assert_never _Fn = TypeVar("_Fn", bound=Callable[..., Any]) else: # pragma: no cover if sys.version_info >= (3, 13): from typing import TypeVar from warnings import deprecated else: from typing import TypeVar as _TypeVar def TypeVar( name: str, *constraints: Any, bound: Any | None = None, covariant: bool = False, contravariant: bool = False, **kwds: Any, ): return _TypeVar( name, *constraints, bound=bound, covariant=covariant, contravariant=contravariant, ) def deprecated(message: str, /) -> Callable[[_Fn], _Fn]: def wrapper(func: _Fn, /) -> _Fn: return func return wrapper _ASSERT_NEVER_REPR_MAX_LENGTH = 100 _BUG_URL = ( "https://github.com/narwhals-dev/narwhals/issues/new?template=bug_report.yml" ) def assert_never(arg: Never, /) -> Never: value = repr(arg) if len(value) > _ASSERT_NEVER_REPR_MAX_LENGTH: value = value[:_ASSERT_NEVER_REPR_MAX_LENGTH] + "..." msg = ( f"Expected code to be unreachable, but got: {value}.\n" f"Please report an issue at {_BUG_URL}" ) raise AssertionError(msg) __all__ = ["TypeVar", "assert_never", "deprecated"] narwhals-2.11.0/narwhals/_utils.py0000644000000000000000000021442413615410400014031 0ustar00from __future__ import annotations import os import re import sys from collections.abc import Collection, Container, Iterable, Iterator, Mapping, Sequence from datetime import timezone from enum import Enum, auto from functools import cache, lru_cache, partial, wraps from importlib.util import find_spec from inspect import getattr_static, getdoc from itertools import chain from operator import attrgetter from secrets import token_hex from typing import ( TYPE_CHECKING, Any, Callable, Final, Generic, Literal, Protocol, TypeVar, Union, cast, overload, ) from narwhals._enum import NoAutoEnum from narwhals._exceptions import issue_deprecation_warning from narwhals._typing_compat import assert_never, deprecated from narwhals.dependencies import ( get_cudf, get_dask_dataframe, get_duckdb, get_ibis, get_modin, get_pandas, get_polars, get_pyarrow, get_pyspark_connect, get_pyspark_sql, get_sqlframe, is_narwhals_series, is_narwhals_series_int, is_numpy_array_1d, is_numpy_array_1d_int, is_pandas_like_dataframe, is_pandas_like_series, ) from narwhals.exceptions import ColumnNotFoundError, DuplicateError, InvalidOperationError if TYPE_CHECKING: from collections.abc import Set # noqa: PYI025 from types import ModuleType import pandas as pd import polars as pl import pyarrow as pa from typing_extensions import ( Concatenate, LiteralString, ParamSpec, Self, TypeAlias, TypeIs, ) from narwhals._compliant import CompliantExprT, CompliantSeriesT, NativeSeriesT_co from narwhals._compliant.any_namespace import NamespaceAccessor from narwhals._compliant.typing import ( Accessor, EvalNames, NativeDataFrameT, NativeLazyFrameT, ) from narwhals._namespace import Namespace from narwhals._native import ( NativeArrow, NativeCuDF, NativeDask, NativeDuckDB, NativeIbis, NativeModin, NativePandas, NativePandasLike, NativePolars, NativePySpark, NativePySparkConnect, NativeSQLFrame, ) from narwhals._translate import ArrowStreamExportable, IntoArrowTable, ToNarwhalsT_co from narwhals._typing import ( Backend, IntoBackend, _ArrowImpl, _CuDFImpl, _DaskImpl, _DuckDBImpl, _EagerAllowedImpl, _IbisImpl, _LazyAllowedImpl, _LazyFrameCollectImpl, _ModinImpl, _PandasImpl, _PandasLikeImpl, _PolarsImpl, _PySparkConnectImpl, _PySparkImpl, _SQLFrameImpl, ) from narwhals.dataframe import DataFrame, LazyFrame from narwhals.dtypes import DType from narwhals.series import Series from narwhals.typing import ( CompliantDataFrame, CompliantLazyFrame, CompliantSeries, DTypes, FileSource, IntoSeriesT, MultiIndexSelector, SingleIndexSelector, SizedMultiIndexSelector, SizeUnit, SupportsNativeNamespace, TimeUnit, _1DArray, _SliceIndex, _SliceName, _SliceNone, ) UnknownBackendName: TypeAlias = str FrameOrSeriesT = TypeVar( "FrameOrSeriesT", bound=Union[LazyFrame[Any], DataFrame[Any], Series[Any]] ) _T1 = TypeVar("_T1") _T2 = TypeVar("_T2") _T3 = TypeVar("_T3") _Fn = TypeVar("_Fn", bound="Callable[..., Any]") P = ParamSpec("P") R = TypeVar("R") R1 = TypeVar("R1") R2 = TypeVar("R2") class _SupportsVersion(Protocol): __version__: str class _SupportsGet(Protocol): # noqa: PYI046 def __get__(self, instance: Any, owner: Any | None = None, /) -> Any: ... class _StoresColumns(Protocol): @property def columns(self) -> Sequence[str]: ... _T = TypeVar("_T") NativeT_co = TypeVar("NativeT_co", covariant=True) CompliantT_co = TypeVar("CompliantT_co", covariant=True) _IntoContext: TypeAlias = "_FullContext | NamespaceAccessor[_FullContext]" _IntoContextT = TypeVar("_IntoContextT", bound=_IntoContext) _Method: TypeAlias = "Callable[Concatenate[_IntoContextT, P], R]" _Constructor: TypeAlias = "Callable[Concatenate[_T, P], R2]" class _StoresNative(Protocol[NativeT_co]): """Provides access to a native object. Native objects have types like: >>> from pandas import Series >>> from pyarrow import Table """ @property def native(self) -> NativeT_co: """Return the native object.""" ... class _StoresCompliant(Protocol[CompliantT_co]): # noqa: PYI046 """Provides access to a compliant object. Compliant objects have types like: >>> from narwhals._pandas_like.series import PandasLikeSeries >>> from narwhals._arrow.dataframe import ArrowDataFrame """ @property def compliant(self) -> CompliantT_co: """Return the compliant object.""" ... class _StoresBackendVersion(Protocol): @property def _backend_version(self) -> tuple[int, ...]: """Version tuple for a native package.""" ... class _StoresVersion(Protocol): _version: Version """Narwhals API version (V1 or MAIN).""" class _StoresImplementation(Protocol): _implementation: Implementation """Implementation of native object (pandas, Polars, PyArrow, ...).""" class _LimitedContext(_StoresImplementation, _StoresVersion, Protocol): """Provides 2 attributes. - `_implementation` - `_version` """ class _FullContext(_StoresImplementation, _StoresBackendVersion, Protocol): """Provides 2 attributes. - `_implementation` - `_backend_version` """ class ValidateBackendVersion(_StoresImplementation, Protocol): """Ensure the target `Implementation` is on a supported version.""" def _validate_backend_version(self) -> None: """Raise if installed version below `nw._utils.MIN_VERSIONS`. **Only use this when moving between backends.** Otherwise, the validation will have taken place already. """ _ = self._implementation._backend_version() class Version(Enum): V1 = auto() V2 = auto() MAIN = auto() @property def namespace(self) -> type[Namespace[Any]]: if self is Version.V1: from narwhals.stable.v1._namespace import Namespace as NamespaceV1 return NamespaceV1 if self is Version.V2: from narwhals.stable.v2._namespace import Namespace as NamespaceV2 return NamespaceV2 from narwhals._namespace import Namespace return Namespace @property def dtypes(self) -> DTypes: if self is Version.V1: from narwhals.stable.v1 import dtypes as dtypes_v1 return dtypes_v1 if self is Version.V2: from narwhals.stable.v2 import dtypes as dtypes_v2 return dtypes_v2 from narwhals import dtypes return dtypes @property def dataframe(self) -> type[DataFrame[Any]]: if self is Version.V1: from narwhals.stable.v1 import DataFrame as DataFrameV1 return DataFrameV1 if self is Version.V2: from narwhals.stable.v2 import DataFrame as DataFrameV2 return DataFrameV2 from narwhals.dataframe import DataFrame return DataFrame @property def lazyframe(self) -> type[LazyFrame[Any]]: if self is Version.V1: from narwhals.stable.v1 import LazyFrame as LazyFrameV1 return LazyFrameV1 if self is Version.V2: from narwhals.stable.v2 import LazyFrame as LazyFrameV2 return LazyFrameV2 from narwhals.dataframe import LazyFrame return LazyFrame @property def series(self) -> type[Series[Any]]: if self is Version.V1: from narwhals.stable.v1 import Series as SeriesV1 return SeriesV1 if self is Version.V2: from narwhals.stable.v2 import Series as SeriesV2 return SeriesV2 from narwhals.series import Series return Series class Implementation(NoAutoEnum): """Implementation of native object (pandas, Polars, PyArrow, ...).""" PANDAS = "pandas" """pandas implementation.""" MODIN = "modin" """Modin implementation.""" CUDF = "cudf" """cuDF implementation.""" PYARROW = "pyarrow" """PyArrow implementation.""" PYSPARK = "pyspark" """PySpark implementation.""" POLARS = "polars" """Polars implementation.""" DASK = "dask" """Dask implementation.""" DUCKDB = "duckdb" """DuckDB implementation.""" IBIS = "ibis" """Ibis implementation.""" SQLFRAME = "sqlframe" """SQLFrame implementation.""" PYSPARK_CONNECT = "pyspark[connect]" """PySpark Connect implementation.""" UNKNOWN = "unknown" """Unknown implementation.""" def __str__(self) -> str: return str(self.value) @classmethod def from_native_namespace( cls: type[Self], native_namespace: ModuleType ) -> Implementation: # pragma: no cover """Instantiate Implementation object from a native namespace module. Arguments: native_namespace: Native namespace. """ mapping = { get_pandas(): Implementation.PANDAS, get_modin(): Implementation.MODIN, get_cudf(): Implementation.CUDF, get_pyarrow(): Implementation.PYARROW, get_pyspark_sql(): Implementation.PYSPARK, get_polars(): Implementation.POLARS, get_dask_dataframe(): Implementation.DASK, get_duckdb(): Implementation.DUCKDB, get_ibis(): Implementation.IBIS, get_sqlframe(): Implementation.SQLFRAME, get_pyspark_connect(): Implementation.PYSPARK_CONNECT, } return mapping.get(native_namespace, Implementation.UNKNOWN) @classmethod def from_string(cls: type[Self], backend_name: str) -> Implementation: """Instantiate Implementation object from a native namespace module. Arguments: backend_name: Name of backend, expressed as string. """ try: return cls(backend_name) except ValueError: return Implementation.UNKNOWN @classmethod def from_backend( cls: type[Self], backend: IntoBackend[Backend] | UnknownBackendName ) -> Implementation: """Instantiate from native namespace module, string, or Implementation. Arguments: backend: Backend to instantiate Implementation from. """ return ( cls.from_string(backend) if isinstance(backend, str) else backend if isinstance(backend, Implementation) else cls.from_native_namespace(backend) ) def to_native_namespace(self) -> ModuleType: """Return the native namespace module corresponding to Implementation.""" if self is Implementation.UNKNOWN: msg = "Cannot return native namespace from UNKNOWN Implementation" raise AssertionError(msg) self._backend_version() module_name = _IMPLEMENTATION_TO_MODULE_NAME.get(self, self.value) return _import_native_namespace(module_name) def is_pandas(self) -> bool: """Return whether implementation is pandas. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2, 3]}) >>> df = nw.from_native(df_native) >>> df.implementation.is_pandas() True """ return self is Implementation.PANDAS def is_pandas_like(self) -> bool: """Return whether implementation is pandas, Modin, or cuDF. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2, 3]}) >>> df = nw.from_native(df_native) >>> df.implementation.is_pandas_like() True """ return self in {Implementation.PANDAS, Implementation.MODIN, Implementation.CUDF} def is_spark_like(self) -> bool: """Return whether implementation is pyspark or sqlframe. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2, 3]}) >>> df = nw.from_native(df_native) >>> df.implementation.is_spark_like() False """ return self in { Implementation.PYSPARK, Implementation.SQLFRAME, Implementation.PYSPARK_CONNECT, } def is_polars(self) -> bool: """Return whether implementation is Polars. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"a": [1, 2, 3]}) >>> df = nw.from_native(df_native) >>> df.implementation.is_polars() True """ return self is Implementation.POLARS def is_cudf(self) -> bool: """Return whether implementation is cuDF. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"a": [1, 2, 3]}) >>> df = nw.from_native(df_native) >>> df.implementation.is_cudf() False """ return self is Implementation.CUDF # pragma: no cover def is_modin(self) -> bool: """Return whether implementation is Modin. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"a": [1, 2, 3]}) >>> df = nw.from_native(df_native) >>> df.implementation.is_modin() False """ return self is Implementation.MODIN # pragma: no cover def is_pyspark(self) -> bool: """Return whether implementation is PySpark. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"a": [1, 2, 3]}) >>> df = nw.from_native(df_native) >>> df.implementation.is_pyspark() False """ return self is Implementation.PYSPARK # pragma: no cover def is_pyspark_connect(self) -> bool: """Return whether implementation is PySpark. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"a": [1, 2, 3]}) >>> df = nw.from_native(df_native) >>> df.implementation.is_pyspark_connect() False """ return self is Implementation.PYSPARK_CONNECT # pragma: no cover def is_pyarrow(self) -> bool: """Return whether implementation is PyArrow. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"a": [1, 2, 3]}) >>> df = nw.from_native(df_native) >>> df.implementation.is_pyarrow() False """ return self is Implementation.PYARROW # pragma: no cover def is_dask(self) -> bool: """Return whether implementation is Dask. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"a": [1, 2, 3]}) >>> df = nw.from_native(df_native) >>> df.implementation.is_dask() False """ return self is Implementation.DASK # pragma: no cover def is_duckdb(self) -> bool: """Return whether implementation is DuckDB. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"a": [1, 2, 3]}) >>> df = nw.from_native(df_native) >>> df.implementation.is_duckdb() False """ return self is Implementation.DUCKDB # pragma: no cover def is_ibis(self) -> bool: """Return whether implementation is Ibis. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"a": [1, 2, 3]}) >>> df = nw.from_native(df_native) >>> df.implementation.is_ibis() False """ return self is Implementation.IBIS # pragma: no cover def is_sqlframe(self) -> bool: """Return whether implementation is SQLFrame. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"a": [1, 2, 3]}) >>> df = nw.from_native(df_native) >>> df.implementation.is_sqlframe() False """ return self is Implementation.SQLFRAME # pragma: no cover def _backend_version(self) -> tuple[int, ...]: """Returns backend version.""" return backend_version(self) MIN_VERSIONS: Mapping[Implementation, tuple[int, ...]] = { Implementation.PANDAS: (1, 1, 3), Implementation.MODIN: (0, 8, 2), Implementation.CUDF: (24, 10), Implementation.PYARROW: (13,), Implementation.PYSPARK: (3, 5), Implementation.PYSPARK_CONNECT: (3, 5), Implementation.POLARS: (0, 20, 4), Implementation.DASK: (2024, 8), Implementation.DUCKDB: (1, 1), Implementation.IBIS: (6,), Implementation.SQLFRAME: (3, 22, 0), } _IMPLEMENTATION_TO_MODULE_NAME: Mapping[Implementation, str] = { Implementation.DASK: "dask.dataframe", Implementation.MODIN: "modin.pandas", Implementation.PYSPARK: "pyspark.sql", Implementation.PYSPARK_CONNECT: "pyspark.sql.connect", } """Stores non default mapping from Implementation to module name""" @lru_cache(maxsize=16) def _import_native_namespace(module_name: str) -> ModuleType: from importlib import import_module return import_module(module_name) # NOTE: We can safely use an unbounded cache, the size is constrained by `len(Implementation._member_names_)` # Faster than `lru_cache` # https://docs.python.org/3/library/functools.html#functools.cache @cache def backend_version(implementation: Implementation, /) -> tuple[int, ...]: if not isinstance(implementation, Implementation): assert_never(implementation) if implementation is Implementation.UNKNOWN: return (0, 0, 0) into_version: ModuleType | str impl = implementation module_name = _IMPLEMENTATION_TO_MODULE_NAME.get(impl, impl.value) native_namespace = _import_native_namespace(module_name) if impl.is_sqlframe(): import sqlframe._version into_version = sqlframe._version elif impl.is_pyspark() or impl.is_pyspark_connect(): # pragma: no cover import pyspark # ignore-banned-import into_version = pyspark elif impl.is_dask(): import dask # ignore-banned-import into_version = dask else: into_version = native_namespace version = parse_version(into_version) if version < (min_version := MIN_VERSIONS[impl]): msg = f"Minimum version of {impl} supported by Narwhals is {min_version}, found: {version}" raise ValueError(msg) return version def flatten(args: Any) -> list[Any]: return list(args[0] if (len(args) == 1 and _is_iterable(args[0])) else args) def tupleify(arg: Any) -> Any: if not isinstance(arg, (list, tuple)): # pragma: no cover return (arg,) return arg def _is_iterable(arg: Any | Iterable[Any]) -> bool: from narwhals.series import Series if ( (pd := get_pandas()) is not None and isinstance(arg, (pd.Series, pd.DataFrame)) ) or ( (pl := get_polars()) is not None and isinstance(arg, (pl.Series, pl.Expr, pl.DataFrame, pl.LazyFrame)) ): # Non-exhaustive check for common potential mistakes. msg = ( f"Expected Narwhals class or scalar, got: {qualified_type_name(arg)!r}.\n\n" "Hint: Perhaps you\n" "- forgot a `nw.from_native` somewhere?\n" "- used `pl.col` instead of `nw.col`?" ) raise TypeError(msg) return isinstance(arg, Iterable) and not isinstance(arg, (str, bytes, Series)) def is_iterator(val: Iterable[_T] | Any) -> TypeIs[Iterator[_T]]: return isinstance(val, Iterator) def parse_version(version: str | ModuleType | _SupportsVersion) -> tuple[int, ...]: """Simple version parser; split into a tuple of ints for comparison. Arguments: version: Version string, or object with one, to parse. """ # lifted from Polars # [marco]: Take care of DuckDB pre-releases which end with e.g. `-dev4108` # and pandas pre-releases which end with e.g. .dev0+618.gb552dc95c9 version_str = version if isinstance(version, str) else version.__version__ version_str = re.sub(r"(\D?dev.*$)", "", version_str) return tuple(int(re.sub(r"\D", "", v)) for v in version_str.split(".")) @overload def isinstance_or_issubclass( obj_or_cls: type, cls_or_tuple: type[_T] ) -> TypeIs[type[_T]]: ... @overload def isinstance_or_issubclass( obj_or_cls: object | type, cls_or_tuple: type[_T] ) -> TypeIs[_T | type[_T]]: ... @overload def isinstance_or_issubclass( obj_or_cls: type, cls_or_tuple: tuple[type[_T1], type[_T2]] ) -> TypeIs[type[_T1 | _T2]]: ... @overload def isinstance_or_issubclass( obj_or_cls: object | type, cls_or_tuple: tuple[type[_T1], type[_T2]] ) -> TypeIs[_T1 | _T2 | type[_T1 | _T2]]: ... @overload def isinstance_or_issubclass( obj_or_cls: type, cls_or_tuple: tuple[type[_T1], type[_T2], type[_T3]] ) -> TypeIs[type[_T1 | _T2 | _T3]]: ... @overload def isinstance_or_issubclass( obj_or_cls: object | type, cls_or_tuple: tuple[type[_T1], type[_T2], type[_T3]] ) -> TypeIs[_T1 | _T2 | _T3 | type[_T1 | _T2 | _T3]]: ... @overload def isinstance_or_issubclass( obj_or_cls: Any, cls_or_tuple: tuple[type, ...] ) -> TypeIs[Any]: ... def isinstance_or_issubclass(obj_or_cls: Any, cls_or_tuple: Any) -> bool: from narwhals.dtypes import DType if isinstance(obj_or_cls, DType): return isinstance(obj_or_cls, cls_or_tuple) return isinstance(obj_or_cls, cls_or_tuple) or ( isinstance(obj_or_cls, type) and issubclass(obj_or_cls, cls_or_tuple) ) def validate_laziness(items: Iterable[Any]) -> None: from narwhals.dataframe import DataFrame, LazyFrame if all(isinstance(item, DataFrame) for item in items) or ( all(isinstance(item, LazyFrame) for item in items) ): return msg = f"The items to concatenate should either all be eager, or all lazy, got: {[type(item) for item in items]}" raise TypeError(msg) def maybe_align_index( lhs: FrameOrSeriesT, rhs: Series[Any] | DataFrame[Any] | LazyFrame[Any] ) -> FrameOrSeriesT: """Align `lhs` to the Index of `rhs`, if they're both pandas-like. Arguments: lhs: Dataframe or Series. rhs: Dataframe or Series to align with. Notes: This is only really intended for backwards-compatibility purposes, for example if your library already aligns indices for users. If you're designing a new library, we highly encourage you to not rely on the Index. For non-pandas-like inputs, this only checks that `lhs` and `rhs` are the same length. Examples: >>> import pandas as pd >>> import polars as pl >>> import narwhals as nw >>> df_pd = pd.DataFrame({"a": [1, 2]}, index=[3, 4]) >>> s_pd = pd.Series([6, 7], index=[4, 3]) >>> df = nw.from_native(df_pd) >>> s = nw.from_native(s_pd, series_only=True) >>> nw.to_native(nw.maybe_align_index(df, s)) a 4 2 3 1 """ from narwhals._pandas_like.dataframe import PandasLikeDataFrame from narwhals._pandas_like.series import PandasLikeSeries def _validate_index(index: Any) -> None: if not index.is_unique: msg = "given index doesn't have a unique index" raise ValueError(msg) lhs_any = cast("Any", lhs) rhs_any = cast("Any", rhs) if isinstance( getattr(lhs_any, "_compliant_frame", None), PandasLikeDataFrame ) and isinstance(getattr(rhs_any, "_compliant_frame", None), PandasLikeDataFrame): _validate_index(lhs_any._compliant_frame.native.index) _validate_index(rhs_any._compliant_frame.native.index) return lhs_any._with_compliant( lhs_any._compliant_frame._with_native( lhs_any._compliant_frame.native.loc[rhs_any._compliant_frame.native.index] ) ) if isinstance( getattr(lhs_any, "_compliant_frame", None), PandasLikeDataFrame ) and isinstance(getattr(rhs_any, "_compliant_series", None), PandasLikeSeries): _validate_index(lhs_any._compliant_frame.native.index) _validate_index(rhs_any._compliant_series.native.index) return lhs_any._with_compliant( lhs_any._compliant_frame._with_native( lhs_any._compliant_frame.native.loc[ rhs_any._compliant_series.native.index ] ) ) if isinstance( getattr(lhs_any, "_compliant_series", None), PandasLikeSeries ) and isinstance(getattr(rhs_any, "_compliant_frame", None), PandasLikeDataFrame): _validate_index(lhs_any._compliant_series.native.index) _validate_index(rhs_any._compliant_frame.native.index) return lhs_any._with_compliant( lhs_any._compliant_series._with_native( lhs_any._compliant_series.native.loc[ rhs_any._compliant_frame.native.index ] ) ) if isinstance( getattr(lhs_any, "_compliant_series", None), PandasLikeSeries ) and isinstance(getattr(rhs_any, "_compliant_series", None), PandasLikeSeries): _validate_index(lhs_any._compliant_series.native.index) _validate_index(rhs_any._compliant_series.native.index) return lhs_any._with_compliant( lhs_any._compliant_series._with_native( lhs_any._compliant_series.native.loc[ rhs_any._compliant_series.native.index ] ) ) if len(lhs_any) != len(rhs_any): msg = f"Expected `lhs` and `rhs` to have the same length, got {len(lhs_any)} and {len(rhs_any)}" raise ValueError(msg) return lhs def maybe_get_index(obj: DataFrame[Any] | LazyFrame[Any] | Series[Any]) -> Any | None: """Get the index of a DataFrame or a Series, if it's pandas-like. Arguments: obj: Dataframe or Series. Notes: This is only really intended for backwards-compatibility purposes, for example if your library already aligns indices for users. If you're designing a new library, we highly encourage you to not rely on the Index. For non-pandas-like inputs, this returns `None`. Examples: >>> import pandas as pd >>> import polars as pl >>> import narwhals as nw >>> df_pd = pd.DataFrame({"a": [1, 2], "b": [4, 5]}) >>> df = nw.from_native(df_pd) >>> nw.maybe_get_index(df) RangeIndex(start=0, stop=2, step=1) >>> series_pd = pd.Series([1, 2]) >>> series = nw.from_native(series_pd, series_only=True) >>> nw.maybe_get_index(series) RangeIndex(start=0, stop=2, step=1) """ obj_any = cast("Any", obj) native_obj = obj_any.to_native() if is_pandas_like_dataframe(native_obj) or is_pandas_like_series(native_obj): return native_obj.index return None def maybe_set_index( obj: FrameOrSeriesT, column_names: str | list[str] | None = None, *, index: Series[IntoSeriesT] | list[Series[IntoSeriesT]] | None = None, ) -> FrameOrSeriesT: """Set the index of a DataFrame or a Series, if it's pandas-like. Arguments: obj: object for which maybe set the index (can be either a Narwhals `DataFrame` or `Series`). column_names: name or list of names of the columns to set as index. For dataframes, only one of `column_names` and `index` can be specified but not both. If `column_names` is passed and `df` is a Series, then a `ValueError` is raised. index: series or list of series to set as index. Raises: ValueError: If one of the following conditions happens - none of `column_names` and `index` are provided - both `column_names` and `index` are provided - `column_names` is provided and `df` is a Series Notes: This is only really intended for backwards-compatibility purposes, for example if your library already aligns indices for users. If you're designing a new library, we highly encourage you to not rely on the Index. For non-pandas-like inputs, this is a no-op. Examples: >>> import pandas as pd >>> import polars as pl >>> import narwhals as nw >>> df_pd = pd.DataFrame({"a": [1, 2], "b": [4, 5]}) >>> df = nw.from_native(df_pd) >>> nw.to_native(nw.maybe_set_index(df, "b")) # doctest: +NORMALIZE_WHITESPACE a b 4 1 5 2 """ from narwhals.translate import to_native df_any = cast("Any", obj) native_obj = df_any.to_native() if column_names is not None and index is not None: msg = "Only one of `column_names` or `index` should be provided" raise ValueError(msg) if not column_names and index is None: msg = "Either `column_names` or `index` should be provided" raise ValueError(msg) if index is not None: keys = ( [to_native(idx, pass_through=True) for idx in index] if _is_iterable(index) else to_native(index, pass_through=True) ) else: keys = column_names if is_pandas_like_dataframe(native_obj): return df_any._with_compliant( df_any._compliant_frame._with_native(native_obj.set_index(keys)) ) if is_pandas_like_series(native_obj): from narwhals._pandas_like.utils import set_index if column_names: msg = "Cannot set index using column names on a Series" raise ValueError(msg) native_obj = set_index( native_obj, keys, implementation=obj._compliant_series._implementation, # type: ignore[union-attr] ) return df_any._with_compliant(df_any._compliant_series._with_native(native_obj)) return df_any def maybe_reset_index(obj: FrameOrSeriesT) -> FrameOrSeriesT: """Reset the index to the default integer index of a DataFrame or a Series, if it's pandas-like. Arguments: obj: Dataframe or Series. Notes: This is only really intended for backwards-compatibility purposes, for example if your library already resets the index for users. If you're designing a new library, we highly encourage you to not rely on the Index. For non-pandas-like inputs, this is a no-op. Examples: >>> import pandas as pd >>> import polars as pl >>> import narwhals as nw >>> df_pd = pd.DataFrame({"a": [1, 2], "b": [4, 5]}, index=([6, 7])) >>> df = nw.from_native(df_pd) >>> nw.to_native(nw.maybe_reset_index(df)) a b 0 1 4 1 2 5 >>> series_pd = pd.Series([1, 2]) >>> series = nw.from_native(series_pd, series_only=True) >>> nw.maybe_get_index(series) RangeIndex(start=0, stop=2, step=1) """ obj_any = cast("Any", obj) native_obj = obj_any.to_native() if is_pandas_like_dataframe(native_obj): native_namespace = obj_any.__native_namespace__() if _has_default_index(native_obj, native_namespace): return obj_any return obj_any._with_compliant( obj_any._compliant_frame._with_native(native_obj.reset_index(drop=True)) ) if is_pandas_like_series(native_obj): native_namespace = obj_any.__native_namespace__() if _has_default_index(native_obj, native_namespace): return obj_any return obj_any._with_compliant( obj_any._compliant_series._with_native(native_obj.reset_index(drop=True)) ) return obj_any if TYPE_CHECKING: zip_strict = partial(zip, strict=True) else: import sys if sys.version_info >= (3, 10): zip_strict = partial(zip, strict=True) else: # pragma: no cover # https://stackoverflow.com/questions/32954486/zip-iterators-asserting-for-equal-length-in-python/69485272#69485272 def zip_strict(*iterables: Iterable[Any]) -> Iterable[tuple[Any, ...]]: # For trivial cases, use pure zip. if len(iterables) < 2: return zip(*iterables) # Tail for the first iterable first_stopped = False def first_tail() -> Any: nonlocal first_stopped first_stopped = True return yield # Tail for the zip def zip_tail() -> Any: if not first_stopped: # pragma: no cover msg = "zip_strict: first iterable is longer" raise ValueError(msg) for _ in chain.from_iterable(rest): # pragma: no cover msg = "zip_strict: first iterable is shorter" raise ValueError(msg) yield # Put the pieces together iterables_it = iter(iterables) first = chain(next(iterables_it), first_tail()) rest = list(map(iter, iterables_it)) return chain(zip(first, *rest), zip_tail()) def _is_range_index(obj: Any, native_namespace: Any) -> TypeIs[pd.RangeIndex]: return isinstance(obj, native_namespace.RangeIndex) def _has_default_index( native_frame_or_series: pd.Series[Any] | pd.DataFrame, native_namespace: Any ) -> bool: index = native_frame_or_series.index return ( _is_range_index(index, native_namespace) and index.start == 0 and index.stop == len(index) and index.step == 1 ) def maybe_convert_dtypes( obj: FrameOrSeriesT, *args: bool, **kwargs: bool | str ) -> FrameOrSeriesT: """Convert columns or series to the best possible dtypes using dtypes supporting ``pd.NA``, if df is pandas-like. Arguments: obj: DataFrame or Series. *args: Additional arguments which gets passed through. **kwargs: Additional arguments which gets passed through. Notes: For non-pandas-like inputs, this is a no-op. Also, `args` and `kwargs` just get passed down to the underlying library as-is. Examples: >>> import pandas as pd >>> import polars as pl >>> import narwhals as nw >>> import numpy as np >>> df_pd = pd.DataFrame( ... { ... "a": pd.Series([1, 2, 3], dtype=np.dtype("int32")), ... "b": pd.Series([True, False, np.nan], dtype=np.dtype("O")), ... } ... ) >>> df = nw.from_native(df_pd) >>> nw.to_native( ... nw.maybe_convert_dtypes(df) ... ).dtypes # doctest: +NORMALIZE_WHITESPACE a Int32 b boolean dtype: object """ if not obj.implementation.is_pandas_like(): return obj result = obj._with_compliant( obj._compliant._with_native(obj.to_native().convert_dtypes(*args, **kwargs)) ) return cast("FrameOrSeriesT", result) def scale_bytes(sz: int, unit: SizeUnit) -> int | float: """Scale size in bytes to other size units (eg: "kb", "mb", "gb", "tb"). Arguments: sz: original size in bytes unit: size unit to convert into """ if unit in {"b", "bytes"}: return sz if unit in {"kb", "kilobytes"}: return sz / 1024 if unit in {"mb", "megabytes"}: return sz / 1024**2 if unit in {"gb", "gigabytes"}: return sz / 1024**3 if unit in {"tb", "terabytes"}: return sz / 1024**4 msg = f"`unit` must be one of {{'b', 'kb', 'mb', 'gb', 'tb'}}, got {unit!r}" raise ValueError(msg) def is_ordered_categorical(series: Series[Any]) -> bool: """Return whether indices of categories are semantically meaningful. This is a convenience function to accessing what would otherwise be the `is_ordered` property from the DataFrame Interchange Protocol, see https://data-apis.org/dataframe-protocol/latest/API.html. - For Polars: - Enums are always ordered. - Categoricals are ordered if `dtype.ordering == "physical"`. - For pandas-like APIs: - Categoricals are ordered if `dtype.cat.ordered == True`. - For PyArrow table: - Categoricals are ordered if `dtype.type.ordered == True`. Arguments: series: Input Series. Examples: >>> import narwhals as nw >>> import pandas as pd >>> import polars as pl >>> data = ["x", "y"] >>> s_pd = pd.Series(data, dtype=pd.CategoricalDtype(ordered=True)) >>> s_pl = pl.Series(data, dtype=pl.Categorical(ordering="lexical")) Let's define a library-agnostic function: >>> @nw.narwhalify ... def func(s): ... return nw.is_ordered_categorical(s) Then, we can pass any supported library to `func`: >>> func(s_pd) True >>> func(s_pl) False """ from narwhals._interchange.series import InterchangeSeries dtypes = series._compliant_series._version.dtypes compliant = series._compliant_series # If it doesn't match any branches, let's just play it safe and return False. result: bool = False if isinstance(compliant, InterchangeSeries) and isinstance( series.dtype, dtypes.Categorical ): result = compliant.native.describe_categorical["is_ordered"] elif series.dtype == dtypes.Enum: result = True elif series.dtype != dtypes.Categorical: result = False else: native = series.to_native() impl = series.implementation if impl.is_polars() and impl._backend_version() < (1, 32): # NOTE: Deprecated https://github.com/pola-rs/polars/pull/23779 # Since version 1.32.0, ordering parameter is ignored and # it always behaves as if 'lexical' was passed. result = cast("pl.Categorical", native.dtype).ordering == "physical" elif impl.is_pandas_like(): result = bool(native.cat.ordered) elif impl.is_pyarrow(): from narwhals._arrow.utils import is_dictionary result = is_dictionary(native.type) and native.type.ordered return result def generate_unique_token( n_bytes: int, columns: Container[str], prefix: str = "nw" ) -> str: # pragma: no cover msg = ( "Use `generate_temporary_column_name` instead. `generate_unique_token` is " "deprecated and it will be removed in future versions" ) issue_deprecation_warning(msg, _version="1.13.0") return generate_temporary_column_name(n_bytes=n_bytes, columns=columns, prefix=prefix) def generate_temporary_column_name( n_bytes: int, columns: Container[str], prefix: str = "nw" ) -> str: """Generates a unique column name that is not present in the given list of columns. It relies on [python secrets token_hex](https://docs.python.org/3/library/secrets.html#secrets.token_hex) function to return a string nbytes random bytes. Arguments: n_bytes: The number of bytes to generate for the token. columns: The list of columns to check for uniqueness. prefix: prefix with which the temporary column name should start with. Returns: A unique token that is not present in the given list of columns. Raises: AssertionError: If a unique token cannot be generated after 100 attempts. Examples: >>> import narwhals as nw >>> columns = ["abc", "xyz"] >>> nw.generate_temporary_column_name(n_bytes=8, columns=columns) not in columns True >>> temp_name = nw.generate_temporary_column_name( ... n_bytes=8, columns=columns, prefix="foo" ... ) >>> temp_name not in columns and temp_name.startswith("foo") True """ counter = 0 while True: token = f"{prefix}{token_hex(n_bytes - 1)}" if token not in columns: return token counter += 1 if counter > 100: msg = ( "Internal Error: Narwhals was not able to generate a column name with " f"{n_bytes=} and not in {columns}" ) raise AssertionError(msg) def parse_columns_to_drop( frame: _StoresColumns, subset: Iterable[str], /, *, strict: bool ) -> list[str]: if not strict: return list(set(frame.columns).intersection(subset)) to_drop = list(subset) if error := check_columns_exist(to_drop, available=frame.columns): raise error return to_drop def is_sequence_but_not_str(sequence: Sequence[_T] | Any) -> TypeIs[Sequence[_T]]: return isinstance(sequence, Sequence) and not isinstance(sequence, str) def is_slice_none(obj: Any) -> TypeIs[_SliceNone]: return isinstance(obj, slice) and obj == slice(None) def is_sized_multi_index_selector( obj: Any, ) -> TypeIs[SizedMultiIndexSelector[Series[Any] | CompliantSeries[Any]]]: return ( ( is_sequence_but_not_str(obj) and ((len(obj) > 0 and isinstance(obj[0], int)) or (len(obj) == 0)) ) or is_numpy_array_1d_int(obj) or is_narwhals_series_int(obj) or is_compliant_series_int(obj) ) def is_sequence_like( obj: Sequence[_T] | Any, ) -> TypeIs[Sequence[_T] | Series[Any] | _1DArray]: return ( is_sequence_but_not_str(obj) or is_numpy_array_1d(obj) or is_narwhals_series(obj) or is_compliant_series(obj) ) def is_slice_index(obj: Any) -> TypeIs[_SliceIndex]: return isinstance(obj, slice) and ( isinstance(obj.start, int) or isinstance(obj.stop, int) or (isinstance(obj.step, int) and obj.start is None and obj.stop is None) ) def is_range(obj: Any) -> TypeIs[range]: return isinstance(obj, range) def is_single_index_selector(obj: Any) -> TypeIs[SingleIndexSelector]: return bool(isinstance(obj, int) and not isinstance(obj, bool)) def is_index_selector( obj: Any, ) -> TypeIs[SingleIndexSelector | MultiIndexSelector[Series[Any] | CompliantSeries[Any]]]: return ( is_single_index_selector(obj) or is_sized_multi_index_selector(obj) or is_slice_index(obj) ) def is_list_of(obj: Any, tp: type[_T]) -> TypeIs[list[_T]]: # Check if an object is a list of `tp`, only sniffing the first element. return bool(isinstance(obj, list) and obj and isinstance(obj[0], tp)) def predicates_contains_list_of_bool( predicates: Collection[Any], ) -> TypeIs[Collection[list[bool]]]: return any(is_list_of(pred, bool) for pred in predicates) def is_sequence_of(obj: Any, tp: type[_T]) -> TypeIs[Sequence[_T]]: # Check if an object is a sequence of `tp`, only sniffing the first element. return bool( is_sequence_but_not_str(obj) and (first := next(iter(obj), None)) and isinstance(first, tp) ) def validate_strict_and_pass_though( strict: bool | None, # noqa: FBT001 pass_through: bool | None, # noqa: FBT001 *, pass_through_default: bool, ) -> bool: if strict is None and pass_through is None: pass_through = pass_through_default elif strict is not None and pass_through is None: pass_through = not strict elif strict is None and pass_through is not None: pass else: msg = "Cannot pass both `strict` and `pass_through`" raise ValueError(msg) return pass_through def deprecate_native_namespace( *, warn_version: str = "", required: bool = False ) -> Callable[[Callable[P, R]], Callable[P, R]]: """Decorator to transition from `native_namespace` to `backend` argument. Arguments: warn_version: Emit a deprecation warning from this version. required: Raise when both `native_namespace`, `backend` are `None`. Returns: Wrapped function, with `native_namespace` **removed**. """ def decorate(fn: Callable[P, R], /) -> Callable[P, R]: @wraps(fn) def wrapper(*args: P.args, **kwds: P.kwargs) -> R: backend = kwds.pop("backend", None) native_namespace = kwds.pop("native_namespace", None) if native_namespace is not None and backend is None: if warn_version: msg = ( "`native_namespace` is deprecated, please use `backend` instead.\n\n" "Note: `native_namespace` will remain available in `narwhals.stable.v1`.\n" "See https://narwhals-dev.github.io/narwhals/backcompat/ for more information.\n" ) issue_deprecation_warning(msg, _version=warn_version) backend = native_namespace elif native_namespace is not None and backend is not None: msg = "Can't pass both `native_namespace` and `backend`" raise ValueError(msg) elif native_namespace is None and backend is None and required: msg = f"`backend` must be specified in `{fn.__name__}`." raise ValueError(msg) kwds["backend"] = backend return fn(*args, **kwds) return wrapper return decorate def _validate_rolling_arguments( window_size: int, min_samples: int | None ) -> tuple[int, int]: ensure_type(window_size, int, param_name="window_size") ensure_type(min_samples, int, type(None), param_name="min_samples") if window_size < 1: msg = "window_size must be greater or equal than 1" raise ValueError(msg) if min_samples is not None: if min_samples < 1: msg = "min_samples must be greater or equal than 1" raise ValueError(msg) if min_samples > window_size: msg = "`min_samples` must be less or equal than `window_size`" raise InvalidOperationError(msg) else: min_samples = window_size return window_size, min_samples def generate_repr(header: str, native_repr: str) -> str: try: terminal_width = os.get_terminal_size().columns except OSError: terminal_width = int(os.getenv("COLUMNS", 80)) # noqa: PLW1508 native_lines = native_repr.expandtabs().splitlines() max_native_width = max(len(line) for line in native_lines) if max_native_width + 2 <= terminal_width: length = max(max_native_width, len(header)) output = f"┌{'─' * length}┐\n" header_extra = length - len(header) output += f"|{' ' * (header_extra // 2)}{header}{' ' * (header_extra // 2 + header_extra % 2)}|\n" output += f"|{'-' * (length)}|\n" start_extra = (length - max_native_width) // 2 end_extra = (length - max_native_width) // 2 + (length - max_native_width) % 2 for line in native_lines: output += f"|{' ' * (start_extra)}{line}{' ' * (end_extra + max_native_width - len(line))}|\n" output += f"└{'─' * length}┘" return output diff = 39 - len(header) return ( f"┌{'─' * (39)}┐\n" f"|{' ' * (diff // 2)}{header}{' ' * (diff // 2 + diff % 2)}|\n" "| Use `.to_native` to see native output |\n└" f"{'─' * 39}┘" ) def check_columns_exist( subset: Collection[str], /, *, available: Collection[str] ) -> ColumnNotFoundError | None: if missing := set(subset).difference(available): return ColumnNotFoundError.from_missing_and_available_column_names( missing, available ) return None def check_column_names_are_unique(columns: Collection[str]) -> None: if len(columns) != len(set(columns)): from collections import Counter counter = Counter(columns) duplicates = {k: v for k, v in counter.items() if v > 1} msg = "".join(f"\n- '{k}' {v} times" for k, v in duplicates.items()) msg = f"Expected unique column names, got:{msg}" raise DuplicateError(msg) def _parse_time_unit_and_time_zone( time_unit: TimeUnit | Iterable[TimeUnit] | None, time_zone: str | timezone | Iterable[str | timezone | None] | None, ) -> tuple[Set[TimeUnit], Set[str | None]]: time_units: Set[TimeUnit] = ( {"ms", "us", "ns", "s"} if time_unit is None else {time_unit} if isinstance(time_unit, str) else set(time_unit) ) time_zones: Set[str | None] = ( {None} if time_zone is None else {str(time_zone)} if isinstance(time_zone, (str, timezone)) else {str(tz) if tz is not None else None for tz in time_zone} ) return time_units, time_zones def dtype_matches_time_unit_and_time_zone( dtype: DType, dtypes: DTypes, time_units: Set[TimeUnit], time_zones: Set[str | None] ) -> bool: return ( isinstance(dtype, dtypes.Datetime) and (dtype.time_unit in time_units) and ( dtype.time_zone in time_zones or ("*" in time_zones and dtype.time_zone is not None) ) ) def get_column_names(frame: _StoresColumns, /) -> Sequence[str]: return frame.columns def exclude_column_names(frame: _StoresColumns, names: Container[str]) -> Sequence[str]: return [col_name for col_name in frame.columns if col_name not in names] def passthrough_column_names(names: Sequence[str], /) -> EvalNames[Any]: def fn(_frame: Any, /) -> Sequence[str]: return names return fn _SENTINEL: Final = object() def _hasattr_static(obj: Any, attr: str) -> bool: return getattr_static(obj, attr, _SENTINEL) is not _SENTINEL def is_compliant_dataframe( obj: CompliantDataFrame[ CompliantSeriesT, CompliantExprT, NativeDataFrameT, ToNarwhalsT_co ] | Any, ) -> TypeIs[ CompliantDataFrame[CompliantSeriesT, CompliantExprT, NativeDataFrameT, ToNarwhalsT_co] ]: return _hasattr_static(obj, "__narwhals_dataframe__") def is_compliant_lazyframe( obj: CompliantLazyFrame[CompliantExprT, NativeLazyFrameT, ToNarwhalsT_co] | Any, ) -> TypeIs[CompliantLazyFrame[CompliantExprT, NativeLazyFrameT, ToNarwhalsT_co]]: return _hasattr_static(obj, "__narwhals_lazyframe__") def is_compliant_series( obj: CompliantSeries[NativeSeriesT_co] | Any, ) -> TypeIs[CompliantSeries[NativeSeriesT_co]]: return _hasattr_static(obj, "__narwhals_series__") def is_compliant_series_int( obj: CompliantSeries[NativeSeriesT_co] | Any, ) -> TypeIs[CompliantSeries[NativeSeriesT_co]]: return is_compliant_series(obj) and obj.dtype.is_integer() def _is_namespace_accessor(obj: _IntoContext) -> TypeIs[NamespaceAccessor[_FullContext]]: # NOTE: Only `compliant` has false positives **internally** # - https://github.com/narwhals-dev/narwhals/blob/cc69bac35eb8c81a1106969c49bfba9fd569b856/narwhals/_compliant/group_by.py#L44-L49 # - https://github.com/narwhals-dev/narwhals/blob/cc69bac35eb8c81a1106969c49bfba9fd569b856/narwhals/_namespace.py#L166-L168 # NOTE: Only `_accessor` has false positives **upstream** # - https://github.com/pandas-dev/pandas/blob/e209a35403f8835bbcff97636b83d2fc39b51e68/pandas/core/accessor.py#L200-L233 # - https://github.com/pola-rs/polars/blob/a60c5019f7b694c97009ef9208d25aaa4cc1d8a6/py-polars/polars/api.py#L29-L42 return _hasattr_static(obj, "compliant") and _hasattr_static(obj, "_accessor") def is_eager_allowed(impl: Implementation, /) -> TypeIs[_EagerAllowedImpl]: """Return True if `impl` allows eager operations.""" return impl in { Implementation.CUDF, Implementation.MODIN, Implementation.PANDAS, Implementation.POLARS, Implementation.PYARROW, } def can_lazyframe_collect(impl: Implementation, /) -> TypeIs[_LazyFrameCollectImpl]: """Return True if `LazyFrame.collect(impl)` is allowed.""" return impl in {Implementation.PANDAS, Implementation.POLARS, Implementation.PYARROW} def is_lazy_allowed(impl: Implementation, /) -> TypeIs[_LazyAllowedImpl]: """Return True if `DataFrame.lazy(impl)` is allowed.""" return impl in { Implementation.DASK, Implementation.DUCKDB, Implementation.IBIS, Implementation.POLARS, Implementation.PYSPARK, Implementation.PYSPARK_CONNECT, Implementation.SQLFRAME, } def has_native_namespace(obj: Any) -> TypeIs[SupportsNativeNamespace]: return _hasattr_static(obj, "__native_namespace__") def supports_arrow_c_stream(obj: Any) -> TypeIs[ArrowStreamExportable]: return _hasattr_static(obj, "__arrow_c_stream__") def _remap_full_join_keys( left_on: Collection[str], right_on: Collection[str], suffix: str ) -> dict[str, str]: """Remap join keys to avoid collisions. If left keys collide with the right keys, append the suffix. If there's no collision, let the right keys be. Arguments: left_on: Left keys. right_on: Right keys. suffix: Suffix to append to right keys. Returns: A map of old to new right keys. """ right_keys_suffixed = ( f"{key}{suffix}" if key in left_on else key for key in right_on ) return dict(zip(right_on, right_keys_suffixed)) def _into_arrow_table(data: IntoArrowTable, context: _LimitedContext, /) -> pa.Table: """Guards `ArrowDataFrame.from_arrow` w/ safer imports. Arguments: data: Object which implements `__arrow_c_stream__`. context: Initialized compliant object. """ if find_spec("pyarrow"): ns = context._version.namespace.from_backend("pyarrow").compliant return ns._dataframe.from_arrow(data, context=ns).native msg = f"'pyarrow>=14.0.0' is required for `from_arrow` for object of type {qualified_type_name(data)!r}." # pragma: no cover raise ModuleNotFoundError(msg) # pragma: no cover # TODO @dangotbanned: Extend with runtime behavior for `v1.*` # See `narwhals.exceptions.NarwhalsUnstableWarning` def unstable(fn: _Fn, /) -> _Fn: """Visual-only marker for unstable functionality. Arguments: fn: Function to decorate. Returns: Decorated function (unchanged). Examples: >>> from narwhals._utils import unstable >>> @unstable ... def a_work_in_progress_feature(*args): ... return args >>> >>> a_work_in_progress_feature.__name__ 'a_work_in_progress_feature' >>> a_work_in_progress_feature(1, 2, 3) (1, 2, 3) """ return fn def _is_naive_format(format: str) -> bool: """Determines if a datetime format string is 'naive', i.e., does not include timezone information. A format is considered naive if it does not contain any of the following - '%s': Unix timestamp - '%z': UTC offset - 'Z' : UTC timezone designator Arguments: format: The datetime format string to check. Returns: bool: True if the format is naive (does not include timezone info), False otherwise. """ return not any(x in format for x in ("%s", "%z", "Z")) class not_implemented: # noqa: N801 """Mark some functionality as unsupported. Arguments: alias: optional name used instead of the data model hook [`__set_name__`]. Returns: An exception-raising [descriptor]. Notes: - Attribute/method name *doesn't* need to be declared twice - Allows different behavior when looked up on the class vs instance - Allows us to use `isinstance(...)` instead of monkeypatching an attribute to the function Examples: >>> from narwhals._utils import not_implemented >>> class Thing: ... def totally_ready(self) -> str: ... return "I'm ready!" ... ... not_ready_yet = not_implemented() >>> >>> thing = Thing() >>> thing.totally_ready() "I'm ready!" >>> thing.not_ready_yet() Traceback (most recent call last): ... NotImplementedError: 'not_ready_yet' is not implemented for: 'Thing'. ... >>> isinstance(Thing.not_ready_yet, not_implemented) True [`__set_name__`]: https://docs.python.org/3/reference/datamodel.html#object.__set_name__ [descriptor]: https://docs.python.org/3/howto/descriptor.html """ def __init__(self, alias: str | None = None, /) -> None: # NOTE: Don't like this # Trying to workaround `mypy` requiring `@property` everywhere self._alias: str | None = alias def __repr__(self) -> str: return f"<{type(self).__name__}>: {self._name_owner}.{self._name}" def __set_name__(self, owner: type[_T], name: str) -> None: # https://docs.python.org/3/howto/descriptor.html#customized-names self._name_owner: str = owner.__name__ self._name: str = self._alias or name def __get__( self, instance: _T | Literal["raise"] | None, owner: type[_T] | None = None, / ) -> Any: if instance is None: # NOTE: Branch for `cls._name` # We can check that to see if an instance of `type(self)` for # https://narwhals-dev.github.io/narwhals/api-completeness/expr/ return self # NOTE: Prefer not exposing the actual class we're defining in # `_implementation` may not be available everywhere implementation = getattr(instance, "_implementation", Implementation.UNKNOWN) if implementation is not Implementation.UNKNOWN: who = repr(implementation) else: who = self._name_owner _raise_not_implemented_error(self._name, who) return None # pragma: no cover def __call__(self, *args: Any, **kwds: Any) -> Any: # NOTE: Purely to duck-type as assignable to **any** instance method # Wouldn't be reachable through *regular* attribute access return self.__get__("raise") @classmethod def deprecated(cls, message: LiteralString, /) -> Self: # pragma: no cover """Alt constructor, wraps with `@deprecated`. Arguments: message: **Static-only** deprecation message, emitted in an IDE. [descriptor]: https://docs.python.org/3/howto/descriptor.html """ obj = cls() return deprecated(message)(obj) def _raise_not_implemented_error(what: str, who: str, /) -> NotImplementedError: msg = ( f"{what!r} is not implemented for: {who!r}.\n\n" "If you would like to see this functionality in `narwhals`, " "please open an issue at: https://github.com/narwhals-dev/narwhals/issues" ) raise NotImplementedError(msg) class requires: # noqa: N801 """Method decorator for raising under certain constraints. Attributes: _min_version: Minimum backend version. _hint: Optional suggested alternative. Examples: >>> from narwhals._utils import requires, Implementation >>> class SomeBackend: ... _implementation = Implementation.PYARROW ... _backend_version = 20, 0, 0 ... ... @requires.backend_version((9000, 0, 0)) ... def really_complex_feature(self) -> str: ... return "hello" >>> backend = SomeBackend() >>> backend.really_complex_feature() Traceback (most recent call last): ... NotImplementedError: `really_complex_feature` is only available in 'pyarrow>=9000.0.0', found version '20.0.0'. """ _min_version: tuple[int, ...] _hint: str _wrapped_name: str """(Unqualified) decorated method name. When used in a namespace accessor, it will be prefixed by the property name. """ @classmethod def backend_version(cls, minimum: tuple[int, ...], /, hint: str = "") -> Self: """Method decorator for raising below a minimum `_backend_version`. Arguments: minimum: Minimum backend version. hint: Optional suggested alternative. """ obj = cls.__new__(cls) obj._min_version = minimum obj._hint = hint return obj @staticmethod def _unparse_version(backend_version: tuple[int, ...], /) -> str: return ".".join(f"{d}" for d in backend_version) def _qualify_accessor_name(self, prefix: Accessor, /) -> None: # NOTE: Should only need to do this once per class (the first time the method is called) if "." not in self._wrapped_name: self._wrapped_name = f"{prefix}.{self._wrapped_name}" def _unwrap_context(self, instance: _IntoContext, /) -> tuple[tuple[int, ...], str]: if _is_namespace_accessor(instance): self._qualify_accessor_name(instance._accessor) compliant = instance.compliant else: compliant = instance return compliant._backend_version, str(compliant._implementation) def _ensure_version(self, instance: _IntoContext, /) -> None: version, backend = self._unwrap_context(instance) if version >= self._min_version: return minimum = self._unparse_version(self._min_version) found = self._unparse_version(version) msg = f"`{self._wrapped_name}` is only available in '{backend}>={minimum}', found version {found!r}." if self._hint: msg = f"{msg}\n{self._hint}" raise NotImplementedError(msg) def __call__( self, fn: _Method[_IntoContextT, P, R], / ) -> _Method[_IntoContextT, P, R]: self._wrapped_name = fn.__name__ @wraps(fn) def wrapper(instance: _IntoContextT, *args: P.args, **kwds: P.kwargs) -> R: self._ensure_version(instance) return fn(instance, *args, **kwds) # NOTE: Only getting a complaint from `mypy` return wrapper # type: ignore[return-value] def convert_str_slice_to_int_slice( str_slice: _SliceName, columns: Sequence[str] ) -> tuple[int | None, int | None, Any]: start = columns.index(str_slice.start) if str_slice.start is not None else None stop = columns.index(str_slice.stop) + 1 if str_slice.stop is not None else None step = str_slice.step return (start, stop, step) def inherit_doc( tp_parent: Callable[P, R1], / ) -> Callable[[_Constructor[_T, P, R2]], _Constructor[_T, P, R2]]: """Steal the class-level docstring from parent and attach to child `__init__`. Returns: Decorated constructor. Notes: - Passes static typing (mostly) - Passes at runtime """ def decorate(init_child: _Constructor[_T, P, R2], /) -> _Constructor[_T, P, R2]: if init_child.__name__ == "__init__" and issubclass(type(tp_parent), type): init_child.__doc__ = getdoc(tp_parent) return init_child msg = ( # pragma: no cover f"`@{inherit_doc.__name__}` is only allowed to decorate an `__init__` with a class-level doc.\n" f"Method: {init_child.__qualname__!r}\n" f"Parent: {tp_parent!r}" ) raise TypeError(msg) # pragma: no cover return decorate def qualified_type_name(obj: object | type[Any], /) -> str: tp = obj if isinstance(obj, type) else type(obj) module = tp.__module__ if tp.__module__ != "builtins" else "" return f"{module}.{tp.__name__}".lstrip(".") def ensure_type(obj: Any, /, *valid_types: type[Any], param_name: str = "") -> None: """Validate that an object is an instance of one or more specified types. Parameters: obj: The object to validate. *valid_types: One or more valid types that `obj` is expected to match. param_name: The name of the parameter being validated. Used to improve error message clarity. Raises: TypeError: If `obj` is not an instance of any of the provided `valid_types`. Examples: >>> from narwhals._utils import ensure_type >>> ensure_type(42, int, float) >>> ensure_type("hello", str) >>> ensure_type("hello", int, param_name="test") Traceback (most recent call last): ... TypeError: Expected 'int', got: 'str' test='hello' ^^^^^^^ >>> import polars as pl >>> import pandas as pd >>> df = pl.DataFrame([[1], [2], [3], [4], [5]], schema=[*"abcde"]) >>> ensure_type(df, pd.DataFrame, param_name="df") Traceback (most recent call last): ... TypeError: Expected 'pandas.core.frame.DataFrame', got: 'polars.dataframe.frame.DataFrame' df=polars.dataframe.frame.DataFrame(...) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ """ if not isinstance(obj, valid_types): # pragma: no cover tp_names = " | ".join(qualified_type_name(tp) for tp in valid_types) msg = f"Expected {tp_names!r}, got: {qualified_type_name(obj)!r}" if param_name: left_pad = " " * 4 val = repr(obj) if len(val) > 40: # truncate long reprs val = f"{qualified_type_name(obj)}(...)" assign = f"{left_pad}{param_name}=" underline = (" " * len(assign)) + ("^" * len(val)) msg = f"{msg}\n{assign}{val}\n{underline}" raise TypeError(msg) class _DeferredIterable(Generic[_T]): """Store a callable producing an iterable to defer collection until we need it.""" def __init__(self, into_iter: Callable[[], Iterable[_T]], /) -> None: self._into_iter: Callable[[], Iterable[_T]] = into_iter def __iter__(self) -> Iterator[_T]: yield from self._into_iter() def to_tuple(self) -> tuple[_T, ...]: # Collect and return as a `tuple`. it = self._into_iter() return it if isinstance(it, tuple) else tuple(it) @lru_cache(maxsize=64) def deep_attrgetter(attr: str, *nested: str) -> attrgetter[Any]: name = ".".join((attr, *nested)) if nested else attr return attrgetter(name) def deep_getattr(obj: Any, name_1: str, *nested: str) -> Any: """Perform a nested attribute lookup on `obj`.""" return deep_attrgetter(name_1, *nested)(obj) class Compliant( _StoresNative[NativeT_co], _StoresImplementation, Protocol[NativeT_co] ): ... class Narwhals(Protocol[NativeT_co]): """Minimal *Narwhals-level* protocol. Provides access to a compliant object: obj: Narwhals[NativeT_co]] compliant: Compliant[NativeT_co] = obj._compliant Which itself exposes: implementation: Implementation = compliant.implementation native: NativeT_co = compliant.native This interface is used for revealing which `Implementation` member is associated with **either**: - One or more [nominal] native type(s) - One or more [structural] type(s) - where the true native type(s) are [assignable to] *at least* one of them These relationships are defined in the `@overload`s of `_Implementation.__get__(...)`. [nominal]: https://typing.python.org/en/latest/spec/glossary.html#term-nominal [structural]: https://typing.python.org/en/latest/spec/glossary.html#term-structural [assignable to]: https://typing.python.org/en/latest/spec/glossary.html#term-assignable """ @property def _compliant(self) -> Compliant[NativeT_co]: ... class _Implementation: """Descriptor for matching an opaque `Implementation` on a generic class. Based on [pyright comment](https://github.com/microsoft/pyright/issues/3071#issuecomment-1043978070) """ def __set_name__(self, owner: type[Any], name: str) -> None: self.__name__: str = name @overload def __get__(self, instance: Narwhals[NativePolars], owner: Any) -> _PolarsImpl: ... @overload def __get__(self, instance: Narwhals[NativePandas], owner: Any) -> _PandasImpl: ... @overload def __get__(self, instance: Narwhals[NativeModin], owner: Any) -> _ModinImpl: ... @overload def __get__(self, instance: Narwhals[NativeCuDF], owner: Any) -> _CuDFImpl: ... @overload def __get__( self, instance: Narwhals[NativePandasLike], owner: Any ) -> _PandasLikeImpl: ... @overload def __get__(self, instance: Narwhals[NativeArrow], owner: Any) -> _ArrowImpl: ... @overload def __get__( self, instance: Narwhals[NativePolars | NativeArrow | NativePandas], owner: Any ) -> _PolarsImpl | _PandasImpl | _ArrowImpl: ... @overload def __get__(self, instance: Narwhals[NativeDuckDB], owner: Any) -> _DuckDBImpl: ... @overload def __get__( self, instance: Narwhals[NativeSQLFrame], owner: Any ) -> _SQLFrameImpl: ... @overload def __get__(self, instance: Narwhals[NativeDask], owner: Any) -> _DaskImpl: ... @overload def __get__(self, instance: Narwhals[NativeIbis], owner: Any) -> _IbisImpl: ... @overload def __get__( self, instance: Narwhals[NativePySpark | NativePySparkConnect], owner: Any ) -> _PySparkImpl | _PySparkConnectImpl: ... # NOTE: https://docs.python.org/3/howto/descriptor.html#invocation-from-a-class @overload def __get__(self, instance: None, owner: type[Narwhals[Any]]) -> Self: ... @overload def __get__( self, instance: DataFrame[Any] | Series[Any], owner: Any ) -> _EagerAllowedImpl: ... @overload def __get__(self, instance: LazyFrame[Any], owner: Any) -> _LazyAllowedImpl: ... def __get__(self, instance: Narwhals[Any] | None, owner: Any) -> Any: return self if instance is None else instance._compliant._implementation def to_pyarrow_table(tbl: pa.Table | pa.RecordBatchReader) -> pa.Table: import pyarrow as pa # ignore-banned-import if isinstance(tbl, pa.RecordBatchReader): # pragma: no cover return pa.Table.from_batches(tbl) return tbl def normalize_path(source: FileSource, /) -> str: if isinstance(source, str): return source from pathlib import Path return str(Path(source)) def extend_bool( value: bool | Iterable[bool], # noqa: FBT001 n_match: int, ) -> Sequence[bool]: """Ensure the given bool or sequence of bools is the correct length. Stolen from https://github.com/pola-rs/polars/blob/b8bfb07a4a37a8d449d6d1841e345817431142df/py-polars/polars/_utils/various.py#L580-L594 """ return (value,) * n_match if isinstance(value, bool) else tuple(value) class _NoDefault(Enum): # "borrowed" from # https://github.com/pandas-dev/pandas/blob/e7859983a814b1823cf26e3b491ae2fa3be47c53/pandas/_libs/lib.pyx#L2736-L2748 no_default = "NO_DEFAULT" def __repr__(self) -> str: # pragma: no cover return "" # the "no_default" sentinel should typically be used when one of the valid parameter # values is None, as otherwise we cannot determine if the caller has set that value. no_default = _NoDefault.no_default narwhals-2.11.0/narwhals/dataframe.py0000644000000000000000000041337013615410400014457 0ustar00from __future__ import annotations from abc import abstractmethod from functools import partial from itertools import chain from typing import ( TYPE_CHECKING, Any, Callable, ClassVar, Generic, Literal, NoReturn, TypeVar, get_args, overload, ) from narwhals._exceptions import issue_warning from narwhals._expression_parsing import ( _parse_into_expr, check_expressions_preserve_length, is_scalar_like, ) from narwhals._typing import Arrow, Pandas, _LazyAllowedImpl, _LazyFrameCollectImpl from narwhals._utils import ( Implementation, Version, _Implementation, can_lazyframe_collect, check_columns_exist, flatten, generate_repr, is_compliant_dataframe, is_compliant_lazyframe, is_eager_allowed, is_index_selector, is_iterator, is_lazy_allowed, is_list_of, is_sequence_like, is_slice_none, predicates_contains_list_of_bool, qualified_type_name, supports_arrow_c_stream, zip_strict, ) from narwhals.dependencies import is_numpy_array_2d, is_pyarrow_table from narwhals.exceptions import ( ColumnNotFoundError, InvalidOperationError, PerformanceWarning, ) from narwhals.functions import _from_dict_no_backend, _is_into_schema from narwhals.schema import Schema from narwhals.series import Series from narwhals.translate import to_native if TYPE_CHECKING: from collections.abc import Iterable, Iterator, Mapping, Sequence from io import BytesIO from pathlib import Path from types import ModuleType import pandas as pd import polars as pl import pyarrow as pa from typing_extensions import Concatenate, ParamSpec, Self, TypeAlias from narwhals._compliant import CompliantDataFrame, CompliantLazyFrame from narwhals._compliant.typing import CompliantExprAny from narwhals._expression_parsing import ExprMetadata from narwhals._translate import IntoArrowTable from narwhals._typing import EagerAllowed, IntoBackend, LazyAllowed, Polars from narwhals.dtypes import DType from narwhals.group_by import GroupBy, LazyGroupBy from narwhals.typing import ( AsofJoinStrategy, IntoDataFrame, IntoExpr, IntoFrame, IntoLazyFrame, IntoSchema, JoinStrategy, MultiColSelector as _MultiColSelector, MultiIndexSelector as _MultiIndexSelector, PivotAgg, SingleColSelector, SingleIndexSelector, SizeUnit, UniqueKeepStrategy, _2DArray, ) PS = ParamSpec("PS") Incomplete: TypeAlias = Any _FrameT = TypeVar("_FrameT", bound="IntoFrame") LazyFrameT = TypeVar("LazyFrameT", bound="IntoLazyFrame") DataFrameT = TypeVar("DataFrameT", bound="IntoDataFrame") R = TypeVar("R") MultiColSelector: TypeAlias = "_MultiColSelector[Series[Any]]" MultiIndexSelector: TypeAlias = "_MultiIndexSelector[Series[Any]]" class BaseFrame(Generic[_FrameT]): _compliant_frame: Any _level: Literal["full", "lazy", "interchange"] implementation: _Implementation = _Implementation() """Return [`narwhals.Implementation`][] of native frame. This can be useful when you need to use special-casing for features outside of Narwhals' scope - for example, when dealing with pandas' Period Dtype. Examples: >>> import narwhals as nw >>> import pandas as pd >>> df_native = pd.DataFrame({"a": [1, 2, 3]}) >>> df = nw.from_native(df_native) >>> df.implementation >>> df.implementation.is_pandas() True >>> df.implementation.is_pandas_like() True >>> df.implementation.is_polars() False """ @property @abstractmethod def _compliant(self) -> Any: ... def __native_namespace__(self) -> ModuleType: return self._compliant_frame.__native_namespace__() # type: ignore[no-any-return] def __narwhals_namespace__(self) -> Any: return self._compliant_frame.__narwhals_namespace__() def _with_compliant(self, df: Any) -> Self: # construct, preserving properties return self.__class__(df, level=self._level) # type: ignore[call-arg] def _flatten_and_extract( self, *exprs: IntoExpr | Iterable[IntoExpr], **named_exprs: IntoExpr ) -> list[CompliantExprAny]: # Process `args` and `kwargs`, extracting underlying objects as we go. # NOTE: Strings are interpreted as column names. out_exprs = [] ns = self.__narwhals_namespace__() parse = partial( _parse_into_expr, backend=self._compliant._implementation, allow_literal=False ) all_exprs = chain( (parse(x) for x in flatten(exprs)), (parse(expr).alias(alias) for alias, expr in named_exprs.items()), ) for expr in all_exprs: ce = expr._to_compliant_expr(ns) out_exprs.append(ce) self._validate_metadata(ce._metadata) return out_exprs def _extract_compliant_frame(self, other: Self | Any, /) -> Any: if isinstance(other, type(self)): return other._compliant_frame msg = f"Expected `other` to be a {qualified_type_name(self)!r}, got: {qualified_type_name(other)!r}" raise TypeError(msg) def _check_columns_exist(self, subset: Sequence[str]) -> ColumnNotFoundError | None: return check_columns_exist(subset, available=self.columns) @abstractmethod def _validate_metadata(self, metadata: ExprMetadata) -> None: # pragma: no cover pass @property def schema(self) -> Schema: return Schema(self._compliant_frame.schema.items()) def collect_schema(self) -> Schema: native_schema = dict(self._compliant_frame.collect_schema()) return Schema(native_schema) def pipe( self, function: Callable[Concatenate[Self, PS], R], *args: PS.args, **kwargs: PS.kwargs, ) -> R: return function(self, *args, **kwargs) def drop_nulls(self, subset: str | list[str] | None) -> Self: subset = [subset] if isinstance(subset, str) else subset return self._with_compliant(self._compliant_frame.drop_nulls(subset=subset)) @property def columns(self) -> list[str]: return self._compliant_frame.columns # type: ignore[no-any-return] def with_columns( self, *exprs: IntoExpr | Iterable[IntoExpr], **named_exprs: IntoExpr ) -> Self: compliant_exprs = self._flatten_and_extract(*exprs, **named_exprs) compliant_exprs = [ compliant_expr.broadcast() if is_scalar_like(compliant_expr) else compliant_expr for compliant_expr in compliant_exprs ] return self._with_compliant(self._compliant_frame.with_columns(*compliant_exprs)) def select( self, *exprs: IntoExpr | Iterable[IntoExpr], **named_exprs: IntoExpr ) -> Self: flat_exprs = tuple(flatten(exprs)) if flat_exprs and all(isinstance(x, str) for x in flat_exprs) and not named_exprs: # fast path! try: return self._with_compliant( self._compliant_frame.simple_select(*flat_exprs) ) except Exception as e: # Column not found is the only thing that can realistically be raised here. if error := self._check_columns_exist(flat_exprs): raise error from e raise compliant_exprs = self._flatten_and_extract(*flat_exprs, **named_exprs) if compliant_exprs and all(is_scalar_like(x) for x in compliant_exprs): return self._with_compliant(self._compliant_frame.aggregate(*compliant_exprs)) compliant_exprs = [ compliant_expr.broadcast() if is_scalar_like(compliant_expr) else compliant_expr for compliant_expr in compliant_exprs ] return self._with_compliant(self._compliant_frame.select(*compliant_exprs)) def rename(self, mapping: dict[str, str]) -> Self: return self._with_compliant(self._compliant_frame.rename(mapping)) def head(self, n: int) -> Self: return self._with_compliant(self._compliant_frame.head(n)) def tail(self, n: int) -> Self: return self._with_compliant(self._compliant_frame.tail(n)) def drop(self, *columns: Iterable[str], strict: bool) -> Self: return self._with_compliant(self._compliant_frame.drop(columns, strict=strict)) def filter( self, *predicates: IntoExpr | Iterable[IntoExpr], **constraints: Any ) -> Self: from narwhals.functions import col flat_predicates = flatten(predicates) plx = self.__narwhals_namespace__() compliant_predicates = self._flatten_and_extract(*flat_predicates) check_expressions_preserve_length(*compliant_predicates, function_name="filter") compliant_constraints = self._flatten_and_extract( *[col(name) == v for name, v in constraints.items()] ) predicate = plx.all_horizontal( *chain(compliant_predicates, compliant_constraints), ignore_nulls=False ) return self._with_compliant(self._compliant_frame.filter(predicate)) def sort( self, by: str | Iterable[str], *more_by: str, descending: bool | Sequence[bool] = False, nulls_last: bool = False, ) -> Self: by = flatten([*flatten([by]), *more_by]) return self._with_compliant( self._compliant_frame.sort(*by, descending=descending, nulls_last=nulls_last) ) def top_k( self, k: int, *, by: str | Iterable[str], reverse: bool | Sequence[bool] = False ) -> Self: flatten_by = flatten([by]) return self._with_compliant( self._compliant_frame.top_k(k, by=flatten_by, reverse=reverse) ) def join( self, other: Incomplete, on: str | list[str] | None, how: JoinStrategy, *, left_on: str | list[str] | None, right_on: str | list[str] | None, suffix: str, ) -> Self: _supported_joins = ("inner", "left", "full", "cross", "anti", "semi") on = [on] if isinstance(on, str) else on left_on = [left_on] if isinstance(left_on, str) else left_on right_on = [right_on] if isinstance(right_on, str) else right_on compliant = self._compliant_frame other = self._extract_compliant_frame(other) if how not in _supported_joins: msg = f"Only the following join strategies are supported: {_supported_joins}; found '{how}'." raise NotImplementedError(msg) if how == "cross": if left_on is not None or right_on is not None or on is not None: msg = "Can not pass `left_on`, `right_on` or `on` keys for cross join" raise ValueError(msg) result = compliant.join( other, how=how, left_on=None, right_on=None, suffix=suffix ) elif on is None: if left_on is None or right_on is None: msg = f"Either (`left_on` and `right_on`) or `on` keys should be specified for {how}." raise ValueError(msg) if len(left_on) != len(right_on): msg = "`left_on` and `right_on` must have the same length." raise ValueError(msg) result = compliant.join( other, how=how, left_on=left_on, right_on=right_on, suffix=suffix ) else: if left_on is not None or right_on is not None: msg = f"If `on` is specified, `left_on` and `right_on` should be None for {how}." raise ValueError(msg) result = compliant.join( other, how=how, left_on=on, right_on=on, suffix=suffix ) return self._with_compliant(result) def gather_every(self, n: int, offset: int = 0) -> Self: return self._with_compliant( self._compliant_frame.gather_every(n=n, offset=offset) ) def join_asof( self, other: Incomplete, *, left_on: str | None, right_on: str | None, on: str | None, by_left: str | list[str] | None, by_right: str | list[str] | None, by: str | list[str] | None, strategy: AsofJoinStrategy, suffix: str, ) -> Self: _supported_strategies = ("backward", "forward", "nearest") if strategy not in _supported_strategies: msg = f"Only the following strategies are supported: {_supported_strategies}; found '{strategy}'." raise NotImplementedError(msg) if (on is None) and (left_on is None or right_on is None): msg = "Either (`left_on` and `right_on`) or `on` keys should be specified." raise ValueError(msg) if (on is not None) and (left_on is not None or right_on is not None): msg = "If `on` is specified, `left_on` and `right_on` should be None." raise ValueError(msg) if (by is None) and ( (by_left is None and by_right is not None) or (by_left is not None and by_right is None) ): msg = ( "Can not specify only `by_left` or `by_right`, you need to specify both." ) raise ValueError(msg) if (by is not None) and (by_left is not None or by_right is not None): msg = "If `by` is specified, `by_left` and `by_right` should be None." raise ValueError(msg) if on is not None: left_on = right_on = on if by is not None: by_left = by_right = by by_left = [by_left] if isinstance(by_left, str) else by_left by_right = [by_right] if isinstance(by_right, str) else by_right if (isinstance(by_left, list) and isinstance(by_right, list)) and ( len(by_left) != len(by_right) ): msg = "`by_left` and `by_right` must have the same length." raise ValueError(msg) return self._with_compliant( self._compliant_frame.join_asof( self._extract_compliant_frame(other), left_on=left_on, right_on=right_on, by_left=by_left, by_right=by_right, strategy=strategy, suffix=suffix, ) ) def unpivot( self, on: str | list[str] | None, *, index: str | list[str] | None, variable_name: str, value_name: str, ) -> Self: on = [on] if isinstance(on, str) else on index = [index] if isinstance(index, str) else index return self._with_compliant( self._compliant_frame.unpivot( on=on, index=index, variable_name=variable_name, value_name=value_name ) ) def __neq__(self, other: object) -> NoReturn: msg = ( "DataFrame.__neq__ and LazyFrame.__neq__ are not implemented, please " "use expressions instead.\n\n" "Hint: instead of\n" " df != 0\n" "you may want to use\n" " df.select(nw.all() != 0)" ) raise NotImplementedError(msg) def __eq__(self, other: object) -> NoReturn: msg = ( "DataFrame.__eq__ and LazyFrame.__eq__ are not implemented, please " "use expressions instead.\n\n" "Hint: instead of\n" " df == 0\n" "you may want to use\n" " df.select(nw.all() == 0)" ) raise NotImplementedError(msg) def explode(self, columns: str | Sequence[str], *more_columns: str) -> Self: to_explode = ( [columns, *more_columns] if isinstance(columns, str) else [*columns, *more_columns] ) return self._with_compliant(self._compliant_frame.explode(columns=to_explode)) class DataFrame(BaseFrame[DataFrameT]): """Narwhals DataFrame, backed by a native eager dataframe. Warning: This class is not meant to be instantiated directly - instead: - If the native object is a eager dataframe from one of the supported backend (e.g. pandas.DataFrame, polars.DataFrame, pyarrow.Table), you can use [`narwhals.from_native`][]: ```py narwhals.from_native(native_dataframe) narwhals.from_native(native_dataframe, eager_only=True) ``` - If the object is a dictionary of column names and generic sequences mapping (e.g. `dict[str, list]`), you can create a DataFrame via [`narwhals.from_dict`][]: ```py narwhals.from_dict( data={"a": [1, 2, 3]}, backend=narwhals.get_native_namespace(another_object), ) ``` """ _version: ClassVar[Version] = Version.MAIN @property def _compliant(self) -> CompliantDataFrame[Any, Any, DataFrameT, Self]: return self._compliant_frame @property def _series(self) -> type[Series[Any]]: return Series @property def _lazyframe(self) -> type[LazyFrame[Any]]: return LazyFrame def _validate_metadata(self, metadata: ExprMetadata) -> None: # all is valid in eager case. pass def __init__(self, df: Any, *, level: Literal["full", "lazy", "interchange"]) -> None: self._level: Literal["full", "lazy", "interchange"] = level self._compliant_frame: CompliantDataFrame[Any, Any, DataFrameT, Self] if is_compliant_dataframe(df): self._compliant_frame = df.__narwhals_dataframe__() else: # pragma: no cover msg = f"Expected an object which implements `__narwhals_dataframe__`, got: {type(df)}" raise AssertionError(msg) @classmethod def from_arrow( cls, native_frame: IntoArrowTable, *, backend: IntoBackend[EagerAllowed] ) -> DataFrame[Any]: """Construct a DataFrame from an object which supports the PyCapsule Interface. Arguments: native_frame: Object which implements `__arrow_c_stream__`. backend: specifies which eager backend instantiate to. `backend` can be specified in various ways - As `Implementation.` with `BACKEND` being `PANDAS`, `PYARROW`, `POLARS`, `MODIN` or `CUDF`. - As a string: `"pandas"`, `"pyarrow"`, `"polars"`, `"modin"` or `"cudf"`. - Directly as a module `pandas`, `pyarrow`, `polars`, `modin` or `cudf`. Examples: >>> import pandas as pd >>> import polars as pl >>> import narwhals as nw >>> >>> df_native = pd.DataFrame({"a": [1, 2], "b": [4.2, 5.1]}) >>> nw.DataFrame.from_arrow(df_native, backend="polars") ┌──────────────────┐ |Narwhals DataFrame| |------------------| | shape: (2, 2) | | ┌─────┮─────┐ | | │ a ┆ b │ | | │ --- ┆ --- │ | | │ i64 ┆ f64 │ | | ╞═════╩═════╡ | | │ 1 ┆ 4.2 │ | | │ 2 ┆ 5.1 │ | | └─────â”ī─────┘ | └──────────────────┘ """ if not (supports_arrow_c_stream(native_frame) or is_pyarrow_table(native_frame)): msg = f"Given object of type {type(native_frame)} does not support PyCapsule interface" raise TypeError(msg) implementation = Implementation.from_backend(backend) if is_eager_allowed(implementation): ns = cls._version.namespace.from_backend(implementation).compliant compliant = ns._dataframe.from_arrow(native_frame, context=ns) return cls(compliant, level="full") msg = ( f"{implementation} support in Narwhals is lazy-only, but `DataFrame.from_arrow` is an eager-only function.\n\n" "Hint: you may want to use an eager backend and then call `.lazy`, e.g.:\n\n" f" nw.DataFrame.from_arrow(df, backend='pyarrow').lazy('{implementation}')" ) raise ValueError(msg) @classmethod def from_dict( cls, data: Mapping[str, Any], schema: IntoSchema | Mapping[str, DType | None] | None = None, *, backend: IntoBackend[EagerAllowed] | None = None, ) -> DataFrame[Any]: """Instantiate DataFrame from dictionary. Indexes (if present, for pandas-like backends) are aligned following the [left-hand-rule](../concepts/pandas_index.md/). Notes: For pandas-like dataframes, conversion to schema is applied after dataframe creation. Arguments: data: Dictionary to create DataFrame from. schema: The DataFrame schema as Schema or dict of {name: type}. If not specified, the schema will be inferred by the native library. If any `dtype` is `None`, the data type for that column will be inferred by the native library. backend: specifies which eager backend instantiate to. Only necessary if inputs are not Narwhals Series. `backend` can be specified in various ways - As `Implementation.` with `BACKEND` being `PANDAS`, `PYARROW`, `POLARS`, `MODIN` or `CUDF`. - As a string: `"pandas"`, `"pyarrow"`, `"polars"`, `"modin"` or `"cudf"`. - Directly as a module `pandas`, `pyarrow`, `polars`, `modin` or `cudf`. Examples: >>> import pandas as pd >>> import narwhals as nw >>> data = {"c": [5, 2], "d": [1, 4]} >>> nw.DataFrame.from_dict(data, backend="pandas") ┌──────────────────┐ |Narwhals DataFrame| |------------------| | c d | | 0 5 1 | | 1 2 4 | └──────────────────┘ """ if backend is None: data, backend = _from_dict_no_backend(data) implementation = Implementation.from_backend(backend) if is_eager_allowed(implementation): ns = cls._version.namespace.from_backend(implementation).compliant compliant = ns._dataframe.from_dict(data, schema=schema, context=ns) return cls(compliant, level="full") # NOTE: (#2786) needs resolving for extensions msg = ( f"{implementation} support in Narwhals is lazy-only, but `DataFrame.from_dict` is an eager-only function.\n\n" "Hint: you may want to use an eager backend and then call `.lazy`, e.g.:\n\n" f" nw.DataFrame.from_dict({{'a': [1, 2]}}, backend='pyarrow').lazy('{implementation}')" ) raise ValueError(msg) @classmethod def from_dicts( cls, data: Sequence[Mapping[str, Any]], schema: IntoSchema | Mapping[str, DType | None] | None = None, *, backend: IntoBackend[EagerAllowed], ) -> DataFrame[Any]: """Instantiate DataFrame from a sequence of dictionaries representing rows. Notes: For pandas-like dataframes, conversion to schema is applied after dataframe creation. Arguments: data: Sequence with dictionaries mapping column name to value. schema: The DataFrame schema as Schema or dict of {name: type}. If not specified, the schema will be inferred by the native library. If any `dtype` is `None`, the data type for that column will be inferred by the native library. backend: Specifies which eager backend instantiate to. `backend` can be specified in various ways - As `Implementation.` with `BACKEND` being `PANDAS`, `PYARROW`, `POLARS`, `MODIN` or `CUDF`. - As a string: `"pandas"`, `"pyarrow"`, `"polars"`, `"modin"` or `"cudf"`. - Directly as a module `pandas`, `pyarrow`, `polars`, `modin` or `cudf`. Tip: If you expect non-uniform keys in `data`, consider passing `schema` for more consistent results, as **inference varies between backends**: - pandas uses all rows - polars uses the first 100 rows - pyarrow uses only the first row Examples: >>> import polars as pl >>> import narwhals as nw >>> data = [ ... {"item": "apple", "weight": 80, "price": 0.60}, ... {"item": "egg", "weight": 55, "price": 0.40}, ... ] >>> nw.DataFrame.from_dicts(data, backend="polars") ┌──────────────────────────┐ | Narwhals DataFrame | |--------------------------| |shape: (2, 3) | |┌───────┮────────┮───────┐| |│ item ┆ weight ┆ price │| |│ --- ┆ --- ┆ --- │| |│ str ┆ i64 ┆ f64 │| |╞═══════╩════════╩═══════╡| |│ apple ┆ 80 ┆ 0.6 │| |│ egg ┆ 55 ┆ 0.4 │| |└───────â”ī────────â”ī───────┘| └──────────────────────────┘ """ implementation = Implementation.from_backend(backend) if is_eager_allowed(implementation): ns = cls._version.namespace.from_backend(implementation).compliant compliant = ns._dataframe.from_dicts(data, schema=schema, context=ns) return cls(compliant, level="full") # NOTE: (#2786) needs resolving for extensions msg = ( f"{implementation} support in Narwhals is lazy-only, but `DataFrame.from_dicts` is an eager-only function.\n\n" "Hint: you may want to use an eager backend and then call `.lazy`, e.g.:\n\n" f" nw.DataFrame.from_dicts([{{'a': 1}}, {{'a': 2}}], backend='pyarrow').lazy('{implementation}')" ) raise ValueError(msg) @classmethod def from_numpy( cls, data: _2DArray, schema: IntoSchema | Sequence[str] | None = None, *, backend: IntoBackend[EagerAllowed], ) -> DataFrame[Any]: """Construct a DataFrame from a NumPy ndarray. Notes: Only row orientation is currently supported. For pandas-like dataframes, conversion to schema is applied after dataframe creation. Arguments: data: Two-dimensional data represented as a NumPy ndarray. schema: The DataFrame schema as Schema, dict of {name: type}, or a sequence of str. backend: specifies which eager backend instantiate to. `backend` can be specified in various ways - As `Implementation.` with `BACKEND` being `PANDAS`, `PYARROW`, `POLARS`, `MODIN` or `CUDF`. - As a string: `"pandas"`, `"pyarrow"`, `"polars"`, `"modin"` or `"cudf"`. - Directly as a module `pandas`, `pyarrow`, `polars`, `modin` or `cudf`. Examples: >>> import numpy as np >>> import polars as pl >>> import narwhals as nw >>> >>> arr = np.array([[5, 2, 1], [1, 4, 3]]) >>> schema = {"c": nw.Int16(), "d": nw.Float32(), "e": nw.Int8()} >>> nw.DataFrame.from_numpy(arr, schema=schema, backend="polars") ┌───────────────────┐ |Narwhals DataFrame | |-------------------| |shape: (2, 3) | |┌─────┮─────┮─────┐| |│ c ┆ d ┆ e │| |│ --- ┆ --- ┆ --- │| |│ i16 ┆ f32 ┆ i8 │| |╞═════╩═════╩═════╡| |│ 5 ┆ 2.0 ┆ 1 │| |│ 1 ┆ 4.0 ┆ 3 │| |└─────â”ī─────â”ī─────┘| └───────────────────┘ """ if not is_numpy_array_2d(data): msg = "`from_numpy` only accepts 2D numpy arrays" raise ValueError(msg) if not _is_into_schema(schema): msg = ( "`schema` is expected to be one of the following types: " "Mapping[str, DType] | Schema | Sequence[str]. " f"Got {type(schema)}." ) raise TypeError(msg) implementation = Implementation.from_backend(backend) if is_eager_allowed(implementation): ns = cls._version.namespace.from_backend(implementation).compliant return cls(ns.from_numpy(data, schema), level="full") msg = ( f"{implementation} support in Narwhals is lazy-only, but `DataFrame.from_numpy` is an eager-only function.\n\n" "Hint: you may want to use an eager backend and then call `.lazy`, e.g.:\n\n" f" nw.DataFrame.from_numpy(arr, backend='pyarrow').lazy('{implementation}')" ) raise ValueError(msg) def __len__(self) -> int: return self._compliant_frame.__len__() def __array__(self, dtype: Any = None, copy: bool | None = None) -> _2DArray: # noqa: FBT001 return self._compliant_frame.__array__(dtype, copy=copy) def __repr__(self) -> str: # pragma: no cover return generate_repr("Narwhals DataFrame", self.to_native().__repr__()) def __arrow_c_stream__(self, requested_schema: object | None = None) -> object: """Export a DataFrame via the Arrow PyCapsule Interface. - if the underlying dataframe implements the interface, it'll return that - else, it'll call `to_arrow` and then defer to PyArrow's implementation See [PyCapsule Interface](https://arrow.apache.org/docs/dev/format/CDataInterface/PyCapsuleInterface.html) for more. """ native_frame = self._compliant_frame._native_frame if supports_arrow_c_stream(native_frame): return native_frame.__arrow_c_stream__(requested_schema=requested_schema) try: pa_version = Implementation.PYARROW._backend_version() except ModuleNotFoundError as exc: msg = f"'pyarrow>=14.0.0' is required for `DataFrame.__arrow_c_stream__` for object of type {type(native_frame)}" raise ModuleNotFoundError(msg) from exc if pa_version < (14, 0): # pragma: no cover msg = f"'pyarrow>=14.0.0' is required for `DataFrame.__arrow_c_stream__` for object of type {type(native_frame)}" raise ModuleNotFoundError(msg) from None pa_table = self.to_arrow() return pa_table.__arrow_c_stream__(requested_schema=requested_schema) # type: ignore[no-untyped-call] def lazy( self, backend: IntoBackend[LazyAllowed] | None = None, *, session: Any | None = None, ) -> LazyFrame[Any]: """Restrict available API methods to lazy-only ones. If `backend` is specified, then a conversion between different backends might be triggered. If a library does not support lazy execution and `backend` is not specified, then this is will only restrict the API to lazy-only operations. This is useful if you want to ensure that you write dataframe-agnostic code which all has the possibility of running entirely lazily. Note: If `backend` is spark-like, then a valid `session` is required. For instance: ```py import narwhals as nw from sqlframe.duckdb import DuckDBSession df.lazy(backend=nw.Implementation.SQLFRAME, session=DuckDBSession()) ``` Arguments: backend: Which lazy backend collect to. This will be the underlying backend for the resulting Narwhals LazyFrame. If not specified, and the given library does not support lazy execution, then this will restrict the API to lazy-only operations. `backend` can be specified in various ways - As `Implementation.` with `BACKEND` being `DASK`, `DUCKDB`, `IBIS` or `POLARS`. - As a string: `"dask"`, `"duckdb"`, `"ibis"` or `"polars"` - Directly as a module `dask.dataframe`, `duckdb`, `ibis` or `polars`. session: Session to be used if backend is spark-like. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"a": [1, 2], "b": [4, 6]}) >>> df = nw.from_native(df_native) If we call `df.lazy`, we get a `narwhals.LazyFrame` backed by a Polars LazyFrame. >>> df.lazy() # doctest: +SKIP ┌─────────────────────────────┐ | Narwhals LazyFrame | |-----------------------------| || └─────────────────────────────┘ We can also pass DuckDB as the backend, and then we'll get a `narwhals.LazyFrame` backed by a `duckdb.DuckDBPyRelation`. >>> df.lazy(backend=nw.Implementation.DUCKDB) ┌──────────────────┐ |Narwhals LazyFrame| |------------------| |┌───────┮───────┐ | |│ a │ b │ | |│ int64 │ int64 │ | |├───────┾───────â”Ī | |│ 1 │ 4 │ | |│ 2 │ 6 │ | |└───────â”ī───────┘ | └──────────────────┘ """ lazy = self._compliant_frame.lazy if backend is None: return self._lazyframe(lazy(None, session=session), level="lazy") lazy_backend = Implementation.from_backend(backend) if is_lazy_allowed(lazy_backend): return self._lazyframe(lazy(lazy_backend, session=session), level="lazy") msg = f"Not-supported backend.\n\nExpected one of {get_args(_LazyAllowedImpl)} or `None`, got {lazy_backend}" raise ValueError(msg) def to_native(self) -> DataFrameT: """Convert Narwhals DataFrame to native one. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame( ... {"foo": [1, 2, 3], "bar": [6.0, 7.0, 8.0], "ham": ["a", "b", "c"]} ... ) Calling `to_native` on a Narwhals DataFrame returns the native object: >>> nw.from_native(df_native).to_native() foo bar ham 0 1 6.0 a 1 2 7.0 b 2 3 8.0 c """ return self._compliant_frame._native_frame def to_pandas(self) -> pd.DataFrame: """Convert this DataFrame to a pandas DataFrame. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame( ... {"foo": [1, 2, 3], "bar": [6.0, 7.0, 8.0], "ham": ["a", "b", "c"]} ... ) >>> df = nw.from_native(df_native) >>> df.to_pandas() foo bar ham 0 1 6.0 a 1 2 7.0 b 2 3 8.0 c """ return self._compliant_frame.to_pandas() def to_polars(self) -> pl.DataFrame: """Convert this DataFrame to a polars DataFrame. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"foo": [1, 2], "bar": [6.0, 7.0]}) >>> df = nw.from_native(df_native) >>> df.to_polars() shape: (2, 2) ┌─────┮─────┐ │ foo ┆ bar │ │ --- ┆ --- │ │ i64 ┆ f64 │ ╞═════╩═════╡ │ 1 ┆ 6.0 │ │ 2 ┆ 7.0 │ └─────â”ī─────┘ """ return self._compliant_frame.to_polars() @overload def write_csv(self, file: None = None) -> str: ... @overload def write_csv(self, file: str | Path | BytesIO) -> None: ... def write_csv(self, file: str | Path | BytesIO | None = None) -> str | None: r"""Write dataframe to comma-separated values (CSV) file. Arguments: file: String, path object or file-like object to which the dataframe will be written. If None, the resulting csv format is returned as a string. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame( ... {"foo": [1, 2, 3], "bar": [6.0, 7.0, 8.0], "ham": ["a", "b", "c"]} ... ) >>> df = nw.from_native(df_native) >>> df.write_csv() # doctest: +SKIP 'foo,bar,ham\n1,6.0,a\n2,7.0,b\n3,8.0,c\n' If we had passed a file name to `write_csv`, it would have been written to that file. """ return self._compliant_frame.write_csv(file) def write_parquet(self, file: str | Path | BytesIO) -> None: """Write dataframe to parquet file. Arguments: file: String, path object or file-like object to which the dataframe will be written. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"foo": [1, 2], "bar": [6.0, 7.0]}) >>> df = nw.from_native(df_native) >>> df.write_parquet("out.parquet") # doctest:+SKIP """ self._compliant_frame.write_parquet(file) def to_numpy(self) -> _2DArray: """Convert this DataFrame to a NumPy ndarray. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"foo": [1, 2], "bar": [6.5, 7.0]}) >>> df = nw.from_native(df_native) >>> df.to_numpy() array([[1. , 6.5], [2. , 7. ]]) """ return self._compliant_frame.to_numpy(None, copy=None) @property def shape(self) -> tuple[int, int]: """Get the shape of the DataFrame. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"foo": [1, 2]}) >>> df = nw.from_native(df_native) >>> df.shape (2, 1) """ return self._compliant_frame.shape def get_column(self, name: str) -> Series[Any]: """Get a single column by name. Arguments: name: The column name as a string. Notes: Although `name` is typed as `str`, pandas does allow non-string column names, and they will work when passed to this function if the `narwhals.DataFrame` is backed by a pandas dataframe with non-string columns. This function can only be used to extract a column by name, so there is no risk of ambiguity. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2]}) >>> df = nw.from_native(df_native) >>> df.get_column("a").to_native() 0 1 1 2 Name: a, dtype: int64 """ return self._series(self._compliant_frame.get_column(name), level=self._level) def estimated_size(self, unit: SizeUnit = "b") -> int | float: """Return an estimation of the total (heap) allocated size of the `DataFrame`. Estimated size is given in the specified unit (bytes by default). Arguments: unit: 'b', 'kb', 'mb', 'gb', 'tb', 'bytes', 'kilobytes', 'megabytes', 'gigabytes', or 'terabytes'. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"foo": [1, 2], "bar": [6.0, 7.0]}) >>> df = nw.from_native(df_native) >>> df.estimated_size() 32 """ return self._compliant_frame.estimated_size(unit=unit) # `str` overlaps with `Sequence[str]` # We can ignore this but we must keep this overload ordering @overload def __getitem__(self, item: tuple[SingleIndexSelector, SingleColSelector]) -> Any: ... @overload def __getitem__( # type: ignore[overload-overlap] self, item: str | tuple[MultiIndexSelector, SingleColSelector] ) -> Series[Any]: ... @overload def __getitem__( self, item: ( SingleIndexSelector | MultiIndexSelector | MultiColSelector | tuple[SingleIndexSelector, MultiColSelector] | tuple[MultiIndexSelector, MultiColSelector] ), ) -> Self: ... def __getitem__( # noqa: C901, PLR0912 self, item: ( SingleIndexSelector | SingleColSelector | MultiColSelector | MultiIndexSelector | tuple[SingleIndexSelector, SingleColSelector] | tuple[SingleIndexSelector, MultiColSelector] | tuple[MultiIndexSelector, SingleColSelector] | tuple[MultiIndexSelector, MultiColSelector] ), ) -> Series[Any] | Self | Any: """Extract column or slice of DataFrame. Arguments: item: How to slice dataframe. What happens depends on what is passed. It's easiest to explain by example. Suppose we have a Dataframe `df` - `df['a']` extracts column `'a'` and returns a `Series`. - `df[0:2]` extracts the first two rows and returns a `DataFrame`. - `df[0:2, 'a']` extracts the first two rows from column `'a'` and returns a `Series`. - `df[0:2, 0]` extracts the first two rows from the first column and returns a `Series`. - `df[[0, 1], [0, 1, 2]]` extracts the first two rows and the first three columns and returns a `DataFrame` - `df[:, [0, 1, 2]]` extracts all rows from the first three columns and returns a `DataFrame`. - `df[:, ['a', 'c']]` extracts all rows and columns `'a'` and `'c'` and returns a `DataFrame`. - `df[['a', 'c']]` extracts all rows and columns `'a'` and `'c'` and returns a `DataFrame`. - `df[0: 2, ['a', 'c']]` extracts the first two rows and columns `'a'` and `'c'` and returns a `DataFrame` - `df[:, 0: 2]` extracts all rows from the first two columns and returns a `DataFrame` - `df[:, 'a': 'c']` extracts all rows and all columns positioned between `'a'` and `'c'` _inclusive_ and returns a `DataFrame`. For example, if the columns are `'a', 'd', 'c', 'b'`, then that would extract columns `'a'`, `'d'`, and `'c'`. Notes: - Integers are always interpreted as positions - Strings are always interpreted as column names. In contrast with Polars, pandas allows non-string column names. If you don't know whether the column name you're trying to extract is definitely a string (e.g. `df[df.columns[0]]`) then you should use `DataFrame.get_column` instead. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2]}) >>> df = nw.from_native(df_native) >>> df["a"].to_native() 0 1 1 2 Name: a, dtype: int64 """ from narwhals.series import Series msg = ( f"Unexpected type for `DataFrame.__getitem__`, got: {type(item)}.\n\n" "Hints:\n" "- use `df.item` to select a single item.\n" "- Use `df[indices, :]` to select rows positionally.\n" "- Use `df.filter(mask)` to filter rows based on a boolean mask." ) if isinstance(item, tuple): if len(item) > 2: tuple_msg = ( "Tuples cannot be passed to DataFrame.__getitem__ directly.\n\n" "Hint: instead of `df[indices]`, did you mean `df[indices, :]`?" ) raise TypeError(tuple_msg) rows = None if not item or is_slice_none(item[0]) else item[0] columns = None if len(item) < 2 or is_slice_none(item[1]) else item[1] if rows is None and columns is None: return self elif is_index_selector(item): rows = item columns = None elif is_sequence_like(item) or isinstance(item, (slice, str)): rows = None columns = item else: raise TypeError(msg) if isinstance(rows, str): raise TypeError(msg) compliant = self._compliant_frame if isinstance(columns, (int, str)): if isinstance(rows, int): return self.item(rows, columns) col_name = columns if isinstance(columns, str) else self.columns[columns] series = self.get_column(col_name) return series[rows] if rows is not None else series if isinstance(rows, Series): rows = rows._compliant_series if isinstance(columns, Series): columns = columns._compliant_series if rows is None: return self._with_compliant(compliant[:, columns]) if columns is None: return self._with_compliant(compliant[rows, :]) return self._with_compliant(compliant[rows, columns]) def __contains__(self, key: str) -> bool: return key in self.columns @overload def to_dict(self, *, as_series: Literal[True] = ...) -> dict[str, Series[Any]]: ... @overload def to_dict(self, *, as_series: Literal[False]) -> dict[str, list[Any]]: ... @overload def to_dict( self, *, as_series: bool = True ) -> dict[str, Series[Any]] | dict[str, list[Any]]: ... def to_dict( self, *, as_series: bool = True ) -> dict[str, Series[Any]] | dict[str, list[Any]]: """Convert DataFrame to a dictionary mapping column name to values. Arguments: as_series: If set to true ``True``, then the values are Narwhals Series, otherwise the values are Any. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"A": [1, 2], "fruits": ["banana", "apple"]}) >>> df = nw.from_native(df_native) >>> df.to_dict(as_series=False) {'A': [1, 2], 'fruits': ['banana', 'apple']} """ if as_series: return { key: self._series(value, level=self._level) for key, value in self._compliant_frame.to_dict( as_series=as_series ).items() } return self._compliant_frame.to_dict(as_series=as_series) def row(self, index: int) -> tuple[Any, ...]: """Get values at given row. Warning: You should NEVER use this method to iterate over a DataFrame; if you require row-iteration you should strongly prefer use of iter_rows() instead. Arguments: index: Row number. Notes: cuDF doesn't support this method. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"a": [1, 2], "b": [4, 5]}) >>> nw.from_native(df_native).row(1) (, ) """ return self._compliant_frame.row(index) # inherited def pipe( self, function: Callable[Concatenate[Self, PS], R], *args: PS.args, **kwargs: PS.kwargs, ) -> R: """Pipe function call. Arguments: function: Function to apply. args: Positional arguments to pass to function. kwargs: Keyword arguments to pass to function. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2], "ba": [4, 5]}) >>> nw.from_native(df_native).pipe( ... lambda _df: _df.select( ... [x for x in _df.columns if len(x) == 1] ... ).to_native() ... ) a 0 1 1 2 """ return super().pipe(function, *args, **kwargs) def drop_nulls(self, subset: str | list[str] | None = None) -> Self: """Drop rows that contain null values. Arguments: subset: Column name(s) for which null values are considered. If set to None (default), use all columns. Notes: pandas handles null values differently from Polars and PyArrow. See [null_handling](../concepts/null_handling.md) for reference. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"a": [1.0, None], "ba": [1.0, 2.0]}) >>> nw.from_native(df_native).drop_nulls().to_native() pyarrow.Table a: double ba: double ---- a: [[1]] ba: [[1]] """ return super().drop_nulls(subset=subset) def with_row_index( self, name: str = "index", *, order_by: str | Sequence[str] | None = None ) -> Self: """Insert column which enumerates rows. Arguments: name: The name of the column as a string. The default is "index". order_by: Column(s) to order by when computing the row index. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"a": [1, 2], "b": [4, 5]}) >>> nw.from_native(df_native).with_row_index().to_native() pyarrow.Table index: int64 a: int64 b: int64 ---- index: [[0,1]] a: [[1,2]] b: [[4,5]] """ order_by_ = [order_by] if isinstance(order_by, str) else order_by return self._with_compliant( self._compliant_frame.with_row_index(name, order_by=order_by_) ) @property def schema(self) -> Schema: r"""Get an ordered mapping of column names to their data type. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"foo": [1, 2], "bar": [6.0, 7.0]}) >>> nw.from_native(df_native).schema Schema({'foo': Int64, 'bar': Float64}) """ return super().schema def collect_schema(self) -> Schema: r"""Get an ordered mapping of column names to their data type. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"foo": [1, 2], "bar": [6.0, 7.0]}) >>> nw.from_native(df_native).collect_schema() Schema({'foo': Int64, 'bar': Float64}) """ return super().collect_schema() @property def columns(self) -> list[str]: """Get column names. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"foo": [1, 2], "bar": [6.0, 7.0]}) >>> nw.from_native(df_native).columns ['foo', 'bar'] """ return super().columns @overload def rows(self, *, named: Literal[False] = False) -> list[tuple[Any, ...]]: ... @overload def rows(self, *, named: Literal[True]) -> list[dict[str, Any]]: ... @overload def rows(self, *, named: bool) -> list[tuple[Any, ...]] | list[dict[str, Any]]: ... def rows( self, *, named: bool = False ) -> list[tuple[Any, ...]] | list[dict[str, Any]]: """Returns all data in the DataFrame as a list of rows of python-native values. Arguments: named: By default, each row is returned as a tuple of values given in the same order as the frame columns. Setting named=True will return rows of dictionaries instead. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"foo": [1, 2], "bar": [6.0, 7.0]}) >>> nw.from_native(df_native).rows() [(1, 6.0), (2, 7.0)] """ return self._compliant_frame.rows(named=named) # type: ignore[return-value] def iter_columns(self) -> Iterator[Series[Any]]: """Returns an iterator over the columns of this DataFrame. Yields: A Narwhals Series, backed by a native series. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"foo": [1, 2], "bar": [6.0, 7.0]}) >>> iter_columns = nw.from_native(df_native).iter_columns() >>> next(iter_columns) ┌───────────────────────┐ | Narwhals Series | |-----------------------| |0 1 | |1 2 | |Name: foo, dtype: int64| └───────────────────────┘ >>> next(iter_columns) ┌─────────────────────────┐ | Narwhals Series | |-------------------------| |0 6.0 | |1 7.0 | |Name: bar, dtype: float64| └─────────────────────────┘ """ for series in self._compliant_frame.iter_columns(): yield self._series(series, level=self._level) @overload def iter_rows( self, *, named: Literal[False] = ..., buffer_size: int = ... ) -> Iterator[tuple[Any, ...]]: ... @overload def iter_rows( self, *, named: Literal[True], buffer_size: int = ... ) -> Iterator[dict[str, Any]]: ... @overload def iter_rows( self, *, named: bool, buffer_size: int = ... ) -> Iterator[tuple[Any, ...]] | Iterator[dict[str, Any]]: ... def iter_rows( self, *, named: bool = False, buffer_size: int = 512 ) -> Iterator[tuple[Any, ...]] | Iterator[dict[str, Any]]: """Returns an iterator over the DataFrame of rows of python-native values. Arguments: named: By default, each row is returned as a tuple of values given in the same order as the frame columns. Setting named=True will return rows of dictionaries instead. buffer_size: Determines the number of rows that are buffered internally while iterating over the data. See https://docs.pola.rs/api/python/stable/reference/dataframe/api/polars.DataFrame.iter_rows.html Notes: cuDF doesn't support this method. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"foo": [1, 2], "bar": [6.0, 7.0]}) >>> iter_rows = nw.from_native(df_native).iter_rows() >>> next(iter_rows) (1, 6.0) >>> next(iter_rows) (2, 7.0) """ return self._compliant_frame.iter_rows(named=named, buffer_size=buffer_size) # type: ignore[return-value] def with_columns( self, *exprs: IntoExpr | Iterable[IntoExpr], **named_exprs: IntoExpr ) -> Self: r"""Add columns to this DataFrame. Added columns will replace existing columns with the same name. Arguments: *exprs: Column(s) to add, specified as positional arguments. Accepts expression input. Strings are parsed as column names, other non-expression inputs are parsed as literals. **named_exprs: Additional columns to add, specified as keyword arguments. The columns will be renamed to the keyword used. Note: Creating a new DataFrame using this method does not create a new copy of existing data. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2], "b": [0.5, 4.0]}) >>> ( ... nw.from_native(df_native) ... .with_columns((nw.col("a") * 2).alias("a*2")) ... .to_native() ... ) a b a*2 0 1 0.5 2 1 2 4.0 4 """ return super().with_columns(*exprs, **named_exprs) def select( self, *exprs: IntoExpr | Iterable[IntoExpr], **named_exprs: IntoExpr ) -> Self: r"""Select columns from this DataFrame. Arguments: *exprs: Column(s) to select, specified as positional arguments. Accepts expression input. Strings are parsed as column names, other non-expression inputs are parsed as literals. **named_exprs: Additional columns to select, specified as keyword arguments. The columns will be renamed to the keyword used. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"a": [1, 2], "b": [3, 4]}) >>> nw.from_native(df_native).select("a", a_plus_1=nw.col("a") + 1) ┌──────────────────┐ |Narwhals DataFrame| |------------------| |pyarrow.Table | |a: int64 | |a_plus_1: int64 | |---- | |a: [[1,2]] | |a_plus_1: [[2,3]] | └──────────────────┘ """ return super().select(*exprs, **named_exprs) def rename(self, mapping: dict[str, str]) -> Self: """Rename column names. Arguments: mapping: Key value pairs that map from old name to new name. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"foo": [1, 2], "bar": [6, 7]}) >>> nw.from_native(df_native).rename({"foo": "apple"}).to_native() pyarrow.Table apple: int64 bar: int64 ---- apple: [[1,2]] bar: [[6,7]] """ return super().rename(mapping) def head(self, n: int = 5) -> Self: """Get the first `n` rows. Arguments: n: Number of rows to return. If a negative value is passed, return all rows except the last `abs(n)`. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2], "b": [0.5, 4.0]}) >>> nw.from_native(df_native).head(1).to_native() a b 0 1 0.5 """ return super().head(n) def tail(self, n: int = 5) -> Self: """Get the last `n` rows. Arguments: n: Number of rows to return. If a negative value is passed, return all rows except the first `abs(n)`. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2], "b": [0.5, 4.0]}) >>> nw.from_native(df_native).tail(1) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 1 2 4.0 | └──────────────────┘ """ return super().tail(n) def drop(self, *columns: str | Iterable[str], strict: bool = True) -> Self: """Remove columns from the dataframe. Arguments: *columns: Names of the columns that should be removed from the dataframe. strict: Validate that all column names exist in the schema and throw an exception if a column name does not exist in the schema. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame( ... {"foo": [1, 2], "bar": [6.0, 7.0], "ham": ["a", "b"]} ... ) >>> nw.from_native(df_native).drop("ham").to_native() foo bar 0 1 6.0 1 2 7.0 """ return super().drop(*flatten(columns), strict=strict) def unique( self, subset: str | list[str] | None = None, *, keep: UniqueKeepStrategy = "any", maintain_order: bool = False, order_by: str | Sequence[str] | None = None, ) -> Self: """Drop duplicate rows from this dataframe. Arguments: subset: Column name(s) to consider when identifying duplicate rows. keep: {'first', 'last', 'any', 'none'} Which of the duplicate rows to keep. * 'any': Does not give any guarantee of which row is kept. This allows more optimizations. * 'none': Don't keep duplicate rows. * 'first': Keep first unique row. * 'last': Keep last unique row. maintain_order: Keep the same order as the original DataFrame. This may be more expensive to compute. order_by: Column(s) to order by when computing the row index. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame( ... {"foo": [1, 2], "bar": ["a", "a"], "ham": ["b", "b"]} ... ) >>> nw.from_native(df_native).unique(["bar", "ham"]).to_native() foo bar ham 0 1 a b """ if keep not in {"any", "none", "first", "last"}: msg = f"Expected {'any', 'none', 'first', 'last'}, got: {keep}" raise ValueError(msg) if isinstance(subset, str): subset = [subset] return self._with_compliant( self._compliant_frame.unique( subset, keep=keep, maintain_order=maintain_order, order_by=order_by ) ) def filter( self, *predicates: IntoExpr | Iterable[IntoExpr] | list[bool], **constraints: Any ) -> Self: r"""Filter the rows in the DataFrame based on one or more predicate expressions. The original order of the remaining rows is preserved. Arguments: *predicates: Expression(s) that evaluates to a boolean Series. Can also be a boolean list(s). **constraints: Column filters; use `name = value` to filter columns by the supplied value. Each constraint will behave the same as `nw.col(name).eq(value)`, and will be implicitly joined with the other filter conditions using &. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame( ... {"foo": [1, 2, 3], "bar": [6, 7, 8], "ham": ["a", "b", "c"]} ... ) Filter on one condition >>> nw.from_native(df_native).filter(nw.col("foo") > 1).to_native() foo bar ham 1 2 7 b 2 3 8 c Filter on multiple conditions with implicit `&` >>> nw.from_native(df_native).filter( ... nw.col("foo") < 3, nw.col("ham") == "a" ... ).to_native() foo bar ham 0 1 6 a Filter on multiple conditions with `|` >>> nw.from_native(df_native).filter( ... (nw.col("foo") == 1) | (nw.col("ham") == "c") ... ).to_native() foo bar ham 0 1 6 a 2 3 8 c Filter using `**kwargs` syntax >>> nw.from_native(df_native).filter(foo=2, ham="b").to_native() foo bar ham 1 2 7 b """ impl = self.implementation parsed_predicates = ( self._series.from_iterable("", p, backend=impl) if is_list_of(p, bool) else p for p in predicates ) return super().filter(*parsed_predicates, **constraints) @overload def group_by( self, *keys: IntoExpr | Iterable[IntoExpr], drop_null_keys: Literal[False] = ... ) -> GroupBy[Self]: ... @overload def group_by( self, *keys: str | Iterable[str], drop_null_keys: Literal[True] ) -> GroupBy[Self]: ... def group_by( self, *keys: IntoExpr | Iterable[IntoExpr], drop_null_keys: bool = False ) -> GroupBy[Self]: r"""Start a group by operation. Arguments: *keys: Column(s) to group by. Accepts expression input. Strings are parsed as column names. drop_null_keys: if True, then groups where any key is null won't be included in the result. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame( ... { ... "a": ["a", "b", "a", "b", "c"], ... "b": [1, 2, 1, 3, 3], ... "c": [5, 4, 3, 2, 1], ... } ... ) Group by one column and compute the sum of another column >>> nw.from_native(df_native, eager_only=True).group_by("a").agg( ... nw.col("b").sum() ... ).sort("a").to_native() a b 0 a 2 1 b 5 2 c 3 Group by multiple columns and compute the max of another column >>> ( ... nw.from_native(df_native, eager_only=True) ... .group_by(["a", "b"]) ... .agg(nw.max("c")) ... .sort("a", "b") ... .to_native() ... ) a b c 0 a 1 5 1 b 2 4 2 b 3 2 3 c 3 1 Expressions are also accepted. >>> nw.from_native(df_native, eager_only=True).group_by( ... "a", nw.col("b") // 2 ... ).agg(nw.col("c").mean()).to_native() a b c 0 a 0 4.0 1 b 1 3.0 2 c 1 1.0 """ from narwhals.group_by import GroupBy flat_keys = flatten(keys) if all(isinstance(key, str) for key in flat_keys): return GroupBy(self, flat_keys, drop_null_keys=drop_null_keys) from narwhals import col from narwhals.expr import Expr from narwhals.series import Series key_is_expr_or_series = tuple(isinstance(k, (Expr, Series)) for k in flat_keys) if drop_null_keys and any(key_is_expr_or_series): msg = "drop_null_keys cannot be True when keys contains Expr or Series" raise NotImplementedError(msg) _keys = [ k if is_expr else col(k) for k, is_expr in zip_strict(flat_keys, key_is_expr_or_series) ] expr_flat_keys = self._flatten_and_extract(*_keys) check_expressions_preserve_length( *expr_flat_keys, function_name="DataFrame.group_by" ) return GroupBy(self, expr_flat_keys, drop_null_keys=drop_null_keys) def sort( self, by: str | Iterable[str], *more_by: str, descending: bool | Sequence[bool] = False, nulls_last: bool = False, ) -> Self: r"""Sort the dataframe by the given columns. Arguments: by: Column(s) names to sort by. *more_by: Additional columns to sort by, specified as positional arguments. descending: Sort in descending order. When sorting by multiple columns, can be specified per column by passing a sequence of booleans. nulls_last: Place null values last. Note: Unlike Polars, it is not possible to specify a sequence of booleans for `nulls_last` in order to control per-column behaviour. Instead a single boolean is applied for all `by` columns. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame( ... {"foo": [2, 1], "bar": [6.0, 7.0], "ham": ["a", "b"]} ... ) >>> nw.from_native(df_native).sort("foo") ┌──────────────────┐ |Narwhals DataFrame| |------------------| | foo bar ham | | 1 1 7.0 b | | 0 2 6.0 a | └──────────────────┘ """ return super().sort(by, *more_by, descending=descending, nulls_last=nulls_last) def top_k( self, k: int, *, by: str | Iterable[str], reverse: bool | Sequence[bool] = False ) -> Self: r"""Return the `k` largest rows. Non-null elements are always preferred over null elements, regardless of the value of reverse. The output is not guaranteed to be in any particular order, sort the outputs afterwards if you wish the output to be sorted. Arguments: k: Number of rows to return. by: Column(s) used to determine the top rows. Accepts expression input. Strings are parsed as column names. reverse: Consider the k smallest elements of the by column(s) (instead of the k largest). This can be specified per column by passing a sequence of booleans. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame( ... {"a": ["a", "b", "a", "b", None, "c"], "b": [2, 1, 1, 3, 2, 1]} ... ) >>> nw.from_native(df_native).top_k(4, by=["b", "a"]) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 3 b 3 | | 0 a 2 | | 4 None 2 | | 5 c 1 | └──────────────────┘ """ return super().top_k(k, by=by, reverse=reverse) def join( self, other: Self, on: str | list[str] | None = None, how: JoinStrategy = "inner", *, left_on: str | list[str] | None = None, right_on: str | list[str] | None = None, suffix: str = "_right", ) -> Self: r"""Join in SQL-like fashion. Arguments: other: DataFrame to join with. on: Name(s) of the join columns in both DataFrames. If set, `left_on` and `right_on` should be None. how: Join strategy. * *inner*: Returns rows that have matching values in both tables. * *left*: Returns all rows from the left table, and the matched rows from the right table. * *full*: Returns all rows in both dataframes, with the suffix appended to the right join keys. * *cross*: Returns the Cartesian product of rows from both tables. * *semi*: Filter rows that have a match in the right table. * *anti*: Filter rows that do not have a match in the right table. left_on: Join column of the left DataFrame. right_on: Join column of the right DataFrame. suffix: Suffix to append to columns with a duplicate name. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_1_native = pd.DataFrame({"id": ["a", "b"], "price": [6.0, 7.0]}) >>> df_2_native = pd.DataFrame({"id": ["a", "b", "c"], "qty": [1, 2, 3]}) >>> nw.from_native(df_1_native).join(nw.from_native(df_2_native), on="id") ┌──────────────────┐ |Narwhals DataFrame| |------------------| | id price qty | | 0 a 6.0 1 | | 1 b 7.0 2 | └──────────────────┘ """ return super().join( other, how=how, left_on=left_on, right_on=right_on, on=on, suffix=suffix ) def join_asof( self, other: Self, *, left_on: str | None = None, right_on: str | None = None, on: str | None = None, by_left: str | list[str] | None = None, by_right: str | list[str] | None = None, by: str | list[str] | None = None, strategy: AsofJoinStrategy = "backward", suffix: str = "_right", ) -> Self: """Perform an asof join. This is similar to a left-join except that we match on nearest key rather than equal keys. For Polars, both DataFrames must be sorted by the `on` key (within each `by` group if specified). Arguments: other: DataFrame to join with. left_on: Name(s) of the left join column(s). right_on: Name(s) of the right join column(s). on: Join column of both DataFrames. If set, left_on and right_on should be None. by_left: join on these columns before doing asof join. by_right: join on these columns before doing asof join. by: join on these columns before doing asof join. strategy: Join strategy. The default is "backward". suffix: Suffix to append to columns with a duplicate name. * *backward*: selects the last row in the right DataFrame whose "on" key is less than or equal to the left's key. * *forward*: selects the first row in the right DataFrame whose "on" key is greater than or equal to the left's key. * *nearest*: search selects the last row in the right DataFrame whose value is nearest to the left's key. Examples: >>> from datetime import datetime >>> import pandas as pd >>> import narwhals as nw >>> data_gdp = { ... "datetime": [ ... datetime(2016, 1, 1), ... datetime(2017, 1, 1), ... datetime(2018, 1, 1), ... datetime(2019, 1, 1), ... datetime(2020, 1, 1), ... ], ... "gdp": [4164, 4411, 4566, 4696, 4827], ... } >>> data_population = { ... "datetime": [ ... datetime(2016, 3, 1), ... datetime(2018, 8, 1), ... datetime(2019, 1, 1), ... ], ... "population": [82.19, 82.66, 83.12], ... } >>> gdp_native = pd.DataFrame(data_gdp) >>> population_native = pd.DataFrame(data_population) >>> gdp = nw.from_native(gdp_native) >>> population = nw.from_native(population_native) >>> population.join_asof(gdp, on="datetime", strategy="backward") ┌──────────────────────────────┐ | Narwhals DataFrame | |------------------------------| | datetime population gdp| |0 2016-03-01 82.19 4164| |1 2018-08-01 82.66 4566| |2 2019-01-01 83.12 4696| └──────────────────────────────┘ """ return super().join_asof( other, left_on=left_on, right_on=right_on, on=on, by_left=by_left, by_right=by_right, by=by, strategy=strategy, suffix=suffix, ) # --- descriptive --- def is_duplicated(self) -> Series[Any]: r"""Get a mask of all duplicated rows in this DataFrame. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"foo": [2, 2, 2], "bar": [6.0, 6.0, 7.0]}) >>> nw.from_native(df_native).is_duplicated() ┌───────────────┐ |Narwhals Series| |---------------| | 0 True | | 1 True | | 2 False | | dtype: bool | └───────────────┘ """ return ~self.is_unique() def is_empty(self) -> bool: r"""Check if the dataframe is empty. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"foo": [2, 2, 2], "bar": [6.0, 6.0, 7.0]}) >>> nw.from_native(df_native).is_empty() False """ return len(self) == 0 def is_unique(self) -> Series[Any]: r"""Get a mask of all unique rows in this DataFrame. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"foo": [2, 2, 2], "bar": [6.0, 6.0, 7.0]}) >>> nw.from_native(df_native).is_unique() ┌───────────────┐ |Narwhals Series| |---------------| | 0 False | | 1 False | | 2 True | | dtype: bool | └───────────────┘ """ return self._series(self._compliant_frame.is_unique(), level=self._level) def null_count(self) -> Self: r"""Create a new DataFrame that shows the null counts per column. Notes: pandas handles null values differently from Polars and PyArrow. See [null_handling](../concepts/null_handling.md/) for reference. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"foo": [1, None], "bar": [2, 3]}) >>> nw.from_native(df_native).null_count() ┌──────────────────┐ |Narwhals DataFrame| |------------------| | pyarrow.Table | | foo: int64 | | bar: int64 | | ---- | | foo: [[1]] | | bar: [[0]] | └──────────────────┘ """ plx = self._compliant_frame.__narwhals_namespace__() result = self._compliant_frame.select(plx.all().null_count()) return self._with_compliant(result) def item(self, row: int | None = None, column: int | str | None = None) -> Any: r"""Return the DataFrame as a scalar, or return the element at the given row/column. Arguments: row: The *n*-th row. column: The column selected via an integer or a string (column name). Notes: If row/col not provided, this is equivalent to df[0,0], with a check that the shape is (1,1). With row/col, this is equivalent to df[row,col]. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"foo": [1, None], "bar": [2, 3]}) >>> nw.from_native(df_native).item(0, 1) 2 """ return self._compliant_frame.item(row=row, column=column) def clone(self) -> Self: r"""Create a copy of this DataFrame.""" return self._with_compliant(self._compliant_frame.clone()) def gather_every(self, n: int, offset: int = 0) -> Self: r"""Take every nth row in the DataFrame and return as a new DataFrame. Arguments: n: Gather every *n*-th row. offset: Starting index. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"foo": [1, None, 2, 3]}) >>> nw.from_native(df_native).gather_every(2) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | pyarrow.Table | | foo: int64 | | ---- | | foo: [[1,2]] | └──────────────────┘ """ return super().gather_every(n=n, offset=offset) def pivot( self, on: str | list[str], *, index: str | list[str] | None = None, values: str | list[str] | None = None, aggregate_function: PivotAgg | None = None, maintain_order: bool | None = None, sort_columns: bool = False, separator: str = "_", ) -> Self: r"""Create a spreadsheet-style pivot table as a DataFrame. Arguments: on: Name of the column(s) whose values will be used as the header of the output DataFrame. index: One or multiple keys to group by. If None, all remaining columns not specified on `on` and `values` will be used. At least one of `index` and `values` must be specified. values: One or multiple keys to group by. If None, all remaining columns not specified on `on` and `index` will be used. At least one of `index` and `values` must be specified. aggregate_function: Choose from - None: no aggregation takes place, will raise error if multiple values are in group. - A predefined aggregate function string, one of {'min', 'max', 'first', 'last', 'sum', 'mean', 'median', 'len'} maintain_order: Has no effect and is kept around only for backwards-compatibility. sort_columns: Sort the transposed columns by name. Default is by order of discovery. separator: Used as separator/delimiter in generated column names in case of multiple `values` columns. Examples: >>> import pandas as pd >>> import narwhals as nw >>> data = { ... "ix": [1, 1, 2, 2, 1, 2], ... "col": ["a", "a", "a", "a", "b", "b"], ... "foo": [0, 1, 2, 2, 7, 1], ... "bar": [0, 2, 0, 0, 9, 4], ... } >>> df_native = pd.DataFrame(data) >>> nw.from_native(df_native).pivot( ... "col", index="ix", aggregate_function="sum" ... ) ┌─────────────────────────────────┐ | Narwhals DataFrame | |---------------------------------| | ix foo_a foo_b bar_a bar_b| |0 1 1 7 2 9| |1 2 4 1 0 4| └─────────────────────────────────┘ """ if values is None and index is None: msg = "At least one of `values` and `index` must be passed" raise ValueError(msg) if maintain_order is not None: msg = ( "`maintain_order` has no effect and is only kept around for backwards-compatibility. " "You can safely remove this argument." ) issue_warning(msg, UserWarning) on = [on] if isinstance(on, str) else on values = [values] if isinstance(values, str) else values index = [index] if isinstance(index, str) else index return self._with_compliant( self._compliant_frame.pivot( on=on, index=index, values=values, aggregate_function=aggregate_function, sort_columns=sort_columns, separator=separator, ) ) def to_arrow(self) -> pa.Table: r"""Convert to arrow table. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"foo": [1, None], "bar": [2, 3]}) >>> nw.from_native(df_native).to_arrow() pyarrow.Table foo: double bar: int64 ---- foo: [[1,null]] bar: [[2,3]] """ return self._compliant_frame.to_arrow() def sample( self, n: int | None = None, *, fraction: float | None = None, with_replacement: bool = False, seed: int | None = None, ) -> Self: r"""Sample from this DataFrame. Arguments: n: Number of items to return. Cannot be used with fraction. fraction: Fraction of items to return. Cannot be used with n. with_replacement: Allow values to be sampled more than once. seed: Seed for the random number generator. If set to None (default), a random seed is generated for each sample operation. Notes: The results may not be consistent across libraries. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"foo": [1, 2, 3], "bar": [19, 32, 4]}) >>> nw.from_native(df_native).sample(n=2) # doctest:+SKIP ┌──────────────────┐ |Narwhals DataFrame| |------------------| | foo bar | | 2 3 4 | | 1 2 32 | └──────────────────┘ """ return self._with_compliant( self._compliant_frame.sample( n=n, fraction=fraction, with_replacement=with_replacement, seed=seed ) ) def unpivot( self, on: str | list[str] | None = None, *, index: str | list[str] | None = None, variable_name: str = "variable", value_name: str = "value", ) -> Self: r"""Unpivot a DataFrame from wide to long format. Optionally leaves identifiers set. This function is useful to massage a DataFrame into a format where one or more columns are identifier variables (index) while all other columns, considered measured variables (on), are "unpivoted" to the row axis leaving just two non-identifier columns, 'variable' and 'value'. Arguments: on: Column(s) to use as values variables; if `on` is empty all columns that are not in `index` will be used. index: Column(s) to use as identifier variables. variable_name: Name to give to the `variable` column. Defaults to "variable". value_name: Name to give to the `value` column. Defaults to "value". Notes: If you're coming from pandas, this is similar to `pandas.DataFrame.melt`, but with `index` replacing `id_vars` and `on` replacing `value_vars`. In other frameworks, you might know this operation as `pivot_longer`. Examples: >>> import pandas as pd >>> import narwhals as nw >>> data = {"a": ["x", "y", "z"], "b": [1, 3, 5], "c": [2, 4, 6]} >>> df_native = pd.DataFrame(data) >>> nw.from_native(df_native).unpivot(["b", "c"], index="a") ┌────────────────────┐ | Narwhals DataFrame | |--------------------| | a variable value| |0 x b 1| |1 y b 3| |2 z b 5| |3 x c 2| |4 y c 4| |5 z c 6| └────────────────────┘ """ return super().unpivot( on=on, index=index, variable_name=variable_name, value_name=value_name ) def explode(self, columns: str | Sequence[str], *more_columns: str) -> Self: """Explode the dataframe to long format by exploding the given columns. Notes: It is possible to explode multiple columns only if these columns must have matching element counts. Arguments: columns: Column names. The underlying columns being exploded must be of the `List` data type. *more_columns: Additional names of columns to explode, specified as positional arguments. Examples: >>> import polars as pl >>> import narwhals as nw >>> data = {"a": ["x", "y"], "b": [[1, 2], [3]]} >>> df_native = pl.DataFrame(data) >>> nw.from_native(df_native).explode("b").to_native() shape: (3, 2) ┌─────┮─────┐ │ a ┆ b │ │ --- ┆ --- │ │ str ┆ i64 │ ╞═════╩═════╡ │ x ┆ 1 │ │ x ┆ 2 │ │ y ┆ 3 │ └─────â”ī─────┘ """ return super().explode(columns, *more_columns) class LazyFrame(BaseFrame[LazyFrameT]): """Narwhals LazyFrame, backed by a native lazyframe. Warning: This class is not meant to be instantiated directly - instead use [`narwhals.from_native`][] with a native object that is a lazy dataframe from one of the supported backend (e.g. polars.LazyFrame, dask_expr._collection.DataFrame): ```py narwhals.from_native(native_lazyframe) ``` """ @property def _compliant(self) -> CompliantLazyFrame[Any, LazyFrameT, Self]: return self._compliant_frame @property def _dataframe(self) -> type[DataFrame[Any]]: return DataFrame def _validate_metadata(self, metadata: ExprMetadata) -> None: if metadata.n_orderable_ops > 0: msg = ( "Order-dependent expressions are not supported for use in LazyFrame.\n\n" "Hint: To make the expression valid, use `.over` with `order_by` specified.\n\n" "For example, if you wrote `nw.col('price').cum_sum()` and you have a column\n" "`'date'` which orders your data, then replace:\n\n" " nw.col('price').cum_sum()\n\n" " with:\n\n" " nw.col('price').cum_sum().over(order_by='date')\n" " ^^^^^^^^^^^^^^^^^^^^^^\n\n" "See https://narwhals-dev.github.io/narwhals/concepts/order_dependence/." ) raise InvalidOperationError(msg) if metadata.is_filtration: msg = ( "Length-changing expressions are not supported for use in LazyFrame, unless\n" "followed by an aggregation.\n\n" "Hints:\n" "- Instead of `lf.select(nw.col('a').head())`, use `lf.select('a').head()\n" "- Instead of `lf.select(nw.col('a').drop_nulls()).select(nw.sum('a'))`,\n" " use `lf.select(nw.col('a').drop_nulls().sum())\n" ) raise InvalidOperationError(msg) def __init__(self, df: Any, *, level: Literal["full", "lazy", "interchange"]) -> None: self._level = level self._compliant_frame: CompliantLazyFrame[Any, LazyFrameT, Self] if is_compliant_lazyframe(df): self._compliant_frame = df.__narwhals_lazyframe__() else: # pragma: no cover msg = f"Expected Polars LazyFrame or an object that implements `__narwhals_lazyframe__`, got: {type(df)}" raise AssertionError(msg) def __repr__(self) -> str: # pragma: no cover return generate_repr("Narwhals LazyFrame", self.to_native().__repr__()) def __getitem__(self, item: str | slice) -> NoReturn: msg = "Slicing is not supported on LazyFrame" raise TypeError(msg) def collect( self, backend: IntoBackend[Polars | Pandas | Arrow] | None = None, **kwargs: Any ) -> DataFrame[Any]: r"""Materialize this LazyFrame into a DataFrame. As each underlying lazyframe has different arguments to set when materializing the lazyframe into a dataframe, we allow to pass them as kwargs (see examples below for how to generalize the specification). Arguments: backend: specifies which eager backend collect to. This will be the underlying backend for the resulting Narwhals DataFrame. If None, then the following default conversions will be applied - `polars.LazyFrame` -> `polars.DataFrame` - `dask.DataFrame` -> `pandas.DataFrame` - `duckdb.PyRelation` -> `pyarrow.Table` - `pyspark.DataFrame` -> `pyarrow.Table` `backend` can be specified in various ways - As `Implementation.` with `BACKEND` being `PANDAS`, `PYARROW` or `POLARS`. - As a string: `"pandas"`, `"pyarrow"` or `"polars"` - Directly as a module `pandas`, `pyarrow` or `polars`. kwargs: backend specific kwargs to pass along. To know more please check the backend specific documentation - [polars.LazyFrame.collect](https://docs.pola.rs/api/python/dev/reference/lazyframe/api/polars.LazyFrame.collect.html) - [dask.dataframe.DataFrame.compute](https://docs.dask.org/en/stable/generated/dask.dataframe.DataFrame.compute.html) Examples: >>> import duckdb >>> import narwhals as nw >>> lf_native = duckdb.sql("SELECT * FROM VALUES (1, 2), (3, 4) df(a, b)") >>> lf = nw.from_native(lf_native) >>> lf ┌──────────────────┐ |Narwhals LazyFrame| |------------------| |┌───────┮───────┐ | |│ a │ b │ | |│ int32 │ int32 │ | |├───────┾───────â”Ī | |│ 1 │ 2 │ | |│ 3 │ 4 │ | |└───────â”ī───────┘ | └──────────────────┘ >>> lf.collect() ┌──────────────────┐ |Narwhals DataFrame| |------------------| | pyarrow.Table | | a: int32 | | b: int32 | | ---- | | a: [[1,3]] | | b: [[2,4]] | └──────────────────┘ """ collect = self._compliant_frame.collect if backend is None: return self._dataframe(collect(None, **kwargs), level="full") eager_backend = Implementation.from_backend(backend) if can_lazyframe_collect(eager_backend): return self._dataframe(collect(eager_backend, **kwargs), level="full") msg = f"Unsupported `backend` value.\nExpected one of {get_args(_LazyFrameCollectImpl)} or None, got: {eager_backend}." raise ValueError(msg) def to_native(self) -> LazyFrameT: """Convert Narwhals LazyFrame to native one. Examples: >>> import duckdb >>> import narwhals as nw >>> lf_native = duckdb.sql("SELECT * FROM VALUES (1, 2), (3, 4) df(a, b)") >>> nw.from_native(lf_native).to_native() ┌───────┮───────┐ │ a │ b │ │ int32 │ int32 │ ├───────┾───────â”Ī │ 1 │ 2 │ │ 3 │ 4 │ └───────â”ī───────┘ """ return to_native(narwhals_object=self, pass_through=False) # inherited def pipe( self, function: Callable[Concatenate[Self, PS], R], *args: PS.args, **kwargs: PS.kwargs, ) -> R: """Pipe function call. Arguments: function: Function to apply. args: Positional arguments to pass to function. kwargs: Keyword arguments to pass to function. Examples: >>> import duckdb >>> import narwhals as nw >>> lf_native = duckdb.sql("SELECT * FROM VALUES (1, 2), (3, 4) df(a, b)") >>> nw.from_native(lf_native).pipe(lambda x: x.select("a")).to_native() ┌───────┐ │ a │ │ int32 │ ├───────â”Ī │ 1 │ │ 3 │ └───────┘ """ return super().pipe(function, *args, **kwargs) def drop_nulls(self, subset: str | list[str] | None = None) -> Self: """Drop rows that contain null values. Arguments: subset: Column name(s) for which null values are considered. If set to None (default), use all columns. Notes: pandas handles null values differently from Polars and PyArrow. See [null_handling](../concepts/null_handling.md/) for reference. Examples: >>> import duckdb >>> import narwhals as nw >>> lf_native = duckdb.sql("SELECT * FROM VALUES (1, NULL), (3, 4) df(a, b)") >>> nw.from_native(lf_native).drop_nulls() ┌──────────────────┐ |Narwhals LazyFrame| |------------------| |┌───────┮───────┐ | |│ a │ b │ | |│ int32 │ int32 │ | |├───────┾───────â”Ī | |│ 3 │ 4 │ | |└───────â”ī───────┘ | └──────────────────┘ """ return super().drop_nulls(subset=subset) def with_row_index( self, name: str = "index", *, order_by: str | Sequence[str] ) -> Self: """Insert column which enumerates rows. Arguments: name: The name of the column as a string. The default is "index". order_by: Column(s) to order by when computing the row index. Examples: >>> import duckdb >>> import narwhals as nw >>> lf_native = duckdb.sql("SELECT * FROM VALUES (1, 5), (2, 4) df(a, b)") >>> nw.from_native(lf_native).with_row_index(order_by="a").sort("a").collect() ┌──────────────────┐ |Narwhals DataFrame| |------------------| | pyarrow.Table | | index: int64 | | a: int32 | | b: int32 | | ---- | | index: [[0,1]] | | a: [[1,2]] | | b: [[5,4]] | └──────────────────┘ >>> nw.from_native(lf_native).with_row_index(order_by="b").sort("a").collect() ┌──────────────────┐ |Narwhals DataFrame| |------------------| | pyarrow.Table | | index: int64 | | a: int32 | | b: int32 | | ---- | | index: [[1,0]] | | a: [[1,2]] | | b: [[5,4]] | └──────────────────┘ """ order_by_ = [order_by] if isinstance(order_by, str) else order_by return self._with_compliant( self._compliant_frame.with_row_index(name, order_by=order_by_) ) @property def schema(self) -> Schema: r"""Get an ordered mapping of column names to their data type. Examples: >>> import duckdb >>> import narwhals as nw >>> lf_native = duckdb.sql("SELECT * FROM VALUES (1, 4.5), (3, 2.) df(a, b)") >>> nw.from_native(lf_native).schema # doctest:+SKIP Schema({'a': Int32, 'b': Decimal}) """ if self._compliant_frame._version is not Version.V1: msg = ( "Resolving the schema of a LazyFrame is a potentially expensive operation. " "Use `LazyFrame.collect_schema()` to get the schema without this warning." ) issue_warning(msg, PerformanceWarning) return super().schema def collect_schema(self) -> Schema: r"""Get an ordered mapping of column names to their data type. Examples: >>> import duckdb >>> import narwhals as nw >>> lf_native = duckdb.sql("SELECT * FROM VALUES (1, 4.5), (3, 2.) df(a, b)") >>> nw.from_native(lf_native).collect_schema() Schema({'a': Int32, 'b': Decimal}) """ return super().collect_schema() @property def columns(self) -> list[str]: r"""Get column names. Examples: >>> import duckdb >>> import narwhals as nw >>> lf_native = duckdb.sql("SELECT * FROM VALUES (1, 4.5), (3, 2.) df(a, b)") >>> nw.from_native(lf_native).columns ['a', 'b'] """ return super().columns def with_columns( self, *exprs: IntoExpr | Iterable[IntoExpr], **named_exprs: IntoExpr ) -> Self: r"""Add columns to this LazyFrame. Added columns will replace existing columns with the same name. Arguments: *exprs: Column(s) to add, specified as positional arguments. Accepts expression input. Strings are parsed as column names, other non-expression inputs are parsed as literals. **named_exprs: Additional columns to add, specified as keyword arguments. The columns will be renamed to the keyword used. Note: Creating a new LazyFrame using this method does not create a new copy of existing data. Examples: >>> import duckdb >>> import narwhals as nw >>> lf_native = duckdb.sql("SELECT * FROM VALUES (1, 4.5), (3, 2.) df(a, b)") >>> nw.from_native(lf_native).with_columns(c=nw.col("a") + 1) ┌────────────────────────────────┐ | Narwhals LazyFrame | |--------------------------------| |┌───────┮──────────────┮───────┐| |│ a │ b │ c │| |│ int32 │ decimal(2,1) │ int32 │| |├───────┾──────────────┾───────â”Ī| |│ 1 │ 4.5 │ 2 │| |│ 3 │ 2.0 │ 4 │| |└───────â”ī──────────────â”ī───────┘| └────────────────────────────────┘ """ if not exprs and not named_exprs: msg = "At least one expression must be passed to LazyFrame.with_columns" raise ValueError(msg) return super().with_columns(*exprs, **named_exprs) def select( self, *exprs: IntoExpr | Iterable[IntoExpr], **named_exprs: IntoExpr ) -> Self: r"""Select columns from this LazyFrame. Arguments: *exprs: Column(s) to select, specified as positional arguments. Accepts expression input. Strings are parsed as column names. **named_exprs: Additional columns to select, specified as keyword arguments. The columns will be renamed to the keyword used. Notes: If you'd like to select a column whose name isn't a string (for example, if you're working with pandas) then you should explicitly use `nw.col` instead of just passing the column name. For example, to select a column named `0` use `df.select(nw.col(0))`, not `df.select(0)`. Examples: >>> import duckdb >>> import narwhals as nw >>> lf_native = duckdb.sql("SELECT * FROM VALUES (1, 4.5), (3, 2.) df(a, b)") >>> nw.from_native(lf_native).select("a", a_plus_1=nw.col("a") + 1) ┌────────────────────┐ | Narwhals LazyFrame | |--------------------| |┌───────┮──────────┐| |│ a │ a_plus_1 │| |│ int32 │ int32 │| |├───────┾──────────â”Ī| |│ 1 │ 2 │| |│ 3 │ 4 │| |└───────â”ī──────────┘| └────────────────────┘ """ if not exprs and not named_exprs: msg = "At least one expression must be passed to LazyFrame.select" raise ValueError(msg) return super().select(*exprs, **named_exprs) def rename(self, mapping: dict[str, str]) -> Self: r"""Rename column names. Arguments: mapping: Key value pairs that map from old name to new name, or a function that takes the old name as input and returns the new name. Examples: >>> import duckdb >>> import narwhals as nw >>> lf_native = duckdb.sql("SELECT * FROM VALUES (1, 4.5), (3, 2.) df(a, b)") >>> nw.from_native(lf_native).rename({"a": "c"}) ┌────────────────────────┐ | Narwhals LazyFrame | |------------------------| |┌───────┮──────────────┐| |│ c │ b │| |│ int32 │ decimal(2,1) │| |├───────┾──────────────â”Ī| |│ 1 │ 4.5 │| |│ 3 │ 2.0 │| |└───────â”ī──────────────┘| └────────────────────────┘ """ return super().rename(mapping) def head(self, n: int = 5) -> Self: r"""Get `n` rows. Arguments: n: Number of rows to return. Examples: >>> import dask.dataframe as dd >>> import narwhals as nw >>> lf_native = dd.from_dict({"a": [1, 2, 3], "b": [4, 5, 6]}, npartitions=1) >>> nw.from_native(lf_native).head(2).collect() ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 1 4 | | 1 2 5 | └──────────────────┘ """ return super().head(n) def drop(self, *columns: str | Iterable[str], strict: bool = True) -> Self: r"""Remove columns from the LazyFrame. Arguments: *columns: Names of the columns that should be removed from the dataframe. strict: Validate that all column names exist in the schema and throw an exception if a column name does not exist in the schema. Warning: `strict` argument is ignored for `polars<1.0.0`. Please consider upgrading to a newer version or pass to eager mode. Examples: >>> import duckdb >>> import narwhals as nw >>> lf_native = duckdb.sql("SELECT * FROM VALUES (1, 2), (3, 4) df(a, b)") >>> nw.from_native(lf_native).drop("a").to_native() ┌───────┐ │ b │ │ int32 │ ├───────â”Ī │ 2 │ │ 4 │ └───────┘ """ return super().drop(*flatten(columns), strict=strict) def unique( self, subset: str | list[str] | None = None, *, keep: UniqueKeepStrategy = "any", order_by: str | Sequence[str] | None = None, ) -> Self: """Drop duplicate rows from this LazyFrame. Arguments: subset: Column name(s) to consider when identifying duplicate rows. If set to `None`, use all columns. keep: {'any', 'none', 'first', 'last} Which of the duplicate rows to keep. * 'any': Does not give any guarantee of which row is kept. * 'none': Don't keep duplicate rows. * 'first': Keep the first row. Requires `order_by` to be specified. * 'last': Keep the last row. Requires `order_by` to be specified. order_by: Column(s) to order by when computing the row index. Examples: >>> import duckdb >>> import narwhals as nw >>> lf_native = duckdb.sql("SELECT * FROM VALUES (1, 3), (1, 4) df(a, b)") >>> nw.from_native(lf_native).unique("a").sort("a", descending=True) ┌──────────────────┐ |Narwhals LazyFrame| |------------------| |┌───────┮───────┐ | |│ a │ b │ | |│ int32 │ int32 │ | |├───────┾───────â”Ī | |│ 1 │ 3 │ | |└───────â”ī───────┘ | └──────────────────┘ """ if keep not in {"any", "none", "first", "last"}: msg = f"Expected {'any', 'none', 'first', 'last'}, got: {keep}" raise ValueError(msg) if keep in {"first", "last"} and not order_by: msg = ( "narwhals.LazyFrame makes no assumptions about row order, so only " "'first' and 'last' are only supported if `order_by` is passed." ) raise InvalidOperationError(msg) if isinstance(subset, str): subset = [subset] return self._with_compliant( self._compliant_frame.unique(subset=subset, keep=keep, order_by=order_by) ) def filter( self, *predicates: IntoExpr | Iterable[IntoExpr], **constraints: Any ) -> Self: r"""Filter the rows in the LazyFrame based on a predicate expression. The original order of the remaining rows is preserved. Arguments: *predicates: Expression(s) that evaluates to a boolean Series. **constraints: Column filters; use `name = value` to filter columns by the supplied value. Each constraint will behave the same as `nw.col(name).eq(value)`, and will be implicitly joined with the other filter conditions using &. Examples: >>> import duckdb >>> import narwhals as nw >>> df_native = duckdb.sql(''' ... SELECT * FROM VALUES ... (1, 6, 'a'), ... (2, 7, 'b'), ... (3, 8, 'c') ... df(foo, bar, ham) ... ''') Filter on one condition >>> nw.from_native(df_native).filter(nw.col("foo") > 1).to_native() ┌───────┮───────┮─────────┐ │ foo │ bar │ ham │ │ int32 │ int32 │ varchar │ ├───────┾───────┾─────────â”Ī │ 2 │ 7 │ b │ │ 3 │ 8 │ c │ └───────â”ī───────â”ī─────────┘ Filter on multiple conditions with implicit `&` >>> nw.from_native(df_native).filter( ... nw.col("foo") < 3, nw.col("ham") == "a" ... ).to_native() ┌───────┮───────┮─────────┐ │ foo │ bar │ ham │ │ int32 │ int32 │ varchar │ ├───────┾───────┾─────────â”Ī │ 1 │ 6 │ a │ └───────â”ī───────â”ī─────────┘ Filter on multiple conditions with `|` >>> nw.from_native(df_native).filter( ... (nw.col("foo") == 1) | (nw.col("ham") == "c") ... ).to_native() ┌───────┮───────┮─────────┐ │ foo │ bar │ ham │ │ int32 │ int32 │ varchar │ ├───────┾───────┾─────────â”Ī │ 1 │ 6 │ a │ │ 3 │ 8 │ c │ └───────â”ī───────â”ī─────────┘ Filter using `**kwargs` syntax >>> nw.from_native(df_native).filter(foo=2, ham="b").to_native() ┌───────┮───────┮─────────┐ │ foo │ bar │ ham │ │ int32 │ int32 │ varchar │ ├───────┾───────┾─────────â”Ī │ 2 │ 7 │ b │ └───────â”ī───────â”ī─────────┘ """ predicates_ = tuple(tuple(p) if is_iterator(p) else p for p in predicates) if predicates_contains_list_of_bool(predicates_): msg = "`LazyFrame.filter` is not supported with Python boolean masks - use expressions instead." raise TypeError(msg) return super().filter(*predicates_, **constraints) def sink_parquet(self, file: str | Path | BytesIO) -> None: """Write LazyFrame to Parquet file. This may allow larger-than-RAM datasets to be written to disk. Arguments: file: String, path object or file-like object to which the dataframe will be written. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.LazyFrame({"foo": [1, 2], "bar": [6.0, 7.0]}) >>> df = nw.from_native(df_native) >>> df.sink_parquet("out.parquet") # doctest:+SKIP """ self._compliant_frame.sink_parquet(file) @overload def group_by( self, *keys: IntoExpr | Iterable[IntoExpr], drop_null_keys: Literal[False] = ... ) -> LazyGroupBy[Self]: ... @overload def group_by( self, *keys: str | Iterable[str], drop_null_keys: Literal[True] ) -> LazyGroupBy[Self]: ... def group_by( self, *keys: IntoExpr | Iterable[IntoExpr], drop_null_keys: bool = False ) -> LazyGroupBy[Self]: r"""Start a group by operation. Arguments: *keys: Column(s) to group by. Accepts expression input. Strings are parsed as column names. drop_null_keys: if True, then groups where any key is null won't be included in the result. Examples: >>> import duckdb >>> import narwhals as nw >>> df_native = duckdb.sql( ... "SELECT * FROM VALUES (1, 'a'), (2, 'b'), (3, 'a') df(a, b)" ... ) >>> df = nw.from_native(df_native) >>> df.group_by("b").agg(nw.col("a").sum()).sort("b").to_native() ┌─────────┮────────┐ │ b │ a │ │ varchar │ int128 │ ├─────────┾────────â”Ī │ a │ 4 │ │ b │ 2 │ └─────────â”ī────────┘ Expressions are also accepted. >>> df.group_by(nw.col("b").str.len_chars()).agg( ... nw.col("a").sum() ... ).to_native() ┌───────┮────────┐ │ b │ a │ │ int64 │ int128 │ ├───────┾────────â”Ī │ 1 │ 6 │ └───────â”ī────────┘ """ from narwhals.group_by import LazyGroupBy flat_keys = flatten(keys) if all(isinstance(key, str) for key in flat_keys): return LazyGroupBy(self, flat_keys, drop_null_keys=drop_null_keys) from narwhals import col from narwhals.expr import Expr key_is_expr = tuple(isinstance(k, Expr) for k in flat_keys) if drop_null_keys and any(key_is_expr): msg = "drop_null_keys cannot be True when keys contains Expr" raise NotImplementedError(msg) _keys = [ k if is_expr else col(k) for k, is_expr in zip_strict(flat_keys, key_is_expr) ] expr_flat_keys = self._flatten_and_extract(*_keys) check_expressions_preserve_length( *expr_flat_keys, function_name="LazyFrame.group_by" ) return LazyGroupBy(self, expr_flat_keys, drop_null_keys=drop_null_keys) def sort( self, by: str | Iterable[str], *more_by: str, descending: bool | Sequence[bool] = False, nulls_last: bool = False, ) -> Self: r"""Sort the LazyFrame by the given columns. Arguments: by: Column(s) names to sort by. *more_by: Additional columns to sort by, specified as positional arguments. descending: Sort in descending order. When sorting by multiple columns, can be specified per column by passing a sequence of booleans. nulls_last: Place null values last; can specify a single boolean applying to all columns or a sequence of booleans for per-column control. Warning: Unlike Polars, it is not possible to specify a sequence of booleans for `nulls_last` in order to control per-column behaviour. Instead a single boolean is applied for all `by` columns. Examples: >>> import duckdb >>> import narwhals as nw >>> df_native = duckdb.sql( ... "SELECT * FROM VALUES (1, 6.0, 'a'), (2, 5.0, 'c'), (NULL, 4.0, 'b') df(a, b, c)" ... ) >>> df = nw.from_native(df_native) >>> df.sort("a") ┌──────────────────────────────────┐ | Narwhals LazyFrame | |----------------------------------| |┌───────┮──────────────┮─────────┐| |│ a │ b │ c │| |│ int32 │ decimal(2,1) │ varchar │| |├───────┾──────────────┾─────────â”Ī| |│ NULL │ 4.0 │ b │| |│ 1 │ 6.0 │ a │| |│ 2 │ 5.0 │ c │| |└───────â”ī──────────────â”ī─────────┘| └──────────────────────────────────┘ """ return super().sort(by, *more_by, descending=descending, nulls_last=nulls_last) def top_k( self, k: int, *, by: str | Iterable[str], reverse: bool | Sequence[bool] = False ) -> Self: r"""Return the `k` largest rows. Non-null elements are always preferred over null elements, regardless of the value of reverse. The output is not guaranteed to be in any particular order, sort the outputs afterwards if you wish the output to be sorted. Arguments: k: Number of rows to return. by: Column(s) used to determine the top rows. Accepts expression input. Strings are parsed as column names. reverse: Consider the k smallest elements of the by column(s) (instead of the k largest). This can be specified per column by passing a sequence of booleans. Examples: >>> import duckdb >>> import narwhals as nw >>> df_native = duckdb.sql( ... "SELECT * FROM VALUES ('a', 2), ('b', 1), ('a', 1), ('b', 3), (NULL, 2), ('c', 1) df(a, b)" ... ) >>> df = nw.from_native(df_native) >>> df.top_k(4, by=["b", "a"]) ┌───────────────────┐ |Narwhals LazyFrame | |-------------------| |┌─────────┮───────┐| |│ a │ b │| |│ varchar │ int32 │| |├─────────┾───────â”Ī| |│ b │ 3 │| |│ a │ 2 │| |│ NULL │ 2 │| |│ c │ 1 │| |└─────────â”ī───────┘| └───────────────────┘ """ return super().top_k(k, by=by, reverse=reverse) def join( self, other: Self, on: str | list[str] | None = None, how: JoinStrategy = "inner", *, left_on: str | list[str] | None = None, right_on: str | list[str] | None = None, suffix: str = "_right", ) -> Self: r"""Add a join operation to the Logical Plan. Arguments: other: Lazy DataFrame to join with. on: Name(s) of the join columns in both DataFrames. If set, `left_on` and `right_on` should be None. how: Join strategy. * *inner*: Returns rows that have matching values in both tables. * *left*: Returns all rows from the left table, and the matched rows from the right table. * *full*: Returns all rows in both dataframes, with the suffix appended to the right join keys. * *cross*: Returns the Cartesian product of rows from both tables. * *semi*: Filter rows that have a match in the right table. * *anti*: Filter rows that do not have a match in the right table. left_on: Join column of the left DataFrame. right_on: Join column of the right DataFrame. suffix: Suffix to append to columns with a duplicate name. Examples: >>> import duckdb >>> import narwhals as nw >>> df_native1 = duckdb.sql( ... "SELECT * FROM VALUES (1, 'a'), (2, 'b') df(a, b)" ... ) >>> df_native2 = duckdb.sql( ... "SELECT * FROM VALUES (1, 'x'), (3, 'y') df(a, c)" ... ) >>> df1 = nw.from_native(df_native1) >>> df2 = nw.from_native(df_native2) >>> df1.join(df2, on="a") ┌─────────────────────────────┐ | Narwhals LazyFrame | |-----------------------------| |┌───────┮─────────┮─────────┐| |│ a │ b │ c │| |│ int32 │ varchar │ varchar │| |├───────┾─────────┾─────────â”Ī| |│ 1 │ a │ x │| |└───────â”ī─────────â”ī─────────┘| └─────────────────────────────┘ """ return super().join( other, how=how, left_on=left_on, right_on=right_on, on=on, suffix=suffix ) def join_asof( self, other: Self, *, left_on: str | None = None, right_on: str | None = None, on: str | None = None, by_left: str | list[str] | None = None, by_right: str | list[str] | None = None, by: str | list[str] | None = None, strategy: AsofJoinStrategy = "backward", suffix: str = "_right", ) -> Self: """Perform an asof join. This is similar to a left-join except that we match on nearest key rather than equal keys. For Polars, both DataFrames must be sorted by the `on` key (within each `by` group if specified). Arguments: other: DataFrame to join with. left_on: Name(s) of the left join column(s). right_on: Name(s) of the right join column(s). on: Join column of both DataFrames. If set, left_on and right_on should be None. by_left: join on these columns before doing asof join by_right: join on these columns before doing asof join by: join on these columns before doing asof join strategy: Join strategy. The default is "backward". * *backward*: selects the last row in the right DataFrame whose "on" key is less than or equal to the left's key. * *forward*: selects the first row in the right DataFrame whose "on" key is greater than or equal to the left's key. * *nearest*: search selects the last row in the right DataFrame whose value is nearest to the left's key. suffix: Suffix to append to columns with a duplicate name. Examples: >>> from datetime import datetime >>> import polars as pl >>> import narwhals as nw >>> data_gdp = { ... "datetime": [ ... datetime(2016, 1, 1), ... datetime(2017, 1, 1), ... datetime(2018, 1, 1), ... datetime(2019, 1, 1), ... datetime(2020, 1, 1), ... ], ... "gdp": [4164, 4411, 4566, 4696, 4827], ... } >>> data_population = { ... "datetime": [ ... datetime(2016, 3, 1), ... datetime(2018, 8, 1), ... datetime(2019, 1, 1), ... ], ... "population": [82.19, 82.66, 83.12], ... } >>> gdp_native = pl.DataFrame(data_gdp) >>> population_native = pl.DataFrame(data_population) >>> gdp = nw.from_native(gdp_native) >>> population = nw.from_native(population_native) >>> population.join_asof(gdp, on="datetime", strategy="backward").to_native() shape: (3, 3) ┌─────────────────────┮────────────┮──────┐ │ datetime ┆ population ┆ gdp │ │ --- ┆ --- ┆ --- │ │ datetime[Ξs] ┆ f64 ┆ i64 │ ╞═════════════════════╩════════════╩══════╡ │ 2016-03-01 00:00:00 ┆ 82.19 ┆ 4164 │ │ 2018-08-01 00:00:00 ┆ 82.66 ┆ 4566 │ │ 2019-01-01 00:00:00 ┆ 83.12 ┆ 4696 │ └─────────────────────â”ī────────────â”ī──────┘ """ return super().join_asof( other, left_on=left_on, right_on=right_on, on=on, by_left=by_left, by_right=by_right, by=by, strategy=strategy, suffix=suffix, ) def lazy(self) -> Self: """Restrict available API methods to lazy-only ones. This is a no-op, and exists only for compatibility with `DataFrame.lazy`. """ return self def unpivot( self, on: str | list[str] | None = None, *, index: str | list[str] | None = None, variable_name: str = "variable", value_name: str = "value", ) -> Self: r"""Unpivot a DataFrame from wide to long format. Optionally leaves identifiers set. This function is useful to massage a DataFrame into a format where one or more columns are identifier variables (index) while all other columns, considered measured variables (on), are "unpivoted" to the row axis leaving just two non-identifier columns, 'variable' and 'value'. Arguments: on: Column(s) to use as values variables; if `on` is empty all columns that are not in `index` will be used. index: Column(s) to use as identifier variables. variable_name: Name to give to the `variable` column. Defaults to "variable". value_name: Name to give to the `value` column. Defaults to "value". Notes: If you're coming from pandas, this is similar to `pandas.DataFrame.melt`, but with `index` replacing `id_vars` and `on` replacing `value_vars`. In other frameworks, you might know this operation as `pivot_longer`. Examples: >>> import duckdb >>> import narwhals as nw >>> df_native = duckdb.sql( ... "SELECT * FROM VALUES ('x', 1, 2), ('y', 3, 4), ('z', 5, 6) df(a, b, c)" ... ) >>> df = nw.from_native(df_native) >>> df.unpivot(on=["b", "c"], index="a").sort("a", "variable").to_native() ┌─────────┮──────────┮───────┐ │ a │ variable │ value │ │ varchar │ varchar │ int32 │ ├─────────┾──────────┾───────â”Ī │ x │ b │ 1 │ │ x │ c │ 2 │ │ y │ b │ 3 │ │ y │ c │ 4 │ │ z │ b │ 5 │ │ z │ c │ 6 │ └─────────â”ī──────────â”ī───────┘ """ return super().unpivot( on=on, index=index, variable_name=variable_name, value_name=value_name ) def explode(self, columns: str | Sequence[str], *more_columns: str) -> Self: """Explode the dataframe to long format by exploding the given columns. Notes: It is possible to explode multiple columns only if these columns have matching element counts. Arguments: columns: Column names. The underlying columns being exploded must be of the `List` data type. *more_columns: Additional names of columns to explode, specified as positional arguments. Examples: >>> import duckdb >>> import narwhals as nw >>> df_native = duckdb.sql( ... "SELECT * FROM VALUES ('x', [1, 2]), ('y', [3, 4]), ('z', [5, 6]) df(a, b)" ... ) >>> df = nw.from_native(df_native) >>> df.explode("b").to_native() ┌─────────┮───────┐ │ a │ b │ │ varchar │ int32 │ ├─────────┾───────â”Ī │ x │ 1 │ │ x │ 2 │ │ y │ 3 │ │ y │ 4 │ │ z │ 5 │ │ z │ 6 │ └─────────â”ī───────┘ """ return super().explode(columns, *more_columns) narwhals-2.11.0/narwhals/dependencies.py0000644000000000000000000004616313615410400015163 0ustar00# pandas / Polars / etc. : if a user passes a dataframe from one of these # libraries, it means they must already have imported the given module. # So, we can just check sys.modules. from __future__ import annotations import sys from typing import TYPE_CHECKING, Any from narwhals._exceptions import issue_warning if TYPE_CHECKING: import cudf import dask.dataframe as dd import duckdb import ibis import modin.pandas as mpd import pandas as pd import polars as pl import pyarrow as pa import pyspark.sql as pyspark_sql from pyspark.sql.connect.dataframe import DataFrame as PySparkConnectDataFrame from typing_extensions import TypeGuard, TypeIs from narwhals._spark_like.dataframe import SQLFrameDataFrame from narwhals.dataframe import DataFrame, LazyFrame from narwhals.series import Series from narwhals.typing import ( IntoDataFrameT, IntoLazyFrameT, IntoSeriesT, PandasLikeDType, _1DArray, _1DArrayInt, _2DArray, _NDArray, _NumpyScalar, _ShapeT, ) # We silently allow these but - given that they claim # to be drop-in replacements for pandas - testing is # their responsibility. IMPORT_HOOKS = frozenset(["fireducks"]) def get_polars() -> Any: """Get Polars module (if already imported - else return None).""" return sys.modules.get("polars", None) def get_pandas() -> Any: """Get pandas module (if already imported - else return None).""" return sys.modules.get("pandas", None) def get_modin() -> Any: # pragma: no cover """Get modin.pandas module (if already imported - else return None).""" return sys.modules.get("modin.pandas", None) def get_cudf() -> Any: """Get cudf module (if already imported - else return None).""" return sys.modules.get("cudf", None) def get_cupy() -> Any: """Get cupy module (if already imported - else return None).""" return sys.modules.get("cupy", None) def get_pyarrow() -> Any: # pragma: no cover """Get pyarrow module (if already imported - else return None).""" return sys.modules.get("pyarrow", None) def get_numpy() -> Any: """Get numpy module (if already imported - else return None).""" return sys.modules.get("numpy", None) def get_dask() -> Any: # pragma: no cover """Get dask (if already imported - else return None).""" return sys.modules.get("dask", None) def get_dask_dataframe() -> Any: """Get dask.dataframe module (if already imported - else return None).""" return sys.modules.get("dask.dataframe", None) def get_duckdb() -> Any: """Get duckdb module (if already imported - else return None).""" return sys.modules.get("duckdb", None) def get_ibis() -> Any: """Get ibis module (if already imported - else return None).""" return sys.modules.get("ibis", None) def get_dask_expr() -> Any: # pragma: no cover """Get dask_expr module (if already imported - else return None).""" if (dd := get_dask_dataframe()) is not None and hasattr(dd, "dask_expr"): return dd.dask_expr return sys.modules.get("dask_expr", None) def get_pyspark() -> Any: # pragma: no cover """Get pyspark module (if already imported - else return None).""" return sys.modules.get("pyspark", None) def get_pyspark_sql() -> Any: """Get pyspark.sql module (if already imported - else return None).""" return sys.modules.get("pyspark.sql", None) def get_pyspark_connect() -> Any: """Get pyspark.sql.connect module (if already imported - else return None).""" return sys.modules.get("pyspark.sql.connect", None) def get_sqlframe() -> Any: """Get sqlframe module (if already imported - else return None).""" return sys.modules.get("sqlframe", None) def _warn_if_narwhals_df_or_lf(df: Any) -> None: if is_narwhals_dataframe(df) or is_narwhals_lazyframe(df): msg = ( f"You passed a `{type(df)}` to `is_pandas_dataframe`.\n\n" "Hint: Instead of e.g. `is_pandas_dataframe(df)`, " "did you mean `is_pandas_dataframe(df.to_native())`?" ) issue_warning(msg, UserWarning) def _warn_if_narwhals_series(ser: Any) -> None: if is_narwhals_series(ser): msg = ( f"You passed a `{type(ser)}` to `is_pandas_series`.\n\n" "Hint: Instead of e.g. `is_pandas_series(ser)`, " "did you mean `is_pandas_series(ser.to_native())`?" ) issue_warning(msg, UserWarning) def is_pandas_dataframe(df: Any) -> TypeIs[pd.DataFrame]: """Check whether `df` is a pandas DataFrame without importing pandas. Warning: This method cannot be called on a Narwhals DataFrame/LazyFrame. """ _warn_if_narwhals_df_or_lf(df) return ((pd := get_pandas()) is not None and isinstance(df, pd.DataFrame)) or any( (mod := sys.modules.get(module_name, None)) is not None and isinstance(df, mod.pandas.DataFrame) for module_name in IMPORT_HOOKS ) def is_pandas_series(ser: Any) -> TypeIs[pd.Series[Any]]: """Check whether `ser` is a pandas Series without importing pandas. Warning: This method cannot be called on Narwhals Series. """ _warn_if_narwhals_series(ser) return ((pd := get_pandas()) is not None and isinstance(ser, pd.Series)) or any( (mod := sys.modules.get(module_name, None)) is not None and isinstance(ser, mod.pandas.Series) for module_name in IMPORT_HOOKS ) def is_pandas_index(index: Any) -> TypeIs[pd.Index[Any]]: """Check whether `index` is a pandas Index without importing pandas.""" return ((pd := get_pandas()) is not None and isinstance(index, pd.Index)) or any( (mod := sys.modules.get(module_name, None)) is not None and isinstance(index, mod.pandas.Index) for module_name in IMPORT_HOOKS ) def is_modin_dataframe(df: Any) -> TypeIs[mpd.DataFrame]: """Check whether `df` is a modin DataFrame without importing modin. Warning: This method cannot be called on a Narwhals DataFrame/LazyFrame. """ _warn_if_narwhals_df_or_lf(df) return (mpd := get_modin()) is not None and isinstance(df, mpd.DataFrame) def is_modin_series(ser: Any) -> TypeIs[mpd.Series]: """Check whether `ser` is a modin Series without importing modin. Warning: This method cannot be called on Narwhals Series. """ _warn_if_narwhals_series(ser) return (mpd := get_modin()) is not None and isinstance(ser, mpd.Series) def is_modin_index(index: Any) -> TypeIs[mpd.Index[Any]]: # pragma: no cover """Check whether `index` is a modin Index without importing modin.""" return (mpd := get_modin()) is not None and isinstance(index, mpd.Index) def is_cudf_dataframe(df: Any) -> TypeIs[cudf.DataFrame]: """Check whether `df` is a cudf DataFrame without importing cudf. Warning: This method cannot be called on a Narwhals DataFrame/LazyFrame. """ _warn_if_narwhals_df_or_lf(df) return (cudf := get_cudf()) is not None and isinstance(df, cudf.DataFrame) def is_cudf_series(ser: Any) -> TypeIs[cudf.Series[Any]]: """Check whether `ser` is a cudf Series without importing cudf. Warning: This method cannot be called on Narwhals Series. """ _warn_if_narwhals_series(ser) return (cudf := get_cudf()) is not None and isinstance(ser, cudf.Series) def is_cudf_index(index: Any) -> TypeIs[cudf.Index]: """Check whether `index` is a cudf Index without importing cudf.""" return (cudf := get_cudf()) is not None and isinstance( index, cudf.Index ) # pragma: no cover def is_cupy_scalar(obj: Any) -> bool: return ( (cupy := get_cupy()) is not None and isinstance(obj, cupy.ndarray) and obj.size == 1 ) # pragma: no cover def is_dask_dataframe(df: Any) -> TypeIs[dd.DataFrame]: """Check whether `df` is a Dask DataFrame without importing Dask. Warning: This method cannot be called on a Narwhals DataFrame/LazyFrame. """ _warn_if_narwhals_df_or_lf(df) return (dd := get_dask_dataframe()) is not None and isinstance(df, dd.DataFrame) def is_duckdb_relation(df: Any) -> TypeIs[duckdb.DuckDBPyRelation]: """Check whether `df` is a DuckDB Relation without importing DuckDB. Warning: This method cannot be called on Narwhals DataFrame/LazyFrame. """ _warn_if_narwhals_df_or_lf(df) return (duckdb := get_duckdb()) is not None and isinstance( df, duckdb.DuckDBPyRelation ) def is_ibis_table(df: Any) -> TypeIs[ibis.Table]: """Check whether `df` is a Ibis Table without importing Ibis. Warning: This method cannot be called on Narwhals DataFrame/LazyFrame. """ _warn_if_narwhals_df_or_lf(df) return (ibis := get_ibis()) is not None and isinstance(df, ibis.expr.types.Table) def is_polars_dataframe(df: Any) -> TypeIs[pl.DataFrame]: """Check whether `df` is a Polars DataFrame without importing Polars. Warning: This method cannot be called on a Narwhals DataFrame/LazyFrame. """ _warn_if_narwhals_df_or_lf(df) return (pl := get_polars()) is not None and isinstance(df, pl.DataFrame) def is_polars_lazyframe(df: Any) -> TypeIs[pl.LazyFrame]: """Check whether `df` is a Polars LazyFrame without importing Polars. Warning: This method cannot be called on Narwhals DataFrame/LazyFrame. """ _warn_if_narwhals_df_or_lf(df) return (pl := get_polars()) is not None and isinstance(df, pl.LazyFrame) def is_polars_series(ser: Any) -> TypeIs[pl.Series]: """Check whether `ser` is a Polars Series without importing Polars. Warning: This method cannot be called on Narwhals Series. """ _warn_if_narwhals_series(ser) return (pl := get_polars()) is not None and isinstance(ser, pl.Series) def is_polars_schema(obj: Any) -> TypeIs[pl.Schema]: return ( bool(pl := get_polars()) and hasattr(pl, "Schema") and isinstance(obj, pl.Schema) ) # NOTE: For `pl.Schema` only instantiated dtypes are expected def is_polars_data_type(obj: Any) -> TypeIs[pl.DataType]: return bool(pl := get_polars()) and isinstance(obj, pl.DataType) def is_pyarrow_chunked_array(ser: Any) -> TypeIs[pa.ChunkedArray[Any]]: """Check whether `ser` is a PyArrow ChunkedArray without importing PyArrow. Warning: This method cannot be called on Narwhals Series. """ _warn_if_narwhals_series(ser) return (pa := get_pyarrow()) is not None and isinstance(ser, pa.ChunkedArray) def is_pyarrow_table(df: Any) -> TypeIs[pa.Table]: """Check whether `df` is a PyArrow Table without importing PyArrow. Warning: This method cannot be called on Narwhals DataFrame/LazyFrame. """ _warn_if_narwhals_df_or_lf(df) return (pa := get_pyarrow()) is not None and isinstance(df, pa.Table) def is_pyarrow_scalar(obj: Any) -> TypeIs[pa.Scalar[Any]]: return (pa := get_pyarrow()) is not None and isinstance(obj, pa.Scalar) def is_pyarrow_schema(obj: Any) -> TypeIs[pa.Schema]: return bool(pa := get_pyarrow()) and isinstance(obj, pa.Schema) def is_pyarrow_data_type(obj: Any) -> TypeIs[pa.DataType]: return bool(pa := get_pyarrow()) and isinstance(obj, pa.DataType) def is_pyspark_dataframe(df: Any) -> TypeIs[pyspark_sql.DataFrame]: """Check whether `df` is a PySpark DataFrame without importing PySpark. Warning: This method cannot be called on a Narwhals DataFrame/LazyFrame. """ _warn_if_narwhals_df_or_lf(df) return bool( (pyspark_sql := get_pyspark_sql()) is not None and isinstance(df, pyspark_sql.DataFrame) ) def is_pyspark_connect_dataframe(df: Any) -> TypeIs[PySparkConnectDataFrame]: """Check whether `df` is a PySpark Connect DataFrame without importing PySpark. Warning: This method cannot be called on a Narwhals DataFrame/LazyFrame. """ _warn_if_narwhals_df_or_lf(df) if get_pyspark_connect() is not None: # pragma: no cover try: from pyspark.sql.connect.dataframe import DataFrame except ImportError: return False return isinstance(df, DataFrame) return False def is_sqlframe_dataframe(df: Any) -> TypeIs[SQLFrameDataFrame]: """Check whether `df` is a SQLFrame DataFrame without importing SQLFrame. Warning: This method cannot be called on a Narwhals DataFrame/LazyFrame. """ _warn_if_narwhals_df_or_lf(df) if get_sqlframe() is not None: from sqlframe.base.dataframe import BaseDataFrame return isinstance(df, BaseDataFrame) return False # pragma: no cover def is_numpy_array(arr: Any | _NDArray[_ShapeT]) -> TypeIs[_NDArray[_ShapeT]]: """Check whether `arr` is a NumPy Array without importing NumPy.""" return (np := get_numpy()) is not None and isinstance(arr, np.ndarray) def is_numpy_array_1d(arr: Any) -> TypeIs[_1DArray]: """Check whether `arr` is a 1D NumPy Array without importing NumPy.""" return is_numpy_array(arr) and arr.ndim == 1 def is_numpy_array_1d_int(arr: Any) -> TypeIs[_1DArrayInt]: return ( (np := get_numpy()) and is_numpy_array_1d(arr) and np.issubdtype(arr.dtype, np.integer) ) def is_numpy_array_2d(arr: Any) -> TypeIs[_2DArray]: """Check whether `arr` is a 2D NumPy Array without importing NumPy.""" return is_numpy_array(arr) and arr.ndim == 2 def is_numpy_scalar(scalar: Any) -> TypeGuard[_NumpyScalar]: """Check whether `scalar` is a NumPy Scalar without importing NumPy.""" # NOTE: Needs to stay as `TypeGuard` # - Used in `Series.__getitem__`, but not annotated # - `TypeGuard` is *hiding* that the check introduces an intersection return (np := get_numpy()) is not None and isinstance(scalar, np.generic) def is_pandas_like_dataframe(df: Any) -> bool: """Check whether `df` is a pandas-like DataFrame without doing any imports. By "pandas-like", we mean: pandas, Modin, cuDF. Warning: This method cannot be called on a Narwhals DataFrame/LazyFrame. """ _warn_if_narwhals_df_or_lf(df) return is_pandas_dataframe(df) or is_modin_dataframe(df) or is_cudf_dataframe(df) def is_pandas_like_series(ser: Any) -> bool: """Check whether `ser` is a pandas-like Series without doing any imports. By "pandas-like", we mean: pandas, Modin, cuDF. Warning: This method cannot be called on Narwhals Series. """ _warn_if_narwhals_series(ser) return is_pandas_series(ser) or is_modin_series(ser) or is_cudf_series(ser) def is_pandas_like_index(index: Any) -> bool: """Check whether `index` is a pandas-like Index without doing any imports. By "pandas-like", we mean: pandas, Modin, cuDF. """ return ( is_pandas_index(index) or is_modin_index(index) or is_cudf_index(index) ) # pragma: no cover def is_pandas_like_dtype(obj: Any) -> TypeIs[PandasLikeDType]: return bool(pd := get_pandas()) and isinstance( obj, (pd.api.extensions.ExtensionDtype, get_numpy().dtype) ) def is_cudf_dtype( obj: Any, ) -> TypeIs[pd.api.extensions.ExtensionDtype]: # pragma: no cover return ( bool(pd := get_pandas()) and isinstance(obj, (pd.api.extensions.ExtensionDtype)) and hasattr(obj, "to_arrow") ) def is_into_series(native_series: Any | IntoSeriesT) -> TypeIs[IntoSeriesT]: """Check whether `native_series` can be converted to a Narwhals Series. Arguments: native_series: The object to check. Examples: >>> import pandas as pd >>> import polars as pl >>> import numpy as np >>> import narwhals as nw >>> s_pd = pd.Series([1, 2, 3]) >>> s_pl = pl.Series([1, 2, 3]) >>> np_arr = np.array([1, 2, 3]) >>> nw.dependencies.is_into_series(s_pd) True >>> nw.dependencies.is_into_series(s_pl) True >>> nw.dependencies.is_into_series(np_arr) False """ from narwhals.series import Series return ( isinstance(native_series, Series) or hasattr(native_series, "__narwhals_series__") or is_polars_series(native_series) or is_pyarrow_chunked_array(native_series) or is_pandas_like_series(native_series) ) def is_into_dataframe(native_dataframe: Any | IntoDataFrameT) -> TypeIs[IntoDataFrameT]: """Check whether `native_dataframe` can be converted to a Narwhals DataFrame. Arguments: native_dataframe: The object to check. Examples: >>> import pandas as pd >>> import polars as pl >>> import numpy as np >>> from narwhals.dependencies import is_into_dataframe >>> df_pd = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) >>> df_pl = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) >>> np_arr = np.array([[1, 4], [2, 5], [3, 6]]) >>> is_into_dataframe(df_pd) True >>> is_into_dataframe(df_pl) True >>> is_into_dataframe(np_arr) False """ from narwhals.dataframe import DataFrame return ( isinstance(native_dataframe, DataFrame) or hasattr(native_dataframe, "__narwhals_dataframe__") or is_polars_dataframe(native_dataframe) or is_pyarrow_table(native_dataframe) or is_pandas_like_dataframe(native_dataframe) ) def is_narwhals_dataframe( df: DataFrame[IntoDataFrameT] | Any, ) -> TypeIs[DataFrame[IntoDataFrameT]]: """Check whether `df` is a Narwhals DataFrame. This is useful if you expect a user to pass in a Narwhals DataFrame directly, and you want to catch both `narwhals.DataFrame` and `narwhals.stable.v1.DataFrame`. """ from narwhals.dataframe import DataFrame return isinstance(df, DataFrame) def is_narwhals_lazyframe( lf: Any | LazyFrame[IntoLazyFrameT], ) -> TypeIs[LazyFrame[IntoLazyFrameT]]: """Check whether `lf` is a Narwhals LazyFrame. This is useful if you expect a user to pass in a Narwhals LazyFrame directly, and you want to catch both `narwhals.LazyFrame` and `narwhals.stable.v1.LazyFrame`. """ from narwhals.dataframe import LazyFrame return isinstance(lf, LazyFrame) def is_narwhals_series(ser: Any | Series[IntoSeriesT]) -> TypeIs[Series[IntoSeriesT]]: """Check whether `ser` is a Narwhals Series. This is useful if you expect a user to pass in a Narwhals Series directly, and you want to catch both `narwhals.Series` and `narwhals.stable.v1.Series`. """ from narwhals.series import Series return isinstance(ser, Series) def is_narwhals_series_int(ser: Any | Series[IntoSeriesT]) -> TypeIs[Series[IntoSeriesT]]: return is_narwhals_series(ser) and ser.dtype.is_integer() __all__ = [ "get_cudf", "get_ibis", "get_modin", "get_numpy", "get_pandas", "get_polars", "get_pyarrow", "is_cudf_dataframe", "is_cudf_series", "is_dask_dataframe", "is_ibis_table", "is_into_dataframe", "is_into_series", "is_modin_dataframe", "is_modin_series", "is_narwhals_dataframe", "is_narwhals_lazyframe", "is_narwhals_series", "is_numpy_array", "is_pandas_dataframe", "is_pandas_index", "is_pandas_like_dataframe", "is_pandas_like_series", "is_pandas_series", "is_polars_dataframe", "is_polars_lazyframe", "is_polars_series", "is_pyarrow_chunked_array", "is_pyarrow_table", ] narwhals-2.11.0/narwhals/dtypes.py0000644000000000000000000007446713615410400014055 0ustar00from __future__ import annotations import enum from collections import OrderedDict from collections.abc import Iterable, Mapping from datetime import timezone from itertools import starmap from typing import TYPE_CHECKING from narwhals._utils import ( _DeferredIterable, isinstance_or_issubclass, qualified_type_name, ) if TYPE_CHECKING: from collections.abc import Iterator, Sequence from typing import Any import _typeshed from typing_extensions import Self, TypeIs from narwhals.typing import IntoDType, TimeUnit def _validate_dtype(dtype: DType | type[DType]) -> None: if not isinstance_or_issubclass(dtype, DType): msg = ( f"Expected Narwhals dtype, got: {type(dtype)}.\n\n" "Hint: if you were trying to cast to a type, use e.g. nw.Int64 instead of 'int64'." ) raise TypeError(msg) def _is_into_dtype(obj: Any) -> TypeIs[IntoDType]: return isinstance(obj, DType) or ( isinstance(obj, DTypeClass) and not issubclass(obj, NestedType) ) def _is_nested_type(obj: Any) -> TypeIs[type[NestedType]]: return isinstance(obj, DTypeClass) and issubclass(obj, NestedType) def _validate_into_dtype(dtype: Any) -> None: if not _is_into_dtype(dtype): if _is_nested_type(dtype): name = f"nw.{dtype.__name__}" msg = ( f"{name!r} is not valid in this context.\n\n" f"Hint: instead of:\n\n" f" {name}\n\n" "use:\n\n" f" {name}(...)" ) else: msg = f"Expected Narwhals dtype, got: {qualified_type_name(dtype)!r}." raise TypeError(msg) class DTypeClass(type): """Metaclass for DType classes. - Nicely print classes. - Ensure [`__slots__`] are always defined to prevent `__dict__` creation (empty by default). [`__slots__`]: https://docs.python.org/3/reference/datamodel.html#object.__slots__ """ def __repr__(cls) -> str: return cls.__name__ # https://github.com/python/typeshed/blob/776508741d76b58f9dcb2aaf42f7d4596a48d580/stdlib/abc.pyi#L13-L19 # https://github.com/python/typeshed/blob/776508741d76b58f9dcb2aaf42f7d4596a48d580/stdlib/_typeshed/__init__.pyi#L36-L40 # https://github.com/astral-sh/ruff/issues/8353#issuecomment-1786238311 # https://docs.python.org/3/reference/datamodel.html#creating-the-class-object def __new__( metacls: type[_typeshed.Self], cls_name: str, bases: tuple[type, ...], namespace: dict[str, Any], /, **kwds: Any, ) -> _typeshed.Self: namespace.setdefault("__slots__", ()) return super().__new__(metacls, cls_name, bases, namespace, **kwds) # type: ignore[no-any-return, misc] class DType(metaclass=DTypeClass): """Base class for all Narwhals data types.""" __slots__ = () # NOTE: Keep this one defined manually for the type checker def __repr__(self) -> str: return self.__class__.__qualname__ @classmethod def base_type(cls) -> type[Self]: """Return this DType's fundamental/root type class. Examples: >>> import narwhals as nw >>> nw.Datetime("us").base_type() Datetime >>> nw.String.base_type() String >>> nw.List(nw.Int64).base_type() List """ return cls @classmethod def is_numeric(cls: type[Self]) -> bool: """Check whether the data type is a numeric type.""" return issubclass(cls, NumericType) @classmethod def is_integer(cls: type[Self]) -> bool: """Check whether the data type is an integer type.""" return issubclass(cls, IntegerType) @classmethod def is_signed_integer(cls: type[Self]) -> bool: """Check whether the data type is a signed integer type.""" return issubclass(cls, SignedIntegerType) @classmethod def is_unsigned_integer(cls: type[Self]) -> bool: """Check whether the data type is an unsigned integer type.""" return issubclass(cls, UnsignedIntegerType) @classmethod def is_float(cls: type[Self]) -> bool: """Check whether the data type is a floating point type.""" return issubclass(cls, FloatType) @classmethod def is_decimal(cls: type[Self]) -> bool: """Check whether the data type is a decimal type.""" return issubclass(cls, Decimal) @classmethod def is_temporal(cls: type[Self]) -> bool: """Check whether the data type is a temporal type.""" return issubclass(cls, TemporalType) @classmethod def is_nested(cls: type[Self]) -> bool: """Check whether the data type is a nested type.""" return issubclass(cls, NestedType) @classmethod def is_boolean(cls: type[Self]) -> bool: """Check whether the data type is a boolean type.""" return issubclass(cls, Boolean) def __eq__(self, other: DType | type[DType]) -> bool: # type: ignore[override] """Check if this DType is equivalent to another DType. Examples: >>> import narwhals as nw >>> nw.String() == nw.String() True >>> nw.String() == nw.String True >>> nw.Int16() == nw.Int32 False >>> nw.Boolean() == nw.Int8 False >>> nw.Date() == nw.Datetime False """ return isinstance_or_issubclass(other, type(self)) def __hash__(self) -> int: return hash(self.__class__) class NumericType(DType): """Base class for numeric data types.""" class IntegerType(NumericType): """Base class for integer data types.""" class SignedIntegerType(IntegerType): """Base class for signed integer data types.""" class UnsignedIntegerType(IntegerType): """Base class for unsigned integer data types.""" class FloatType(NumericType): """Base class for float data types.""" class TemporalType(DType): """Base class for temporal data types.""" class NestedType(DType): """Base class for nested data types.""" class Decimal(NumericType): """Decimal type. Examples: >>> import polars as pl >>> import narwhals as nw >>> s = pl.Series(["1.5"], dtype=pl.Decimal) >>> nw.from_native(s, series_only=True).dtype Decimal """ class Int128(SignedIntegerType): """128-bit signed integer type. Examples: >>> import polars as pl >>> import pyarrow as pa >>> import narwhals as nw >>> import duckdb >>> s_native = pl.Series([2, 1, 3, 7]) >>> s = nw.from_native(s_native, series_only=True) >>> df_native = pa.table({"a": [2, 1, 3, 7]}) >>> rel = duckdb.sql(" SELECT CAST (a AS INT128) AS a FROM df_native ") >>> s.cast(nw.Int128).dtype Int128 >>> nw.from_native(rel).collect_schema()["a"] Int128 """ class Int64(SignedIntegerType): """64-bit signed integer type. Examples: >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series([2, 1, 3, 7]) >>> s = nw.from_native(s_native, series_only=True) >>> s.cast(nw.Int64).dtype Int64 """ class Int32(SignedIntegerType): """32-bit signed integer type. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> s_native = pa.chunked_array([[2, 1, 3, 7]]) >>> s = nw.from_native(s_native, series_only=True) >>> s.cast(nw.Int32).dtype Int32 """ class Int16(SignedIntegerType): """16-bit signed integer type. Examples: >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series([2, 1, 3, 7]) >>> s = nw.from_native(s_native, series_only=True) >>> s.cast(nw.Int16).dtype Int16 """ class Int8(SignedIntegerType): """8-bit signed integer type. Examples: >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series([2, 1, 3, 7]) >>> s = nw.from_native(s_native, series_only=True) >>> s.cast(nw.Int8).dtype Int8 """ class UInt128(UnsignedIntegerType): """128-bit unsigned integer type. Examples: >>> import pandas as pd >>> import narwhals as nw >>> import duckdb >>> df_native = pd.DataFrame({"a": [2, 1, 3, 7]}) >>> rel = duckdb.sql(" SELECT CAST (a AS UINT128) AS a FROM df_native ") >>> nw.from_native(rel).collect_schema()["a"] UInt128 """ class UInt64(UnsignedIntegerType): """64-bit unsigned integer type. Examples: >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series([2, 1, 3, 7]) >>> s = nw.from_native(s_native, series_only=True) >>> s.cast(nw.UInt64).dtype UInt64 """ class UInt32(UnsignedIntegerType): """32-bit unsigned integer type. Examples: >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series([2, 1, 3, 7]) >>> s = nw.from_native(s_native, series_only=True) >>> s.cast(nw.UInt32).dtype UInt32 """ class UInt16(UnsignedIntegerType): """16-bit unsigned integer type. Examples: >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series([2, 1, 3, 7]) >>> s = nw.from_native(s_native, series_only=True) >>> s.cast(nw.UInt16).dtype UInt16 """ class UInt8(UnsignedIntegerType): """8-bit unsigned integer type. Examples: >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series([2, 1, 3, 7]) >>> s = nw.from_native(s_native, series_only=True) >>> s.cast(nw.UInt8).dtype UInt8 """ class Float64(FloatType): """64-bit floating point type. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> s_native = pa.chunked_array([[0.001, 0.1, 0.01, 0.1]]) >>> s = nw.from_native(s_native, series_only=True) >>> s.cast(nw.Float64).dtype Float64 """ class Float32(FloatType): """32-bit floating point type. Examples: >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series([0.001, 0.1, 0.01, 0.1]) >>> s = nw.from_native(s_native, series_only=True) >>> s.cast(nw.Float32).dtype Float32 """ class String(DType): """UTF-8 encoded string type. Examples: >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series(["beluga", "narwhal", "orca", "vaquita"]) >>> nw.from_native(s_native, series_only=True).dtype String """ class Boolean(DType): """Boolean type. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> s_native = pa.chunked_array([[True, False, False, True]]) >>> nw.from_native(s_native, series_only=True).dtype Boolean """ class Object(DType): """Data type for wrapping arbitrary Python objects. Examples: >>> import pandas as pd >>> import narwhals as nw >>> class Foo: ... >>> s_native = pd.Series([Foo(), Foo()]) >>> nw.from_native(s_native, series_only=True).dtype Object """ class Unknown(DType): """Type representing DataType values that could not be determined statically. Examples: >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series(pd.period_range("2000-01", periods=4, freq="M")) >>> nw.from_native(s_native, series_only=True).dtype Unknown """ class _DatetimeMeta(DTypeClass): @property def time_unit(cls) -> TimeUnit: """Unit of time. Defaults to `'us'` (microseconds).""" return "us" @property def time_zone(cls) -> str | None: """Time zone string. Defaults to `None`.""" return None class Datetime(TemporalType, metaclass=_DatetimeMeta): """Data type representing a calendar date and time of day. Arguments: time_unit: Unit of time. Defaults to `'us'` (microseconds). time_zone: Time zone string, as defined in zoneinfo (to see valid strings run `import zoneinfo; zoneinfo.available_timezones()` for a full list). Notes: Adapted from [Polars implementation](https://github.com/pola-rs/polars/blob/py-1.7.1/py-polars/polars/datatypes/classes.py#L398-L457) Examples: >>> from datetime import datetime, timedelta >>> import polars as pl >>> import narwhals as nw >>> s_native = ( ... pl.Series([datetime(2024, 12, 9) + timedelta(days=n) for n in range(5)]) ... .cast(pl.Datetime("ms")) ... .dt.replace_time_zone("Africa/Accra") ... ) >>> nw.from_native(s_native, series_only=True).dtype Datetime(time_unit='ms', time_zone='Africa/Accra') """ __slots__ = ("time_unit", "time_zone") def __init__( self, time_unit: TimeUnit = "us", time_zone: str | timezone | None = None ) -> None: if time_unit not in {"s", "ms", "us", "ns"}: msg = ( "invalid `time_unit`" f"\n\nExpected one of {{'ns','us','ms', 's'}}, got {time_unit!r}." ) raise ValueError(msg) if isinstance(time_zone, timezone): time_zone = str(time_zone) self.time_unit: TimeUnit = time_unit """Unit of time.""" self.time_zone: str | None = time_zone """Time zone string, as defined in zoneinfo. Notes: To see valid strings run `import zoneinfo; zoneinfo.available_timezones()` for a full list. """ def __eq__(self, other: DType | type[DType]) -> bool: # type: ignore[override] """Check if this Datetime is equivalent to another DType. Examples: >>> import narwhals as nw >>> nw.Datetime("s") == nw.Datetime("s") True >>> nw.Datetime() == nw.Datetime("us") True >>> nw.Datetime("us") == nw.Datetime("ns") False >>> nw.Datetime("us", "UTC") == nw.Datetime(time_unit="us", time_zone="UTC") True >>> nw.Datetime(time_zone="UTC") == nw.Datetime(time_zone="EST") False >>> nw.Datetime() == nw.Duration() False >>> nw.Datetime("ms") == nw.Datetime True """ if type(other) is _DatetimeMeta: return True if isinstance(other, self.__class__): return self.time_unit == other.time_unit and self.time_zone == other.time_zone return False # pragma: no cover def __hash__(self) -> int: # pragma: no cover return hash((self.__class__, self.time_unit, self.time_zone)) def __repr__(self) -> str: # pragma: no cover class_name = self.__class__.__name__ return f"{class_name}(time_unit={self.time_unit!r}, time_zone={self.time_zone!r})" class _DurationMeta(DTypeClass): @property def time_unit(cls) -> TimeUnit: """Unit of time. Defaults to `'us'` (microseconds).""" return "us" class Duration(TemporalType, metaclass=_DurationMeta): """Data type representing a time duration. Arguments: time_unit: Unit of time. Defaults to `'us'` (microseconds). Notes: Adapted from [Polars implementation](https://github.com/pola-rs/polars/blob/py-1.7.1/py-polars/polars/datatypes/classes.py#L460-L502) Examples: >>> from datetime import timedelta >>> import pyarrow as pa >>> import narwhals as nw >>> s_native = pa.chunked_array( ... [[timedelta(seconds=d) for d in range(1, 4)]], type=pa.duration("ms") ... ) >>> nw.from_native(s_native, series_only=True).dtype Duration(time_unit='ms') """ __slots__ = ("time_unit",) def __init__(self, time_unit: TimeUnit = "us") -> None: if time_unit not in {"s", "ms", "us", "ns"}: msg = ( "invalid `time_unit`" f"\n\nExpected one of {{'ns','us','ms', 's'}}, got {time_unit!r}." ) raise ValueError(msg) self.time_unit: TimeUnit = time_unit """Unit of time.""" def __eq__(self, other: DType | type[DType]) -> bool: # type: ignore[override] """Check if this Duration is equivalent to another DType. Examples: >>> import narwhals as nw >>> nw.Duration("us") == nw.Duration("us") True >>> nw.Duration() == nw.Duration("us") True >>> nw.Duration("us") == nw.Duration("ns") False >>> nw.Duration() == nw.Datetime() False >>> nw.Duration("ms") == nw.Duration True """ if type(other) is _DurationMeta: return True if isinstance(other, self.__class__): return self.time_unit == other.time_unit return False # pragma: no cover def __hash__(self) -> int: # pragma: no cover return hash((self.__class__, self.time_unit)) def __repr__(self) -> str: # pragma: no cover class_name = self.__class__.__name__ return f"{class_name}(time_unit={self.time_unit!r})" class Categorical(DType): """A categorical encoding of a set of strings. Examples: >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series(["beluga", "narwhal", "orca"]) >>> nw.from_native(s_native, series_only=True).cast(nw.Categorical).dtype Categorical """ class Enum(DType): """A fixed categorical encoding of a unique set of strings. Polars has an Enum data type. In pandas, ordered categories get mapped to Enum. PyArrow has no Enum equivalent. Examples: >>> import narwhals as nw >>> nw.Enum(["beluga", "narwhal", "orca"]) Enum(categories=['beluga', 'narwhal', 'orca']) """ __slots__ = ("_cached_categories", "_delayed_categories") def __init__(self, categories: Iterable[str] | type[enum.Enum]) -> None: self._delayed_categories: _DeferredIterable[str] | None = None self._cached_categories: tuple[str, ...] | None = None if isinstance(categories, _DeferredIterable): self._delayed_categories = categories elif isinstance(categories, type) and issubclass(categories, enum.Enum): self._cached_categories = tuple(member.value for member in categories) else: self._cached_categories = tuple(categories) @property def categories(self) -> tuple[str, ...]: """The categories in the dataset.""" if (cached := self._cached_categories) is not None: return cached if (delayed := self._delayed_categories) is not None: self._cached_categories = delayed.to_tuple() return self._cached_categories msg = f"Internal structure of {type(self).__name__!r} is invalid." # pragma: no cover raise TypeError(msg) # pragma: no cover def __eq__(self, other: DType | type[DType]) -> bool: # type: ignore[override] """Check if this Enum is equivalent to another DType. Examples: >>> import narwhals as nw >>> nw.Enum(["a", "b", "c"]) == nw.Enum(["a", "b", "c"]) True >>> import polars as pl >>> categories = pl.Series(["a", "b", "c"]) >>> nw.Enum(["a", "b", "c"]) == nw.Enum(categories) True >>> nw.Enum(["a", "b", "c"]) == nw.Enum(["b", "a", "c"]) False >>> nw.Enum(["a", "b", "c"]) == nw.Enum(["a"]) False >>> nw.Enum(["a", "b", "c"]) == nw.Categorical False >>> nw.Enum(["a", "b", "c"]) == nw.Enum True """ if type(other) is DTypeClass: return other is Enum return isinstance(other, type(self)) and self.categories == other.categories def __hash__(self) -> int: return hash((self.__class__, tuple(self.categories))) def __repr__(self) -> str: return f"{type(self).__name__}(categories={list(self.categories)!r})" class Field: """Definition of a single field within a `Struct` DType. Arguments: name: The name of the field within its parent `Struct`. dtype: The `DType` of the field's values. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> data = [{"a": 1, "b": ["narwhal", "beluga"]}, {"a": 2, "b": ["orca"]}] >>> ser_pa = pa.chunked_array([data]) >>> nw.from_native(ser_pa, series_only=True).dtype.fields [Field('a', Int64), Field('b', List(String))] """ __slots__ = ("dtype", "name") name: str """The name of the field within its parent `Struct`.""" dtype: IntoDType """The `DType` of the field's values.""" def __init__(self, name: str, dtype: IntoDType) -> None: self.name = name self.dtype = dtype def __eq__(self, other: Field) -> bool: # type: ignore[override] """Check if this Field is equivalent to another Field. Two fields are equivalent if they have the same name and the same dtype. Examples: >>> import narwhals as nw >>> nw.Field("a", nw.String) == nw.Field("a", nw.String()) True >>> nw.Field("a", nw.String) == nw.Field("a", nw.String) True >>> nw.Field("a", nw.String) == nw.Field("a", nw.Datetime) False >>> nw.Field("a", nw.String) == nw.Field("b", nw.String) False >>> nw.Field("a", nw.String) == nw.String False """ return ( isinstance(other, Field) and (self.name == other.name) and (self.dtype == other.dtype) ) def __hash__(self) -> int: return hash((self.name, self.dtype)) def __repr__(self) -> str: class_name = self.__class__.__name__ return f"{class_name}({self.name!r}, {self.dtype})" class Struct(NestedType): """Struct composite type. Arguments: fields: The fields that make up the struct. Can be either a sequence of Field objects or a mapping of column names to data types. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> s_native = pa.chunked_array( ... [[{"a": 1, "b": ["narwhal", "beluga"]}, {"a": 2, "b": ["orca"]}]] ... ) >>> nw.from_native(s_native, series_only=True).dtype Struct({'a': Int64, 'b': List(String)}) """ __slots__ = ("fields",) fields: list[Field] """The fields that make up the struct.""" def __init__(self, fields: Sequence[Field] | Mapping[str, IntoDType]) -> None: if isinstance(fields, Mapping): self.fields = list(starmap(Field, fields.items())) else: self.fields = list(fields) def __eq__(self, other: DType | type[DType]) -> bool: # type: ignore[override] """Check if this Struct is equivalent to another DType. Examples: >>> import narwhals as nw >>> nw.Struct({"a": nw.Int64}) == nw.Struct({"a": nw.Int64}) True >>> nw.Struct({"a": nw.Int64}) == nw.Struct({"a": nw.Boolean}) False >>> nw.Struct({"a": nw.Int64}) == nw.Struct({"b": nw.Int64}) False >>> nw.Struct({"a": nw.Int64}) == nw.Struct([nw.Field("a", nw.Int64)]) True If a parent type is not specific about its inner type, we infer it as equal >>> nw.Struct({"a": nw.Int64}) == nw.Struct True """ if type(other) is DTypeClass and issubclass(other, self.__class__): return True if isinstance(other, self.__class__): return self.fields == other.fields return False def __hash__(self) -> int: return hash((self.__class__, tuple(self.fields))) def __iter__(self) -> Iterator[tuple[str, IntoDType]]: # pragma: no cover for fld in self.fields: yield fld.name, fld.dtype def __reversed__(self) -> Iterator[tuple[str, IntoDType]]: for fld in reversed(self.fields): yield fld.name, fld.dtype def __repr__(self) -> str: class_name = self.__class__.__name__ return f"{class_name}({dict(self)})" def to_schema(self) -> OrderedDict[str, IntoDType]: """Return Struct dtype as a schema dict.""" return OrderedDict(self) class List(NestedType): """Variable length list type. Examples: >>> import pandas as pd >>> import pyarrow as pa >>> import narwhals as nw >>> s_native = pd.Series( ... [["narwhal", "orca"], ["beluga", "vaquita"]], ... dtype=pd.ArrowDtype(pa.large_list(pa.large_string())), ... ) >>> nw.from_native(s_native, series_only=True).dtype List(String) """ __slots__ = ("inner",) inner: IntoDType """The DType of the values within each list.""" def __init__(self, inner: IntoDType) -> None: self.inner = inner def __eq__(self, other: DType | type[DType]) -> bool: # type: ignore[override] """Check if this List is equivalent to another DType. Examples: >>> import narwhals as nw >>> nw.List(nw.Int64) == nw.List(nw.Int64) True >>> nw.List(nw.Int64) == nw.List(nw.Float32) False If a parent type is not specific about its inner type, we infer it as equal >>> nw.List(nw.Int64) == nw.List True """ if type(other) is DTypeClass and issubclass(other, self.__class__): return True if isinstance(other, self.__class__): return self.inner == other.inner return False def __hash__(self) -> int: return hash((self.__class__, self.inner)) def __repr__(self) -> str: class_name = self.__class__.__name__ return f"{class_name}({self.inner!r})" class Array(NestedType): """Fixed length list type. Arguments: inner: The datatype of the values within each array. shape: The shape of the arrays. Examples: >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series([[1, 2], [3, 4], [5, 6]], dtype=pl.Array(pl.Int32, 2)) >>> nw.from_native(s_native, series_only=True).dtype Array(Int32, shape=(2,)) """ __slots__ = ("inner", "shape", "size") inner: IntoDType """The DType of the values within each array.""" size: int """The size of the Array.""" shape: tuple[int, ...] """The shape of the arrays.""" def __init__(self, inner: IntoDType, shape: int | tuple[int, ...]) -> None: inner_shape: tuple[int, ...] = inner.shape if isinstance(inner, Array) else () if isinstance(shape, int): self.inner = inner self.size = shape self.shape = (shape, *inner_shape) elif isinstance(shape, tuple) and len(shape) != 0 and isinstance(shape[0], int): if len(shape) > 1: inner = Array(inner, shape[1:]) self.inner = inner self.size = shape[0] self.shape = shape + inner_shape else: msg = f"invalid input for shape: {shape!r}" raise TypeError(msg) def __eq__(self, other: DType | type[DType]) -> bool: # type: ignore[override] """Check if this Array is equivalent to another DType. Examples: >>> import narwhals as nw >>> nw.Array(nw.Int64, 2) == nw.Array(nw.Int64, 2) True >>> nw.Array(nw.Int64, 2) == nw.Array(nw.String, 2) False >>> nw.Array(nw.Int64, 2) == nw.Array(nw.Int64, 4) False If a parent type is not specific about its inner type, we infer it as equal >>> nw.Array(nw.Int64, 2) == nw.Array True """ if type(other) is DTypeClass and issubclass(other, self.__class__): return True if isinstance(other, self.__class__): if self.shape != other.shape: return False return self.inner == other.inner return False def __hash__(self) -> int: return hash((self.__class__, self.inner, self.shape)) def __repr__(self) -> str: # Get leaf type dtype_ = self for _ in self.shape: dtype_ = dtype_.inner # type: ignore[assignment] class_name = self.__class__.__name__ return f"{class_name}({dtype_!r}, shape={self.shape})" class Date(TemporalType): """Data type representing a calendar date. Examples: >>> from datetime import date, timedelta >>> import pyarrow as pa >>> import narwhals as nw >>> s_native = pa.chunked_array( ... [[date(2024, 12, 1) + timedelta(days=d) for d in range(4)]] ... ) >>> nw.from_native(s_native, series_only=True).dtype Date """ class Time(TemporalType): """Data type representing the time of day. Examples: >>> import polars as pl >>> import pyarrow as pa >>> import narwhals as nw >>> import duckdb >>> from datetime import time >>> data = [time(9, 0), time(9, 1, 10), time(9, 2)] >>> ser_pl = pl.Series(data) >>> ser_pa = pa.chunked_array([pa.array(data, type=pa.time64("ns"))]) >>> rel = duckdb.sql( ... " SELECT * FROM (VALUES (TIME '12:00:00'), (TIME '14:30:15')) df(t)" ... ) >>> nw.from_native(ser_pl, series_only=True).dtype Time >>> nw.from_native(ser_pa, series_only=True).dtype Time >>> nw.from_native(rel).collect_schema()["t"] Time """ class Binary(DType): """Binary type. Examples: >>> import polars as pl >>> import narwhals as nw >>> import pyarrow as pa >>> import duckdb >>> data = [b"test1", b"test2"] >>> ser_pl = pl.Series(data, dtype=pl.Binary) >>> ser_pa = pa.chunked_array([pa.array(data, type=pa.binary())]) >>> rel = duckdb.sql( ... "SELECT * FROM (VALUES (BLOB 'test1'), (BLOB 'test2')) AS df(t)" ... ) >>> nw.from_native(ser_pl, series_only=True).dtype Binary >>> nw.from_native(ser_pa, series_only=True).dtype Binary >>> nw.from_native(rel).collect_schema()["t"] Binary """ narwhals-2.11.0/narwhals/exceptions.py0000644000000000000000000000717013615410400014711 0ustar00from __future__ import annotations from typing import TYPE_CHECKING if TYPE_CHECKING: from collections.abc import Collection, Iterable class NarwhalsError(ValueError): """Base class for all Narwhals exceptions.""" class FormattedKeyError(KeyError): """KeyError with formatted error message. Python's `KeyError` has special casing around formatting (see https://bugs.python.org/issue2651). Use this class when the error message has newlines and other special format characters. Needed by https://github.com/tensorflow/tensorflow/issues/36857. """ def __init__(self, message: str) -> None: self.message = message def __str__(self) -> str: return self.message class ColumnNotFoundError(FormattedKeyError, NarwhalsError): """Exception raised when column name isn't present.""" def __init__(self, message: str) -> None: self.message = message super().__init__(self.message) @classmethod def from_missing_and_available_column_names( cls, missing_columns: Iterable[str], available_columns: Collection[str], / ) -> ColumnNotFoundError: message = ( f"The following columns were not found: {sorted(missing_columns)}" f"\n\nHint: Did you mean one of these columns: {list(available_columns)}?" ) return ColumnNotFoundError(message) @classmethod def from_available_column_names( cls, available_columns: Collection[str] ) -> ColumnNotFoundError: message = ( "The selected columns were not found." f"\n\nHint: Did you mean one of these columns: {list(available_columns)}?" ) return ColumnNotFoundError(message) class ComputeError(NarwhalsError): """Exception raised when the underlying computation could not be evaluated.""" class ShapeError(NarwhalsError): """Exception raised when trying to perform operations on data structures with incompatible shapes.""" class MultiOutputExpressionError(NarwhalsError): """Exception raised when using multi-output expression in unsupported context.""" class DuplicateError(NarwhalsError): """Exception when duplicate column names are encountered.""" class InvalidOperationError(NarwhalsError): """Exception raised during invalid operations.""" class InvalidIntoExprError(TypeError, NarwhalsError): """Exception raised when object can't be converted to expression.""" def __init__(self, message: str) -> None: self.message = message super().__init__(self.message) @classmethod def from_invalid_type(cls: type, invalid_type: type) -> InvalidIntoExprError: message = ( f"Expected an object which can be converted into an expression, got {invalid_type}\n\n" "Hint:\n" "- if you were trying to select a column which does not have a string\n" " column name, then you should explicitly use `nw.col`.\n" " For example, `df.select(nw.col(0))` if you have a column named `0`.\n" "- if you were trying to create a new literal column, then you \n" " should explicitly use `nw.lit`.\n" " For example, `df.select(nw.lit(0))` if you want to create a new\n" " column with literal value `0`." ) return InvalidIntoExprError(message) class UnsupportedDTypeError(NarwhalsError): """Exception raised when trying to convert to a DType which is not supported by the given backend.""" class NarwhalsUnstableWarning(UserWarning): """Warning issued when a method or function is considered unstable in the stable api.""" class PerformanceWarning(Warning): ... narwhals-2.11.0/narwhals/expr.py0000644000000000000000000031556213615410400013515 0ustar00from __future__ import annotations import math from collections.abc import Iterable, Mapping, Sequence from typing import TYPE_CHECKING, Any, Callable from narwhals._expression_parsing import ( ExprKind, ExprNode, evaluate_node, evaluate_root_node, ) from narwhals._utils import _validate_rolling_arguments, ensure_type, flatten, no_default from narwhals.dtypes import _validate_dtype from narwhals.exceptions import ComputeError, InvalidOperationError from narwhals.expr_cat import ExprCatNamespace from narwhals.expr_dt import ExprDateTimeNamespace from narwhals.expr_list import ExprListNamespace from narwhals.expr_name import ExprNameNamespace from narwhals.expr_str import ExprStringNamespace from narwhals.expr_struct import ExprStructNamespace from narwhals.translate import to_native if TYPE_CHECKING: from typing import NoReturn, TypeVar from typing_extensions import Concatenate, ParamSpec, Self from narwhals._compliant import CompliantExpr, CompliantNamespace from narwhals._typing import NoDefault from narwhals.dtypes import DType from narwhals.series import Series from narwhals.typing import ( ClosedInterval, FillNullStrategy, IntoDType, IntoExpr, ModeKeepStrategy, NonNestedLiteral, NumericLiteral, RankMethod, RollingInterpolationMethod, TemporalLiteral, ) PS = ParamSpec("PS") R = TypeVar("R") class Expr: def __init__(self, *nodes: ExprNode) -> None: self._nodes = nodes def _to_compliant_expr( self, ns: CompliantNamespace[Any, Any] ) -> CompliantExpr[Any, Any]: nodes = self._nodes ce = evaluate_root_node(nodes[0], ns) for node in nodes[1:]: ce = evaluate_node(ce, node, ns) return ce def _append_node(self, node: ExprNode) -> Self: return self.__class__(*self._nodes, node) def _with_over_node(self, node: ExprNode) -> Self: # insert `over` before any elementwise operations. # check "how it works" page in docs for why we do this. new_nodes = list(self._nodes) kwargs_no_order_by = { key: value if key != "order_by" else [] for (key, value) in node.kwargs.items() } node_without_order_by = node._with_kwargs(**kwargs_no_order_by) n = len(new_nodes) i = n while i > 0 and (_node := new_nodes[i - 1]).kind is ExprKind.ELEMENTWISE: i -= 1 _node._push_down_over_node_in_place(node, node_without_order_by) if i == n: # node could not be pushed down, just append as-is new_nodes.append(node) return self.__class__(*new_nodes) if node.kwargs["order_by"] and any(node.is_orderable() for node in new_nodes[:i]): new_nodes.insert(i, node) elif node.kwargs["partition_by"] and any( not node.is_elementwise() for node in new_nodes[:i] ): new_nodes.insert(i, node_without_order_by) elif all(node.is_elementwise() for node in new_nodes): msg = "Cannot apply `over` to elementwise expression." raise InvalidOperationError(msg) return self.__class__(*new_nodes) def __repr__(self) -> str: """Pretty-print the expression by combining all nodes in the metadata.""" return ".".join(repr(node) for node in self._nodes) def __bool__(self) -> NoReturn: msg = ( f"the truth value of {type(self)} is ambiguous" "\n\n" "You probably got here by using a Python standard library function instead " "of the native expressions API.\n" "Here are some things you might want to try:\n" "- instead of `nw.col('a') and nw.col('b')`, use `nw.col('a') & nw.col('b')`\n" "- instead of `nw.col('a') in [y, z]`, use `nw.col('a').is_in([y, z])`\n" "- instead of `max(nw.col('a'), nw.col('b'))`, use `nw.max_horizontal(nw.col('a'), nw.col('b'))`\n" ) raise TypeError(msg) def _taxicab_norm(self) -> Self: # This is just used to test out the stable api feature in a realistic-ish way. # It's not intended to be used. return self.abs().sum() # --- convert --- def alias(self, name: str) -> Self: """Rename the expression. Arguments: name: The new name. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2], "b": [4, 5]}) >>> df = nw.from_native(df_native) >>> df.select((nw.col("b") + 10).alias("c")) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | c | | 0 14 | | 1 15 | └──────────────────┘ """ return self._append_node(ExprNode(ExprKind.ELEMENTWISE, "alias", name=name)) def pipe( self, function: Callable[Concatenate[Self, PS], R], *args: PS.args, **kwargs: PS.kwargs, ) -> R: """Pipe function call. Arguments: function: Function to apply. args: Positional arguments to pass to function. kwargs: Keyword arguments to pass to function. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2, 3, 4]}) >>> df = nw.from_native(df_native) >>> df.with_columns(a_piped=nw.col("a").pipe(lambda x: x + 1)) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a a_piped | | 0 1 2 | | 1 2 3 | | 2 3 4 | | 3 4 5 | └──────────────────┘ """ return function(self, *args, **kwargs) def cast(self, dtype: IntoDType) -> Self: """Redefine an object's data type. Arguments: dtype: Data type that the object will be cast into. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"foo": [1, 2, 3], "bar": [6.0, 7.0, 8.0]}) >>> df = nw.from_native(df_native) >>> df.select(nw.col("foo").cast(nw.Float32), nw.col("bar").cast(nw.UInt8)) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | foo bar | | 0 1.0 6 | | 1 2.0 7 | | 2 3.0 8 | └──────────────────┘ """ _validate_dtype(dtype) return self._append_node(ExprNode(ExprKind.ELEMENTWISE, "cast", dtype=dtype)) # --- binary --- def _with_binary(self, attr: str, other: Self | Any) -> Self: node = ExprNode(ExprKind.ELEMENTWISE, attr, other, str_as_lit=True) return self._append_node(node) def __eq__(self, other: Self | Any) -> Self: # type: ignore[override] return self._with_binary("__eq__", other) def __ne__(self, other: Self | Any) -> Self: # type: ignore[override] return self._with_binary("__ne__", other) def __and__(self, other: Any) -> Self: return self._with_binary("__and__", other) def __rand__(self, other: Any) -> Self: return (self & other).alias("literal") # type: ignore[no-any-return] def __or__(self, other: Any) -> Self: return self._with_binary("__or__", other) def __ror__(self, other: Any) -> Self: return (self | other).alias("literal") # type: ignore[no-any-return] def __add__(self, other: Any) -> Self: return self._with_binary("__add__", other) def __radd__(self, other: Any) -> Self: return (self + other).alias("literal") # type: ignore[no-any-return] def __sub__(self, other: Any) -> Self: return self._with_binary("__sub__", other) def __rsub__(self, other: Any) -> Self: return self._with_binary("__rsub__", other) def __truediv__(self, other: Any) -> Self: return self._with_binary("__truediv__", other) def __rtruediv__(self, other: Any) -> Self: return self._with_binary("__rtruediv__", other) def __mul__(self, other: Any) -> Self: return self._with_binary("__mul__", other) def __rmul__(self, other: Any) -> Self: return (self * other).alias("literal") # type: ignore[no-any-return] def __le__(self, other: Any) -> Self: return self._with_binary("__le__", other) def __lt__(self, other: Any) -> Self: return self._with_binary("__lt__", other) def __gt__(self, other: Any) -> Self: return self._with_binary("__gt__", other) def __ge__(self, other: Any) -> Self: return self._with_binary("__ge__", other) def __pow__(self, other: Any) -> Self: return self._with_binary("__pow__", other) def __rpow__(self, other: Any) -> Self: return self._with_binary("__rpow__", other) def __floordiv__(self, other: Any) -> Self: return self._with_binary("__floordiv__", other) def __rfloordiv__(self, other: Any) -> Self: return self._with_binary("__rfloordiv__", other) def __mod__(self, other: Any) -> Self: return self._with_binary("__mod__", other) def __rmod__(self, other: Any) -> Self: return self._with_binary("__rmod__", other) # --- unary --- def __invert__(self) -> Self: return self._append_node(ExprNode(ExprKind.ELEMENTWISE, "__invert__")) def any(self) -> Self: """Return whether any of the values in the column are `True`. If there are no non-null elements, the result is `False`. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [True, False], "b": [True, True]}) >>> df = nw.from_native(df_native) >>> df.select(nw.col("a", "b").any()) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 True True | └──────────────────┘ """ return self._append_node(ExprNode(ExprKind.AGGREGATION, "any")) def all(self) -> Self: """Return whether all values in the column are `True`. If there are no non-null elements, the result is `True`. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [True, False], "b": [True, True]}) >>> df = nw.from_native(df_native) >>> df.select(nw.col("a", "b").all()) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 False True | └──────────────────┘ """ return self._append_node(ExprNode(ExprKind.AGGREGATION, "all")) def ewm_mean( self, *, com: float | None = None, span: float | None = None, half_life: float | None = None, alpha: float | None = None, adjust: bool = True, min_samples: int = 1, ignore_nulls: bool = False, ) -> Self: r"""Compute exponentially-weighted moving average. Arguments: com: Specify decay in terms of center of mass, $\gamma$, with
$\alpha = \frac{1}{1+\gamma}\forall\gamma\geq0$ span: Specify decay in terms of span, $\theta$, with
$\alpha = \frac{2}{\theta + 1} \forall \theta \geq 1$ half_life: Specify decay in terms of half-life, $\tau$, with
$\alpha = 1 - \exp \left\{ \frac{ -\ln(2) }{ \tau } \right\} \forall \tau > 0$ alpha: Specify smoothing factor alpha directly, $0 < \alpha \leq 1$. adjust: Divide by decaying adjustment factor in beginning periods to account for imbalance in relative weightings - When `adjust=True` (the default) the EW function is calculated using weights $w_i = (1 - \alpha)^i$ - When `adjust=False` the EW function is calculated recursively by $$ y_0=x_0 $$ $$ y_t = (1 - \alpha)y_{t - 1} + \alpha x_t $$ min_samples: Minimum number of observations in window required to have a value, (otherwise result is null). ignore_nulls: Ignore missing values when calculating weights. - When `ignore_nulls=False` (default), weights are based on absolute positions. For example, the weights of $x_0$ and $x_2$ used in calculating the final weighted average of $[x_0, None, x_2]$ are $(1-\alpha)^2$ and $1$ if `adjust=True`, and $(1-\alpha)^2$ and $\alpha$ if `adjust=False`. - When `ignore_nulls=True`, weights are based on relative positions. For example, the weights of $x_0$ and $x_2$ used in calculating the final weighted average of $[x_0, None, x_2]$ are $1-\alpha$ and $1$ if `adjust=True`, and $1-\alpha$ and $\alpha$ if `adjust=False`. Examples: >>> import pandas as pd >>> import polars as pl >>> import narwhals as nw >>> from narwhals.typing import IntoFrameT >>> >>> data = {"a": [1, 2, 3]} >>> df_pd = pd.DataFrame(data) >>> df_pl = pl.DataFrame(data) We define a library agnostic function: >>> def agnostic_ewm_mean(df_native: IntoFrameT) -> IntoFrameT: ... df = nw.from_native(df_native) ... return df.select( ... nw.col("a").ewm_mean(com=1, ignore_nulls=False) ... ).to_native() We can then pass either pandas or Polars to `agnostic_ewm_mean`: >>> agnostic_ewm_mean(df_pd) a 0 1.000000 1 1.666667 2 2.428571 >>> agnostic_ewm_mean(df_pl) # doctest: +NORMALIZE_WHITESPACE shape: (3, 1) ┌──────────┐ │ a │ │ --- │ │ f64 │ ╞══════════╡ │ 1.0 │ │ 1.666667 │ │ 2.428571 │ └──────────┘ """ return self._append_node( ExprNode( ExprKind.ORDERABLE_WINDOW, "ewm_mean", com=com, span=span, half_life=half_life, alpha=alpha, adjust=adjust, min_samples=min_samples, ignore_nulls=ignore_nulls, ) ) def mean(self) -> Self: """Get mean value. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [-1, 0, 1], "b": [2, 4, 6]}) >>> df = nw.from_native(df_native) >>> df.select(nw.col("a", "b").mean()) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 0.0 4.0 | └──────────────────┘ """ return self._append_node(ExprNode(ExprKind.AGGREGATION, "mean")) def median(self) -> Self: """Get median value. Notes: Results might slightly differ across backends due to differences in the underlying algorithms used to compute the median. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 8, 3], "b": [4, 5, 2]}) >>> df = nw.from_native(df_native) >>> df.select(nw.col("a", "b").median()) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 3.0 4.0 | └──────────────────┘ """ return self._append_node(ExprNode(ExprKind.AGGREGATION, "median")) def std(self, *, ddof: int = 1) -> Self: """Get standard deviation. Arguments: ddof: "Delta Degrees of Freedom": the divisor used in the calculation is N - ddof, where N represents the number of elements. By default ddof is 1. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [20, 25, 60], "b": [1.5, 1, -1.4]}) >>> df = nw.from_native(df_native) >>> df.select(nw.col("a", "b").std(ddof=0)) ┌─────────────────────┐ | Narwhals DataFrame | |---------------------| | a b| |0 17.79513 1.265789| └─────────────────────┘ """ return self._append_node(ExprNode(ExprKind.AGGREGATION, "std", ddof=ddof)) def var(self, *, ddof: int = 1) -> Self: """Get variance. Arguments: ddof: "Delta Degrees of Freedom": the divisor used in the calculation is N - ddof, where N represents the number of elements. By default ddof is 1. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [20, 25, 60], "b": [1.5, 1, -1.4]}) >>> df = nw.from_native(df_native) >>> df.select(nw.col("a", "b").var(ddof=0)) ┌───────────────────────┐ | Narwhals DataFrame | |-----------------------| | a b| |0 316.666667 1.602222| └───────────────────────┘ """ return self._append_node(ExprNode(ExprKind.AGGREGATION, "var", ddof=ddof)) def map_batches( self, function: Callable[[Any], CompliantExpr[Any, Any]], return_dtype: DType | None = None, *, returns_scalar: bool = False, ) -> Self: """Apply a custom python function to a whole Series or sequence of Series. The output of this custom function is presumed to be either a Series, or a NumPy array (in which case it will be automatically converted into a Series). Arguments: function: Function to apply to Series. return_dtype: Dtype of the output Series. If not set, the dtype will be inferred based on the first non-null value that is returned by the function. returns_scalar: If the function returns a scalar, by default it will be wrapped in a list in the output, since the assumption is that the function always returns something Series-like. If you want to keep the result as a scalar, set this argument to True. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) >>> df = nw.from_native(df_native) >>> df.with_columns( ... nw.col("a", "b") ... .map_batches(lambda s: s.to_numpy() + 1, return_dtype=nw.Float64) ... .name.suffix("_mapped") ... ) ┌───────────────────────────┐ | Narwhals DataFrame | |---------------------------| | a b a_mapped b_mapped| |0 1 4 2.0 5.0| |1 2 5 3.0 6.0| |2 3 6 4.0 7.0| └───────────────────────────┘ """ kind = ( ExprKind.ORDERABLE_AGGREGATION if returns_scalar else ExprKind.ORDERABLE_FILTRATION ) return self._append_node( ExprNode( kind, "map_batches", function=function, return_dtype=return_dtype, returns_scalar=returns_scalar, ) ) def skew(self) -> Self: """Calculate the sample skewness of a column. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2, 3, 4, 5], "b": [1, 1, 2, 10, 100]}) >>> df = nw.from_native(df_native) >>> df.select(nw.col("a", "b").skew()) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 0.0 1.472427 | └──────────────────┘ """ return self._append_node(ExprNode(ExprKind.AGGREGATION, "skew")) def kurtosis(self) -> Self: """Compute the kurtosis (Fisher's definition) without bias correction. Kurtosis is the fourth central moment divided by the square of the variance. The Fisher's definition is used where 3.0 is subtracted from the result to give 0.0 for a normal distribution. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2, 3, 4, 5], "b": [1, 1, 2, 10, 100]}) >>> df = nw.from_native(df_native) >>> df.select(nw.col("a", "b").kurtosis()) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 -1.3 0.210657 | └──────────────────┘ """ return self._append_node(ExprNode(ExprKind.AGGREGATION, "kurtosis")) def sum(self) -> Self: """Return the sum value. If there are no non-null elements, the result is zero. Examples: >>> import duckdb >>> import narwhals as nw >>> df_native = duckdb.sql("SELECT * FROM VALUES (5, 50), (10, 100) df(a, b)") >>> df = nw.from_native(df_native) >>> df.select(nw.col("a", "b").sum()) ┌───────────────────┐ |Narwhals LazyFrame | |-------------------| |┌────────┮────────┐| |│ a │ b │| |│ int128 │ int128 │| |├────────┾────────â”Ī| |│ 15 │ 150 │| |└────────â”ī────────┘| └───────────────────┘ """ return self._append_node(ExprNode(ExprKind.AGGREGATION, "sum")) def min(self) -> Self: """Returns the minimum value(s) from a column(s). Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2], "b": [4, 3]}) >>> df = nw.from_native(df_native) >>> df.select(nw.min("a", "b")) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 1 3 | └──────────────────┘ """ return self._append_node(ExprNode(ExprKind.AGGREGATION, "min")) def max(self) -> Self: """Returns the maximum value(s) from a column(s). Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [10, 20], "b": [50, 100]}) >>> df = nw.from_native(df_native) >>> df.select(nw.max("a", "b")) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 20 100 | └──────────────────┘ """ return self._append_node(ExprNode(ExprKind.AGGREGATION, "max")) def count(self) -> Self: """Returns the number of non-null elements in the column. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2, 3], "b": [None, 4, 4]}) >>> df = nw.from_native(df_native) >>> df.select(nw.all().count()) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 3 2 | └──────────────────┘ """ return self._append_node(ExprNode(ExprKind.AGGREGATION, "count")) def n_unique(self) -> Self: """Returns count of unique values. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2, 3, 4, 5], "b": [1, 1, 3, 3, 5]}) >>> df = nw.from_native(df_native) >>> df.select(nw.col("a", "b").n_unique()) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 5 3 | └──────────────────┘ """ return self._append_node(ExprNode(ExprKind.AGGREGATION, "n_unique")) def unique(self) -> Self: """Return unique values of this expression. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 1, 3, 5, 5], "b": [2, 4, 4, 6, 6]}) >>> df = nw.from_native(df_native) >>> df.select(nw.col("a", "b").unique().sum()) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 9 12 | └──────────────────┘ """ return self._append_node(ExprNode(ExprKind.FILTRATION, "unique")) def abs(self) -> Self: """Return absolute value of each element. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, -2], "b": [-3, 4]}) >>> df = nw.from_native(df_native) >>> df.with_columns(nw.col("a", "b").abs().name.suffix("_abs")) ┌─────────────────────┐ | Narwhals DataFrame | |---------------------| | a b a_abs b_abs| |0 1 -3 1 3| |1 -2 4 2 4| └─────────────────────┘ """ return self._append_node(ExprNode(ExprKind.ELEMENTWISE, "abs")) def cum_sum(self, *, reverse: bool = False) -> Self: """Return cumulative sum. Info: For lazy backends, this operation must be followed by `Expr.over` with `order_by` specified, see [order-dependence](../concepts/order_dependence.md). Arguments: reverse: reverse the operation Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 1, 3, 5, 5], "b": [2, 4, 4, 6, 6]}) >>> df = nw.from_native(df_native) >>> df.with_columns(a_cum_sum=nw.col("a").cum_sum()) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b a_cum_sum| |0 1 2 1| |1 1 4 2| |2 3 4 5| |3 5 6 10| |4 5 6 15| └──────────────────┘ """ return self._append_node( ExprNode(ExprKind.ORDERABLE_WINDOW, "cum_sum", reverse=reverse) ) def diff(self) -> Self: """Returns the difference between each element and the previous one. Info: For lazy backends, this operation must be followed by `Expr.over` with `order_by` specified, see [order-dependence](../concepts/order_dependence.md). Notes: pandas may change the dtype here, for example when introducing missing values in an integer column. To ensure, that the dtype doesn't change, you may want to use `fill_null` and `cast`. For example, to calculate the diff and fill missing values with `0` in a Int64 column, you could do: nw.col("a").diff().fill_null(0).cast(nw.Int64) Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"a": [1, 1, 3, 5, 5]}) >>> df = nw.from_native(df_native) >>> df.with_columns(a_diff=nw.col("a").diff()) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | shape: (5, 2) | | ┌─────┮────────┐ | | │ a ┆ a_diff │ | | │ --- ┆ --- │ | | │ i64 ┆ i64 │ | | ╞═════╩════════╡ | | │ 1 ┆ null │ | | │ 1 ┆ 0 │ | | │ 3 ┆ 2 │ | | │ 5 ┆ 2 │ | | │ 5 ┆ 0 │ | | └─────â”ī────────┘ | └──────────────────┘ """ return self._append_node(ExprNode(ExprKind.ORDERABLE_WINDOW, "diff")) def shift(self, n: int) -> Self: """Shift values by `n` positions. Info: For lazy backends, this operation must be followed by `Expr.over` with `order_by` specified, see [order-dependence](../concepts/order_dependence.md). Arguments: n: Number of positions to shift values by. Notes: pandas may change the dtype here, for example when introducing missing values in an integer column. To ensure, that the dtype doesn't change, you may want to use `fill_null` and `cast`. For example, to shift and fill missing values with `0` in a Int64 column, you could do: nw.col("a").shift(1).fill_null(0).cast(nw.Int64) Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"a": [1, 1, 3, 5, 5]}) >>> df = nw.from_native(df_native) >>> df.with_columns(a_shift=nw.col("a").shift(n=1)) ┌──────────────────┐ |Narwhals DataFrame| |------------------| |shape: (5, 2) | |┌─────┮─────────┐ | |│ a ┆ a_shift │ | |│ --- ┆ --- │ | |│ i64 ┆ i64 │ | |╞═════╩═════════╡ | |│ 1 ┆ null │ | |│ 1 ┆ 1 │ | |│ 3 ┆ 1 │ | |│ 5 ┆ 3 │ | |│ 5 ┆ 5 │ | |└─────â”ī─────────┘ | └──────────────────┘ """ ensure_type(n, int, param_name="n") return self._append_node(ExprNode(ExprKind.ORDERABLE_WINDOW, "shift", n=n)) def replace_strict( self, old: Sequence[Any] | Mapping[Any, Any], new: Sequence[Any] | None = None, *, default: Any | NoDefault = no_default, return_dtype: IntoDType | None = None, ) -> Self: """Replace all values by different values. Arguments: old: Sequence of values to replace. It also accepts a mapping of values to their replacement as syntactic sugar for `replace_strict(old=list(mapping.keys()), new=list(mapping.values()))`. new: Sequence of values to replace by. Length must match the length of `old`. default: Set values that were not replaced to this value. If no default is specified, (default), an error is raised if any values were not replaced. Accepts expression input. Non-expression inputs are parsed as literals. return_dtype: The data type of the resulting expression. If set to `None` (default), the data type is determined automatically based on the other inputs. Raises: InvalidOperationError: If any non-null values in the original column were not replaced, and no default was specified. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [3, 0, 1, 2]}) >>> df = nw.from_native(df_native) >>> df.with_columns( ... b=nw.col("a").replace_strict( ... [0, 1, 2, 3], ... ["zero", "one", "two", "three"], ... return_dtype=nw.String, ... ) ... ) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 3 three | | 1 0 zero | | 2 1 one | | 3 2 two | └──────────────────┘ Replace values and set a default for values not in the mapping: >>> data = {"a": [1, 2, 3, 4], "b": ["beluga", "narwhal", "orca", "vaquita"]} >>> df = nw.from_native(pd.DataFrame(data)) >>> df.with_columns( ... a_replaced=nw.col("a").replace_strict( ... {1: "one", 2: "two"}, ... default=nw.concat_str(nw.lit("default_"), nw.col("b")), ... return_dtype=nw.String, ... ) ... ) ┌──────────────────────────────┐ | Narwhals DataFrame | |------------------------------| | a b a_replaced| |0 1 beluga one| |1 2 narwhal two| |2 3 orca default_orca| |3 4 vaquita default_vaquita| └──────────────────────────────┘ """ if new is None: if not isinstance(old, Mapping): msg = "`new` argument is required if `old` argument is not a Mapping type" raise TypeError(msg) new = list(old.values()) old = list(old.keys()) if default is no_default: node = ExprNode( ExprKind.ELEMENTWISE, "replace_strict", default=default, old=old, new=new, return_dtype=return_dtype, ) else: node = ExprNode( ExprKind.ELEMENTWISE, "replace_strict", default, old=old, new=new, return_dtype=return_dtype, str_as_lit=True, ) return self._append_node(node) # --- transform --- def is_between( self, lower_bound: Any | IntoExpr, upper_bound: Any | IntoExpr, closed: ClosedInterval = "both", ) -> Self: """Check if this expression is between the given lower and upper bounds. Arguments: lower_bound: Lower bound value. String literals are interpreted as column names. upper_bound: Upper bound value. String literals are interpreted as column names. closed: Define which sides of the interval are closed (inclusive). Options are {"left", "right", "none", "both"}. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2, 3, 4, 5]}) >>> df = nw.from_native(df_native) >>> df.with_columns(b=nw.col("a").is_between(2, 4, "right")) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 1 False | | 1 2 False | | 2 3 True | | 3 4 True | | 4 5 False | └──────────────────┘ """ node = ExprNode( ExprKind.ELEMENTWISE, "is_between", lower_bound, upper_bound, closed=closed ) return self._append_node(node) def is_in(self, other: Any) -> Self: """Check if elements of this expression are present in the other iterable. Arguments: other: iterable Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2, 9, 10]}) >>> df = nw.from_native(df_native) >>> df.with_columns(b=nw.col("a").is_in([1, 2])) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 1 True | | 1 2 True | | 2 9 False | | 3 10 False | └──────────────────┘ """ if isinstance(other, Iterable) and not isinstance(other, (str, bytes)): return self._append_node( ExprNode( ExprKind.ELEMENTWISE, "is_in", other=to_native(other, pass_through=True), ) ) msg = "Narwhals `is_in` doesn't accept expressions as an argument, as opposed to Polars. You should provide an iterable instead." raise NotImplementedError(msg) def filter(self, *predicates: Any) -> Self: """Filters elements based on a condition, returning a new expression. Arguments: predicates: Conditions to filter by (which get AND-ed together). Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame( ... {"a": [2, 3, 4, 5, 6, 7], "b": [10, 11, 12, 13, 14, 15]} ... ) >>> df = nw.from_native(df_native) >>> df.select( ... nw.col("a").filter(nw.col("a") > 4), ... nw.col("b").filter(nw.col("b") < 13), ... ) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 3 5 10 | | 4 6 11 | | 5 7 12 | └──────────────────┘ """ return self._append_node(ExprNode(ExprKind.FILTRATION, "filter", *predicates)) def is_null(self) -> Self: """Returns a boolean Series indicating which values are null. Notes: pandas handles null values differently from Polars and PyArrow. See [null_handling](../concepts/null_handling.md/) for reference. Examples: >>> import duckdb >>> import narwhals as nw >>> df_native = duckdb.sql( ... "SELECT * FROM VALUES (null, CAST('NaN' AS DOUBLE)), (2, 2.) df(a, b)" ... ) >>> df = nw.from_native(df_native) >>> df.with_columns( ... a_is_null=nw.col("a").is_null(), b_is_null=nw.col("b").is_null() ... ) ┌──────────────────────────────────────────┐ | Narwhals LazyFrame | |------------------------------------------| |┌───────┮────────┮───────────┮───────────┐| |│ a │ b │ a_is_null │ b_is_null │| |│ int32 │ double │ boolean │ boolean │| |├───────┾────────┾───────────┾───────────â”Ī| |│ NULL │ nan │ true │ false │| |│ 2 │ 2.0 │ false │ false │| |└───────â”ī────────â”ī───────────â”ī───────────┘| └──────────────────────────────────────────┘ """ return self._append_node(ExprNode(ExprKind.ELEMENTWISE, "is_null")) def is_nan(self) -> Self: """Indicate which values are NaN. Notes: pandas handles null values differently from Polars and PyArrow. See [null_handling](../concepts/null_handling.md/) for reference. Examples: >>> import duckdb >>> import narwhals as nw >>> df_native = duckdb.sql( ... "SELECT * FROM VALUES (null, CAST('NaN' AS DOUBLE)), (2, 2.) df(a, b)" ... ) >>> df = nw.from_native(df_native) >>> df.with_columns( ... a_is_nan=nw.col("a").is_nan(), b_is_nan=nw.col("b").is_nan() ... ) ┌────────────────────────────────────────┐ | Narwhals LazyFrame | |----------------------------------------| |┌───────┮────────┮──────────┮──────────┐| |│ a │ b │ a_is_nan │ b_is_nan │| |│ int32 │ double │ boolean │ boolean │| |├───────┾────────┾──────────┾──────────â”Ī| |│ NULL │ nan │ NULL │ true │| |│ 2 │ 2.0 │ false │ false │| |└───────â”ī────────â”ī──────────â”ī──────────┘| └────────────────────────────────────────┘ """ return self._append_node(ExprNode(ExprKind.ELEMENTWISE, "is_nan")) def fill_null( self, value: Expr | NonNestedLiteral = None, strategy: FillNullStrategy | None = None, limit: int | None = None, ) -> Self: """Fill null values with given value. Arguments: value: Value or expression used to fill null values. strategy: Strategy used to fill null values. limit: Number of consecutive null values to fill when using the 'forward' or 'backward' strategy. Notes: - pandas handles null values differently from other libraries. See [null_handling](../concepts/null_handling.md/) for reference. - For pandas Series of `object` dtype, `fill_null` will not automatically change the Series' dtype as pandas used to do. Explicitly call `cast` if you want the dtype to change. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame( ... { ... "a": [2, None, None, 3], ... "b": [2.0, float("nan"), float("nan"), 3.0], ... "c": [1, 2, 3, 4], ... } ... ) >>> df = nw.from_native(df_native) >>> df.with_columns( ... nw.col("a", "b").fill_null(0).name.suffix("_filled"), ... nw.col("a").fill_null(nw.col("c")).name.suffix("_filled_with_c"), ... ) ┌────────────────────────────────────────────────────────────┐ | Narwhals DataFrame | |------------------------------------------------------------| |shape: (4, 6) | |┌──────┮─────┮─────┮──────────┮──────────┮─────────────────┐| |│ a ┆ b ┆ c ┆ a_filled ┆ b_filled ┆ a_filled_with_c │| |│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │| |│ i64 ┆ f64 ┆ i64 ┆ i64 ┆ f64 ┆ i64 │| |╞══════╩═════╩═════╩══════════╩══════════╩═════════════════╡| |│ 2 ┆ 2.0 ┆ 1 ┆ 2 ┆ 2.0 ┆ 2 │| |│ null ┆ NaN ┆ 2 ┆ 0 ┆ NaN ┆ 2 │| |│ null ┆ NaN ┆ 3 ┆ 0 ┆ NaN ┆ 3 │| |│ 3 ┆ 3.0 ┆ 4 ┆ 3 ┆ 3.0 ┆ 3 │| |└──────â”ī─────â”ī─────â”ī──────────â”ī──────────â”ī─────────────────┘| └────────────────────────────────────────────────────────────┘ Using a strategy: >>> df.select( ... nw.col("a", "b"), ... nw.col("a", "b") ... .fill_null(strategy="forward", limit=1) ... .name.suffix("_nulls_forward_filled"), ... ) ┌────────────────────────────────────────────────────────────────┐ | Narwhals DataFrame | |----------------------------------------------------------------| |shape: (4, 4) | |┌──────┮─────┮────────────────────────┮────────────────────────┐| |│ a ┆ b ┆ a_nulls_forward_filled ┆ b_nulls_forward_filled │| |│ --- ┆ --- ┆ --- ┆ --- │| |│ i64 ┆ f64 ┆ i64 ┆ f64 │| |╞══════╩═════╩════════════════════════╩════════════════════════╡| |│ 2 ┆ 2.0 ┆ 2 ┆ 2.0 │| |│ null ┆ NaN ┆ 2 ┆ NaN │| |│ null ┆ NaN ┆ null ┆ NaN │| |│ 3 ┆ 3.0 ┆ 3 ┆ 3.0 │| |└──────â”ī─────â”ī────────────────────────â”ī────────────────────────┘| └────────────────────────────────────────────────────────────────┘ """ if value is not None and strategy is not None: msg = "cannot specify both `value` and `strategy`" raise ValueError(msg) if value is None and strategy is None: msg = "must specify either a fill `value` or `strategy`" raise ValueError(msg) if strategy is not None and strategy not in {"forward", "backward"}: msg = f"strategy not supported: {strategy}" raise ValueError(msg) if strategy is not None: node = ExprNode( ExprKind.ORDERABLE_WINDOW, "fill_null", value=value, strategy=strategy, limit=limit, ) else: node = ExprNode( ExprKind.ELEMENTWISE, "fill_null", value, strategy=strategy, limit=limit, str_as_lit=True, ) return self._append_node(node) def fill_nan(self, value: float | None) -> Self: """Fill floating point NaN values with given value. Arguments: value: Value used to fill NaN values. Notes: This function only fills `'NaN'` values, not null ones, except for pandas which doesn't distinguish between them. See [null_handling](../concepts/null_handling.md/) for reference. Examples: >>> import duckdb >>> import narwhals as nw >>> df_native = duckdb.sql( ... "SELECT * FROM VALUES (5.::DOUBLE, 50.::DOUBLE), ('NaN', null) df(a, b)" ... ) >>> df = nw.from_native(df_native) >>> df.with_columns(nw.col("a", "b").fill_nan(0).name.suffix("_nans_filled")) ┌───────────────────────────────────────────────────┐ | Narwhals LazyFrame | |---------------------------------------------------| |┌────────┮────────┮───────────────┮───────────────┐| |│ a │ b │ a_nans_filled │ b_nans_filled │| |│ double │ double │ double │ double │| |├────────┾────────┾───────────────┾───────────────â”Ī| |│ 5.0 │ 50.0 │ 5.0 │ 50.0 │| |│ nan │ NULL │ 0.0 │ NULL │| |└────────â”ī────────â”ī───────────────â”ī───────────────┘| └───────────────────────────────────────────────────┘ """ return self._append_node(ExprNode(ExprKind.ELEMENTWISE, "fill_nan", value=value)) # --- partial reduction --- def drop_nulls(self) -> Self: """Drop null values. Notes: pandas handles null values differently from Polars and PyArrow. See [null_handling](../concepts/null_handling.md/) for reference. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"a": [2.0, 4.0, float("nan"), 3.0, None, 5.0]}) >>> df = nw.from_native(df_native) >>> df.select(nw.col("a").drop_nulls()) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | shape: (5, 1) | | ┌─────┐ | | │ a │ | | │ --- │ | | │ f64 │ | | ╞═════╡ | | │ 2.0 │ | | │ 4.0 │ | | │ NaN │ | | │ 3.0 │ | | │ 5.0 │ | | └─────┘ | └──────────────────┘ """ return self._append_node(ExprNode(ExprKind.FILTRATION, "drop_nulls")) def over( self, *partition_by: str | Sequence[str], order_by: str | Sequence[str] | None = None, ) -> Self: """Compute expressions over the given groups (optionally with given order). Arguments: partition_by: Names of columns to compute window expression over. Must be names of columns, as opposed to expressions - so, this is a bit less flexible than Polars' `Expr.over`. order_by: Column(s) to order window functions by. For lazy backends, this argument is required when `over` is applied to order-dependent functions, see [order-dependence](../concepts/order_dependence.md). Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2, 4], "b": ["x", "x", "y"]}) >>> df = nw.from_native(df_native) >>> df.with_columns(a_min_per_group=nw.col("a").min().over("b")) ┌────────────────────────┐ | Narwhals DataFrame | |------------------------| | a b a_min_per_group| |0 1 x 1| |1 2 x 1| |2 4 y 4| └────────────────────────┘ Cumulative operations are also supported, but (currently) only for pandas and Polars: >>> df.with_columns(a_cum_sum_per_group=nw.col("a").cum_sum().over("b")) ┌────────────────────────────┐ | Narwhals DataFrame | |----------------------------| | a b a_cum_sum_per_group| |0 1 x 1| |1 2 x 3| |2 4 y 4| └────────────────────────────┘ """ flat_partition_by = flatten(partition_by) flat_order_by = [order_by] if isinstance(order_by, str) else (order_by or []) if not flat_partition_by and not flat_order_by: # pragma: no cover msg = "At least one of `partition_by` or `order_by` must be specified." raise ValueError(msg) node = ExprNode( ExprKind.OVER, "over", partition_by=flat_partition_by, order_by=flat_order_by ) return self._with_over_node(node) def is_duplicated(self) -> Self: r"""Return a boolean mask indicating duplicated values. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2, 3, 1], "b": ["a", "a", "b", "c"]}) >>> df = nw.from_native(df_native) >>> df.with_columns(nw.all().is_duplicated().name.suffix("_is_duplicated")) ┌─────────────────────────────────────────┐ | Narwhals DataFrame | |-----------------------------------------| | a b a_is_duplicated b_is_duplicated| |0 1 a True True| |1 2 a False True| |2 3 b False False| |3 1 c True False| └─────────────────────────────────────────┘ """ return self._append_node(ExprNode(ExprKind.WINDOW, "is_duplicated")) def is_unique(self) -> Self: r"""Return a boolean mask indicating unique values. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2, 3, 1], "b": ["a", "a", "b", "c"]}) >>> df = nw.from_native(df_native) >>> df.with_columns(nw.all().is_unique().name.suffix("_is_unique")) ┌─────────────────────────────────┐ | Narwhals DataFrame | |---------------------------------| | a b a_is_unique b_is_unique| |0 1 a False False| |1 2 a True False| |2 3 b True True| |3 1 c False True| └─────────────────────────────────┘ """ return self._append_node(ExprNode(ExprKind.WINDOW, "is_unique")) def null_count(self) -> Self: r"""Count null values. Notes: pandas handles null values differently from Polars and PyArrow. See [null_handling](../concepts/null_handling.md/) for reference. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame( ... {"a": [1, 2, None, 1], "b": ["a", None, "b", None]} ... ) >>> df = nw.from_native(df_native) >>> df.select(nw.all().null_count()) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 1 2 | └──────────────────┘ """ return self._append_node(ExprNode(ExprKind.AGGREGATION, "null_count")) def is_first_distinct(self) -> Self: r"""Return a boolean mask indicating the first occurrence of each distinct value. Info: For lazy backends, this operation must be followed by `Expr.over` with `order_by` specified, see [order-dependence](../concepts/order_dependence.md). Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2, 3, 1], "b": ["a", "a", "b", "c"]}) >>> df = nw.from_native(df_native) >>> df.with_columns( ... nw.all().is_first_distinct().name.suffix("_is_first_distinct") ... ) ┌─────────────────────────────────────────────────┐ | Narwhals DataFrame | |-------------------------------------------------| | a b a_is_first_distinct b_is_first_distinct| |0 1 a True True| |1 2 a True False| |2 3 b True True| |3 1 c False True| └─────────────────────────────────────────────────┘ """ return self._append_node(ExprNode(ExprKind.ORDERABLE_WINDOW, "is_first_distinct")) def is_last_distinct(self) -> Self: r"""Return a boolean mask indicating the last occurrence of each distinct value. Info: For lazy backends, this operation must be followed by `Expr.over` with `order_by` specified, see [order-dependence](../concepts/order_dependence.md). Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2, 3, 1], "b": ["a", "a", "b", "c"]}) >>> df = nw.from_native(df_native) >>> df.with_columns( ... nw.all().is_last_distinct().name.suffix("_is_last_distinct") ... ) ┌───────────────────────────────────────────────┐ | Narwhals DataFrame | |-----------------------------------------------| | a b a_is_last_distinct b_is_last_distinct| |0 1 a False False| |1 2 a True True| |2 3 b True True| |3 1 c True True| └───────────────────────────────────────────────┘ """ return self._append_node(ExprNode(ExprKind.ORDERABLE_WINDOW, "is_last_distinct")) def quantile( self, quantile: float, interpolation: RollingInterpolationMethod ) -> Self: r"""Get quantile value. Arguments: quantile: Quantile between 0.0 and 1.0. interpolation: Interpolation method. Note: - pandas and Polars may have implementation differences for a given interpolation method. - [dask](https://docs.dask.org/en/stable/generated/dask.dataframe.Series.quantile.html) has its own method to approximate quantile and it doesn't implement 'nearest', 'higher', 'lower', 'midpoint' as interpolation method - use 'linear' which is closest to the native 'dask' - method. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame( ... {"a": list(range(50)), "b": list(range(50, 100))} ... ) >>> df = nw.from_native(df_native) >>> df.select(nw.col("a", "b").quantile(0.5, interpolation="linear")) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 24.5 74.5 | └──────────────────┘ """ return self._append_node( ExprNode( ExprKind.AGGREGATION, "quantile", quantile=quantile, interpolation=interpolation, ) ) def round(self, decimals: int = 0) -> Self: r"""Round underlying floating point data by `decimals` digits. Arguments: decimals: Number of decimals to round by. Notes: For values exactly halfway between rounded decimal values pandas behaves differently than Polars and Arrow. pandas rounds to the nearest even value (e.g. -0.5 and 0.5 round to 0.0, 1.5 and 2.5 round to 2.0, 3.5 and 4.5 to 4.0, etc..). Polars and Arrow round away from 0 (e.g. -0.5 to -1.0, 0.5 to 1.0, 1.5 to 2.0, 2.5 to 3.0, etc..). Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1.12345, 2.56789, 3.901234]}) >>> df = nw.from_native(df_native) >>> df.with_columns(a_rounded=nw.col("a").round(1)) ┌──────────────────────┐ | Narwhals DataFrame | |----------------------| | a a_rounded| |0 1.123450 1.1| |1 2.567890 2.6| |2 3.901234 3.9| └──────────────────────┘ """ return self._append_node( ExprNode(ExprKind.ELEMENTWISE, "round", decimals=decimals) ) def floor(self) -> Self: r"""Compute the numerical floor. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"values": [1.1, 4.3, -1.3]}) >>> df = nw.from_native(df_native) >>> result = df.with_columns(floor=nw.col("values").floor()) >>> result ┌────────────────────────┐ | Narwhals DataFrame | |------------------------| |pyarrow.Table | |values: double | |floor: double | |---- | |values: [[1.1,4.3,-1.3]]| |floor: [[1,4,-2]] | └────────────────────────┘ """ return self._append_node(ExprNode(ExprKind.ELEMENTWISE, "floor")) def ceil(self) -> Self: r"""Compute the numerical ceiling. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"values": [1.1, 4.3, -1.3]}) >>> df = nw.from_native(df_native) >>> result = df.with_columns(ceil=nw.col("values").ceil()) >>> result ┌────────────────────────┐ | Narwhals DataFrame | |------------------------| |pyarrow.Table | |values: double | |ceil: double | |---- | |values: [[1.1,4.3,-1.3]]| |ceil: [[2,5,-1]] | └────────────────────────┘ """ return self._append_node(ExprNode(ExprKind.ELEMENTWISE, "ceil")) def len(self) -> Self: r"""Return the number of elements in the column. Null values count towards the total. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": ["x", "y", "z"], "b": [1, 2, 1]}) >>> df = nw.from_native(df_native) >>> df.select( ... nw.col("a").filter(nw.col("b") == 1).len().alias("a1"), ... nw.col("a").filter(nw.col("b") == 2).len().alias("a2"), ... ) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a1 a2 | | 0 2 1 | └──────────────────┘ """ return self._append_node(ExprNode(ExprKind.AGGREGATION, "len")) def clip( self, lower_bound: IntoExpr | NumericLiteral | TemporalLiteral | None = None, upper_bound: IntoExpr | NumericLiteral | TemporalLiteral | None = None, ) -> Self: r"""Clip values in the Series. Arguments: lower_bound: Lower bound value. String literals are treated as column names. upper_bound: Upper bound value. String literals are treated as column names. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 2, 3]}) >>> df = nw.from_native(df_native) >>> df.with_columns(a_clipped=nw.col("a").clip(-1, 3)) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a a_clipped | | 0 1 1 | | 1 2 2 | | 2 3 3 | └──────────────────┘ """ if upper_bound is None: return self._append_node( ExprNode(ExprKind.ELEMENTWISE, "clip_lower", lower_bound) ) if lower_bound is None: return self._append_node( ExprNode(ExprKind.ELEMENTWISE, "clip_upper", upper_bound) ) return self._append_node( ExprNode(ExprKind.ELEMENTWISE, "clip", lower_bound, upper_bound) ) def first(self) -> Self: """Get the first value. Notes: For lazy backends, this can only be used with `over`. We may introduce `min_by` in the future so it can be used as an aggregation. Examples: >>> import pandas as pd >>> import narwhals as nw >>> data = {"a": [1, 1, 2, 2], "b": ["foo", None, None, "baz"]} >>> df_native = pd.DataFrame(data) >>> df = nw.from_native(df_native) >>> df.select(nw.all().first()) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 1 foo | └──────────────────┘ >>> df.group_by("a").agg(nw.col("b").first()) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 1 foo | | 1 2 None | └──────────────────┘ """ return self._append_node(ExprNode(ExprKind.ORDERABLE_AGGREGATION, "first")) def last(self) -> Self: """Get the last value. Notes: For lazy backends, this can only be used with `over`. We may introduce `max_by` in the future so it can be used as an aggregation. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> data = {"a": [1, 1, 2, 2], "b": ["foo", None, None, "baz"]} >>> df_native = pa.table(data) >>> df = nw.from_native(df_native) >>> df.select(nw.all().last()) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | pyarrow.Table | | a: int64 | | b: string | | ---- | | a: [[2]] | | b: [["baz"]] | └──────────────────┘ >>> df.group_by("a").agg(nw.col("b").last()) ┌──────────────────┐ |Narwhals DataFrame| |------------------| |pyarrow.Table | |a: int64 | |b: string | |---- | |a: [[1,2]] | |b: [[null,"baz"]] | └──────────────────┘ """ return self._append_node(ExprNode(ExprKind.ORDERABLE_AGGREGATION, "last")) def mode(self, *, keep: ModeKeepStrategy = "all") -> Self: r"""Compute the most occurring value(s). Can return multiple values. Arguments: keep: Whether to keep all modes or any mode found. Remark that `keep='any'` is not deterministic for multimodal values. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 1, 2, 3], "b": [1, 1, 2, 2]}) >>> df = nw.from_native(df_native) >>> df.select(nw.col("a").mode()).sort("a") ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a | | 0 1 | └──────────────────┘ """ _supported_keep_values = ("all", "any") if keep not in _supported_keep_values: # pragma: no cover msg = f"`keep` must be one of {_supported_keep_values}, found '{keep}'" raise ValueError(msg) kind = ExprKind.AGGREGATION if keep == "any" else ExprKind.FILTRATION return self._append_node(ExprNode(kind, "mode", keep=keep)) def is_finite(self) -> Self: """Returns boolean values indicating which original values are finite. Warning: pandas handles null values differently from Polars and PyArrow. See [null_handling](../concepts/null_handling.md/) for reference. `is_finite` will return False for NaN and Null's in the Dask and pandas non-nullable backend, while for Polars, PyArrow and pandas nullable backends null values are kept as such. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"a": [float("nan"), float("inf"), 2.0, None]}) >>> df = nw.from_native(df_native) >>> df.with_columns(a_is_finite=nw.col("a").is_finite()) ┌──────────────────────┐ | Narwhals DataFrame | |----------------------| |shape: (4, 2) | |┌──────┮─────────────┐| |│ a ┆ a_is_finite │| |│ --- ┆ --- │| |│ f64 ┆ bool │| |╞══════╩═════════════╡| |│ NaN ┆ false │| |│ inf ┆ false │| |│ 2.0 ┆ true │| |│ null ┆ null │| |└──────â”ī─────────────┘| └──────────────────────┘ """ return self._append_node(ExprNode(ExprKind.ELEMENTWISE, "is_finite")) def cum_count(self, *, reverse: bool = False) -> Self: r"""Return the cumulative count of the non-null values in the column. Info: For lazy backends, this operation must be followed by `Expr.over` with `order_by` specified, see [order-dependence](../concepts/order_dependence.md). Arguments: reverse: reverse the operation Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": ["x", "k", None, "d"]}) >>> df = nw.from_native(df_native) >>> df.with_columns( ... nw.col("a").cum_count().alias("a_cum_count"), ... nw.col("a").cum_count(reverse=True).alias("a_cum_count_reverse"), ... ) ┌─────────────────────────────────────────┐ | Narwhals DataFrame | |-----------------------------------------| | a a_cum_count a_cum_count_reverse| |0 x 1 3| |1 k 2 2| |2 None 2 1| |3 d 3 1| └─────────────────────────────────────────┘ """ return self._append_node( ExprNode(ExprKind.ORDERABLE_WINDOW, "cum_count", reverse=reverse) ) def cum_min(self, *, reverse: bool = False) -> Self: r"""Return the cumulative min of the non-null values in the column. Info: For lazy backends, this operation must be followed by `Expr.over` with `order_by` specified, see [order-dependence](../concepts/order_dependence.md). Arguments: reverse: reverse the operation Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [3, 1, None, 2]}) >>> df = nw.from_native(df_native) >>> df.with_columns( ... nw.col("a").cum_min().alias("a_cum_min"), ... nw.col("a").cum_min(reverse=True).alias("a_cum_min_reverse"), ... ) ┌────────────────────────────────────┐ | Narwhals DataFrame | |------------------------------------| | a a_cum_min a_cum_min_reverse| |0 3.0 3.0 1.0| |1 1.0 1.0 1.0| |2 NaN NaN NaN| |3 2.0 1.0 2.0| └────────────────────────────────────┘ """ return self._append_node( ExprNode(ExprKind.ORDERABLE_WINDOW, "cum_min", reverse=reverse) ) def cum_max(self, *, reverse: bool = False) -> Self: r"""Return the cumulative max of the non-null values in the column. Info: For lazy backends, this operation must be followed by `Expr.over` with `order_by` specified, see [order-dependence](../concepts/order_dependence.md). Arguments: reverse: reverse the operation Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 3, None, 2]}) >>> df = nw.from_native(df_native) >>> df.with_columns( ... nw.col("a").cum_max().alias("a_cum_max"), ... nw.col("a").cum_max(reverse=True).alias("a_cum_max_reverse"), ... ) ┌────────────────────────────────────┐ | Narwhals DataFrame | |------------------------------------| | a a_cum_max a_cum_max_reverse| |0 1.0 1.0 3.0| |1 3.0 3.0 3.0| |2 NaN NaN NaN| |3 2.0 3.0 2.0| └────────────────────────────────────┘ """ return self._append_node( ExprNode(ExprKind.ORDERABLE_WINDOW, "cum_max", reverse=reverse) ) def cum_prod(self, *, reverse: bool = False) -> Self: r"""Return the cumulative product of the non-null values in the column. Info: For lazy backends, this operation must be followed by `Expr.over` with `order_by` specified, see [order-dependence](../concepts/order_dependence.md). Arguments: reverse: reverse the operation Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1, 3, None, 2]}) >>> df = nw.from_native(df_native) >>> df.with_columns( ... nw.col("a").cum_prod().alias("a_cum_prod"), ... nw.col("a").cum_prod(reverse=True).alias("a_cum_prod_reverse"), ... ) ┌──────────────────────────────────────┐ | Narwhals DataFrame | |--------------------------------------| | a a_cum_prod a_cum_prod_reverse| |0 1.0 1.0 6.0| |1 3.0 3.0 6.0| |2 NaN NaN NaN| |3 2.0 6.0 2.0| └──────────────────────────────────────┘ """ return self._append_node( ExprNode(ExprKind.ORDERABLE_WINDOW, "cum_prod", reverse=reverse) ) def rolling_sum( self, window_size: int, *, min_samples: int | None = None, center: bool = False ) -> Self: """Apply a rolling sum (moving sum) over the values. A window of length `window_size` will traverse the values. The resulting values will be aggregated to their sum. The window at a given row will include the row itself and the `window_size - 1` elements before it. Info: For lazy backends, this operation must be followed by `Expr.over` with `order_by` specified, see [order-dependence](../concepts/order_dependence.md). Arguments: window_size: The length of the window in number of elements. It must be a strictly positive integer. min_samples: The number of values in the window that should be non-null before computing a result. If set to `None` (default), it will be set equal to `window_size`. If provided, it must be a strictly positive integer, and less than or equal to `window_size` center: Set the labels at the center of the window. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1.0, 2.0, None, 4.0]}) >>> df = nw.from_native(df_native) >>> df.with_columns( ... a_rolling_sum=nw.col("a").rolling_sum(window_size=3, min_samples=1) ... ) ┌─────────────────────┐ | Narwhals DataFrame | |---------------------| | a a_rolling_sum| |0 1.0 1.0| |1 2.0 3.0| |2 NaN 3.0| |3 4.0 6.0| └─────────────────────┘ """ window_size, min_samples = _validate_rolling_arguments( window_size=window_size, min_samples=min_samples ) return self._append_node( ExprNode( ExprKind.ORDERABLE_WINDOW, "rolling_sum", window_size=window_size, min_samples=min_samples, center=center, ) ) def rolling_mean( self, window_size: int, *, min_samples: int | None = None, center: bool = False ) -> Self: """Apply a rolling mean (moving mean) over the values. A window of length `window_size` will traverse the values. The resulting values will be aggregated to their mean. The window at a given row will include the row itself and the `window_size - 1` elements before it. Info: For lazy backends, this operation must be followed by `Expr.over` with `order_by` specified, see [order-dependence](../concepts/order_dependence.md). Arguments: window_size: The length of the window in number of elements. It must be a strictly positive integer. min_samples: The number of values in the window that should be non-null before computing a result. If set to `None` (default), it will be set equal to `window_size`. If provided, it must be a strictly positive integer, and less than or equal to `window_size` center: Set the labels at the center of the window. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1.0, 2.0, None, 4.0]}) >>> df = nw.from_native(df_native) >>> df.with_columns( ... a_rolling_mean=nw.col("a").rolling_mean(window_size=3, min_samples=1) ... ) ┌──────────────────────┐ | Narwhals DataFrame | |----------------------| | a a_rolling_mean| |0 1.0 1.0| |1 2.0 1.5| |2 NaN 1.5| |3 4.0 3.0| └──────────────────────┘ """ window_size, min_samples = _validate_rolling_arguments( window_size=window_size, min_samples=min_samples ) return self._append_node( ExprNode( ExprKind.ORDERABLE_WINDOW, "rolling_mean", window_size=window_size, min_samples=min_samples, center=center, ) ) def rolling_var( self, window_size: int, *, min_samples: int | None = None, center: bool = False, ddof: int = 1, ) -> Self: """Apply a rolling variance (moving variance) over the values. A window of length `window_size` will traverse the values. The resulting values will be aggregated to their variance. The window at a given row will include the row itself and the `window_size - 1` elements before it. Info: For lazy backends, this operation must be followed by `Expr.over` with `order_by` specified, see [order-dependence](../concepts/order_dependence.md). Arguments: window_size: The length of the window in number of elements. It must be a strictly positive integer. min_samples: The number of values in the window that should be non-null before computing a result. If set to `None` (default), it will be set equal to `window_size`. If provided, it must be a strictly positive integer, and less than or equal to `window_size`. center: Set the labels at the center of the window. ddof: Delta Degrees of Freedom; the divisor for a length N window is N - ddof. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1.0, 2.0, None, 4.0]}) >>> df = nw.from_native(df_native) >>> df.with_columns( ... a_rolling_var=nw.col("a").rolling_var(window_size=3, min_samples=1) ... ) ┌─────────────────────┐ | Narwhals DataFrame | |---------------------| | a a_rolling_var| |0 1.0 NaN| |1 2.0 0.5| |2 NaN 0.5| |3 4.0 2.0| └─────────────────────┘ """ window_size, min_samples = _validate_rolling_arguments( window_size=window_size, min_samples=min_samples ) return self._append_node( ExprNode( ExprKind.ORDERABLE_WINDOW, "rolling_var", ddof=ddof, window_size=window_size, min_samples=min_samples, center=center, ) ) def rolling_std( self, window_size: int, *, min_samples: int | None = None, center: bool = False, ddof: int = 1, ) -> Self: """Apply a rolling standard deviation (moving standard deviation) over the values. A window of length `window_size` will traverse the values. The resulting values will be aggregated to their standard deviation. The window at a given row will include the row itself and the `window_size - 1` elements before it. Info: For lazy backends, this operation must be followed by `Expr.over` with `order_by` specified, see [order-dependence](../concepts/order_dependence.md). Arguments: window_size: The length of the window in number of elements. It must be a strictly positive integer. min_samples: The number of values in the window that should be non-null before computing a result. If set to `None` (default), it will be set equal to `window_size`. If provided, it must be a strictly positive integer, and less than or equal to `window_size`. center: Set the labels at the center of the window. ddof: Delta Degrees of Freedom; the divisor for a length N window is N - ddof. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [1.0, 2.0, None, 4.0]}) >>> df = nw.from_native(df_native) >>> df.with_columns( ... a_rolling_std=nw.col("a").rolling_std(window_size=3, min_samples=1) ... ) ┌─────────────────────┐ | Narwhals DataFrame | |---------------------| | a a_rolling_std| |0 1.0 NaN| |1 2.0 0.707107| |2 NaN 0.707107| |3 4.0 1.414214| └─────────────────────┘ """ window_size, min_samples = _validate_rolling_arguments( window_size=window_size, min_samples=min_samples ) return self._append_node( ExprNode( ExprKind.ORDERABLE_WINDOW, "rolling_std", ddof=ddof, window_size=window_size, min_samples=min_samples, center=center, ) ) def rank(self, method: RankMethod = "average", *, descending: bool = False) -> Self: """Assign ranks to data, dealing with ties appropriately. Notes: The resulting dtype may differ between backends. Info: For lazy backends, this operation must be followed by `Expr.over` with `order_by` specified, see [order-dependence](../concepts/order_dependence.md). Arguments: method: The method used to assign ranks to tied elements. The following methods are available (default is 'average') - *"average"*: The average of the ranks that would have been assigned to all the tied values is assigned to each value. - *"min"*: The minimum of the ranks that would have been assigned to all the tied values is assigned to each value. (This is also referred to as "competition" ranking.) - *"max"*: The maximum of the ranks that would have been assigned to all the tied values is assigned to each value. - *"dense"*: Like "min", but the rank of the next highest element is assigned the rank immediately after those assigned to the tied elements. - *"ordinal"*: All values are given a distinct rank, corresponding to the order that the values occur in the Series. descending: Rank in descending order. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"a": [3, 6, 1, 1, 6]}) >>> df = nw.from_native(df_native) >>> result = df.with_columns(rank=nw.col("a").rank(method="dense")) >>> result ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a rank | | 0 3 2.0 | | 1 6 3.0 | | 2 1 1.0 | | 3 1 1.0 | | 4 6 3.0 | └──────────────────┘ """ supported_rank_methods = {"average", "min", "max", "dense", "ordinal"} if method not in supported_rank_methods: msg = ( "Ranking method must be one of {'average', 'min', 'max', 'dense', 'ordinal'}. " f"Found '{method}'" ) raise ValueError(msg) return self._append_node( ExprNode(ExprKind.WINDOW, "rank", method=method, descending=descending) ) def log(self, base: float = math.e) -> Self: r"""Compute the logarithm to a given base. Arguments: base: Given base, defaults to `e` Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"values": [1, 2, 4]}) >>> df = nw.from_native(df_native) >>> result = df.with_columns( ... log=nw.col("values").log(), log_2=nw.col("values").log(base=2) ... ) >>> result ┌────────────────────────────────────────────────┐ | Narwhals DataFrame | |------------------------------------------------| |pyarrow.Table | |values: int64 | |log: double | |log_2: double | |---- | |values: [[1,2,4]] | |log: [[0,0.6931471805599453,1.3862943611198906]]| |log_2: [[0,1,2]] | └────────────────────────────────────────────────┘ """ return self._append_node(ExprNode(ExprKind.ELEMENTWISE, "log", base=base)) def exp(self) -> Self: r"""Compute the exponent. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"values": [-1, 0, 1]}) >>> df = nw.from_native(df_native) >>> result = df.with_columns(exp=nw.col("values").exp()) >>> result ┌────────────────────────────────────────────────┐ | Narwhals DataFrame | |------------------------------------------------| |pyarrow.Table | |values: int64 | |exp: double | |---- | |values: [[-1,0,1]] | |exp: [[0.36787944117144233,1,2.718281828459045]]| └────────────────────────────────────────────────┘ """ return self._append_node(ExprNode(ExprKind.ELEMENTWISE, "exp")) def sqrt(self) -> Self: r"""Compute the square root. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"values": [1, 4, 9]}) >>> df = nw.from_native(df_native) >>> result = df.with_columns(sqrt=nw.col("values").sqrt()) >>> result ┌──────────────────┐ |Narwhals DataFrame| |------------------| |pyarrow.Table | |values: int64 | |sqrt: double | |---- | |values: [[1,4,9]] | |sqrt: [[1,2,3]] | └──────────────────┘ """ return self._append_node(ExprNode(ExprKind.ELEMENTWISE, "sqrt")) def is_close( # noqa: PLR0914 self, other: Expr | Series[Any] | NumericLiteral, *, abs_tol: float = 0.0, rel_tol: float = 1e-09, nans_equal: bool = False, ) -> Self: r"""Check if this expression is close, i.e. almost equal, to the other expression. Two values `a` and `b` are considered close if the following condition holds: $$ |a-b| \le max \{ \text{rel\_tol} \cdot max \{ |a|, |b| \}, \text{abs\_tol} \} $$ Arguments: other: Values to compare with. abs_tol: Absolute tolerance. This is the maximum allowed absolute difference between two values. Must be non-negative. rel_tol: Relative tolerance. This is the maximum allowed difference between two values, relative to the larger absolute value. Must be in the range [0, 1). nans_equal: Whether NaN values should be considered equal. Notes: The implementation of this method is symmetric and mirrors the behavior of `math.isclose`. Specifically note that this behavior is different to `numpy.isclose`. Examples: >>> import duckdb >>> import pyarrow as pa >>> import narwhals as nw >>> >>> data = { ... "x": [1.0, float("inf"), 1.41, None, float("nan")], ... "y": [1.2, float("inf"), 1.40, None, float("nan")], ... } >>> _table = pa.table(data) >>> df_native = duckdb.table("_table") >>> df = nw.from_native(df_native) >>> df.with_columns( ... is_close=nw.col("x").is_close( ... nw.col("y"), abs_tol=0.1, nans_equal=True ... ) ... ) ┌──────────────────────────────┐ | Narwhals LazyFrame | |------------------------------| |┌────────┮────────┮──────────┐| |│ x │ y │ is_close │| |│ double │ double │ boolean │| |├────────┾────────┾──────────â”Ī| |│ 1.0 │ 1.2 │ false │| |│ inf │ inf │ true │| |│ 1.41 │ 1.4 │ true │| |│ NULL │ NULL │ NULL │| |│ nan │ nan │ true │| |└────────â”ī────────â”ī──────────┘| └──────────────────────────────┘ """ if abs_tol < 0: msg = f"`abs_tol` must be non-negative but got {abs_tol}" raise ComputeError(msg) if not (0 <= rel_tol < 1): msg = f"`rel_tol` must be in the range [0, 1) but got {rel_tol}" raise ComputeError(msg) from decimal import Decimal other_abs: Expr | Series[Any] | NumericLiteral other_is_nan: Expr | Series[Any] | bool other_is_inf: Expr | Series[Any] | bool other_is_not_inf: Expr | Series[Any] | bool if isinstance(other, (float, int, Decimal)): from math import isinf, isnan # NOTE: See https://discuss.python.org/t/inferred-type-of-function-that-calls-dunder-abs-abs/101447 other_abs = other.__abs__() other_is_nan = isnan(other) other_is_inf = isinf(other) # Define the other_is_not_inf variable to prevent triggering the following warning: # > DeprecationWarning: Bitwise inversion '~' on bool is deprecated and will be # > removed in Python 3.16. other_is_not_inf = not other_is_inf else: other_abs, other_is_nan = other.abs(), other.is_nan() other_is_not_inf = other.is_finite() | other_is_nan other_is_inf = ~other_is_not_inf rel_threshold = self.abs().clip(lower_bound=other_abs, upper_bound=None) * rel_tol tolerance = rel_threshold.clip(lower_bound=abs_tol, upper_bound=None) self_is_nan = self.is_nan() self_is_not_inf = self.is_finite() | self_is_nan # Values are close if abs_diff <= tolerance, and both finite is_close = ( ((self - other).abs() <= tolerance) & self_is_not_inf & other_is_not_inf ) # Handle infinity cases: infinities are close/equal if they have the same sign self_sign, other_sign = self > 0, other > 0 is_same_inf = (~self_is_not_inf) & other_is_inf & (self_sign == other_sign) # Handle nan cases: # * If any value is NaN, then False (via `& ~either_nan`) # * However, if `nans_equals = True` and if _both_ values are NaN, then True either_nan = self_is_nan | other_is_nan result = (is_close | is_same_inf) & ~either_nan if nans_equal: both_nan = self_is_nan & other_is_nan result = result | both_nan return result @property def str(self) -> ExprStringNamespace[Self]: return ExprStringNamespace(self) @property def dt(self) -> ExprDateTimeNamespace[Self]: return ExprDateTimeNamespace(self) @property def cat(self) -> ExprCatNamespace[Self]: return ExprCatNamespace(self) @property def name(self) -> ExprNameNamespace[Self]: return ExprNameNamespace(self) @property def list(self) -> ExprListNamespace[Self]: return ExprListNamespace(self) @property def struct(self) -> ExprStructNamespace[Self]: return ExprStructNamespace(self) __all__ = ["Expr"] narwhals-2.11.0/narwhals/expr_cat.py0000644000000000000000000000234213615410400014331 0ustar00from __future__ import annotations from typing import TYPE_CHECKING, Generic, TypeVar from narwhals._expression_parsing import ExprKind, ExprNode if TYPE_CHECKING: from narwhals.expr import Expr ExprT = TypeVar("ExprT", bound="Expr") class ExprCatNamespace(Generic[ExprT]): def __init__(self, expr: ExprT) -> None: self._expr = expr def get_categories(self) -> ExprT: """Get unique categories from column. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame( ... {"fruits": ["apple", "mango", "mango"]}, ... schema={"fruits": pl.Categorical}, ... ) >>> df = nw.from_native(df_native) >>> df.select(nw.col("fruits").cat.get_categories()).to_native() shape: (2, 1) ┌────────┐ │ fruits │ │ --- │ │ str │ ╞════════╡ │ apple │ │ mango │ └────────┘ """ return self._expr._append_node( ExprNode(ExprKind.FILTRATION, "cat.get_categories") ) narwhals-2.11.0/narwhals/expr_dt.py0000644000000000000000000007611013615410400014175 0ustar00from __future__ import annotations from typing import TYPE_CHECKING, Generic, TypeVar from narwhals._expression_parsing import ExprKind, ExprNode if TYPE_CHECKING: from narwhals.expr import Expr from narwhals.typing import TimeUnit ExprT = TypeVar("ExprT", bound="Expr") class ExprDateTimeNamespace(Generic[ExprT]): def __init__(self, expr: ExprT) -> None: self._expr = expr def date(self) -> ExprT: """Extract the date from underlying DateTime representation. Raises: NotImplementedError: If pandas default backend is being used. Examples: >>> from datetime import datetime >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame( ... {"a": [datetime(2012, 1, 7, 10), datetime(2027, 12, 13)]} ... ) >>> df = nw.from_native(df_native) >>> df.select(nw.col("a").dt.date()).to_native() shape: (2, 1) ┌────────────┐ │ a │ │ --- │ │ date │ ╞════════════╡ │ 2012-01-07 │ │ 2027-12-13 │ └────────────┘ """ return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "dt.date")) def year(self) -> ExprT: """Extract year from underlying DateTime representation. Returns the year number in the calendar date. Examples: >>> from datetime import datetime >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame( ... {"a": [datetime(1978, 6, 1), datetime(2065, 1, 1)]} ... ) >>> df = nw.from_native(df_native) >>> df.with_columns(nw.col("a").dt.year().alias("year")) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a year| |0 1978-06-01 1978| |1 2065-01-01 2065| └──────────────────┘ """ return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "dt.year")) def month(self) -> ExprT: """Extract month from underlying DateTime representation. Returns the month number starting from 1. The return value ranges from 1 to 12. Examples: >>> from datetime import datetime >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"a": [datetime(1978, 6, 1), datetime(2065, 1, 1)]}) >>> df = nw.from_native(df_native) >>> df.with_columns(nw.col("a").dt.month().alias("month")).to_native() pyarrow.Table a: timestamp[us] month: int64 ---- a: [[1978-06-01 00:00:00.000000,2065-01-01 00:00:00.000000]] month: [[6,1]] """ return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "dt.month")) def day(self) -> ExprT: """Extract day from underlying DateTime representation. Returns the day of month starting from 1. The return value ranges from 1 to 31. (The last day of month differs by months.) Examples: >>> from datetime import datetime >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"a": [datetime(1978, 6, 1), datetime(2065, 1, 1)]}) >>> df = nw.from_native(df_native) >>> df.with_columns(nw.col("a").dt.day().alias("day")).to_native() pyarrow.Table a: timestamp[us] day: int64 ---- a: [[1978-06-01 00:00:00.000000,2065-01-01 00:00:00.000000]] day: [[1,1]] """ return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "dt.day")) def hour(self) -> ExprT: """Extract hour from underlying DateTime representation. Returns the hour number from 0 to 23. Examples: >>> from datetime import datetime >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame( ... {"a": [datetime(1978, 1, 1, 1), datetime(2065, 1, 1, 10)]} ... ) >>> df = nw.from_native(df_native) >>> df.with_columns(nw.col("a").dt.hour().alias("hour")) ┌──────────────────────────────┐ | Narwhals DataFrame | |------------------------------| |shape: (2, 2) | |┌─────────────────────┮──────┐| |│ a ┆ hour │| |│ --- ┆ --- │| |│ datetime[Ξs] ┆ i8 │| |╞═════════════════════╩══════╡| |│ 1978-01-01 01:00:00 ┆ 1 │| |│ 2065-01-01 10:00:00 ┆ 10 │| |└─────────────────────â”ī──────┘| └──────────────────────────────┘ """ return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "dt.hour")) def minute(self) -> ExprT: """Extract minutes from underlying DateTime representation. Returns the minute number from 0 to 59. Examples: >>> from datetime import datetime >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame( ... {"a": [datetime(1978, 1, 1, 1, 1), datetime(2065, 1, 1, 10, 20)]} ... ) >>> df = nw.from_native(df_native) >>> df.with_columns(nw.col("a").dt.minute().alias("minute")).to_native() a minute 0 1978-01-01 01:01:00 1 1 2065-01-01 10:20:00 20 """ return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "dt.minute")) def second(self) -> ExprT: """Extract seconds from underlying DateTime representation. Examples: >>> from datetime import datetime >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table( ... { ... "a": [ ... datetime(1978, 1, 1, 1, 1, 1), ... datetime(2065, 1, 1, 10, 20, 30), ... ] ... } ... ) >>> df = nw.from_native(df_native) >>> df.with_columns(nw.col("a").dt.second().alias("second")).to_native() pyarrow.Table a: timestamp[us] second: int64 ---- a: [[1978-01-01 01:01:01.000000,2065-01-01 10:20:30.000000]] second: [[1,30]] """ return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "dt.second")) def millisecond(self) -> ExprT: """Extract milliseconds from underlying DateTime representation. Examples: >>> from datetime import datetime >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table( ... { ... "a": [ ... datetime(1978, 1, 1, 1, 1, 1, 0), ... datetime(2065, 1, 1, 10, 20, 30, 67000), ... ] ... } ... ) >>> df = nw.from_native(df_native) >>> df.with_columns( ... nw.col("a").dt.millisecond().alias("millisecond") ... ).to_native() pyarrow.Table a: timestamp[us] millisecond: int64 ---- a: [[1978-01-01 01:01:01.000000,2065-01-01 10:20:30.067000]] millisecond: [[0,67]] """ return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "dt.millisecond")) def microsecond(self) -> ExprT: """Extract microseconds from underlying DateTime representation. Examples: >>> from datetime import datetime >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table( ... { ... "a": [ ... datetime(1978, 1, 1, 1, 1, 1, 0), ... datetime(2065, 1, 1, 10, 20, 30, 67000), ... ] ... } ... ) >>> df = nw.from_native(df_native) >>> df.with_columns( ... nw.col("a").dt.microsecond().alias("microsecond") ... ).to_native() pyarrow.Table a: timestamp[us] microsecond: int64 ---- a: [[1978-01-01 01:01:01.000000,2065-01-01 10:20:30.067000]] microsecond: [[0,67000]] """ return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "dt.microsecond")) def nanosecond(self) -> ExprT: """Extract Nanoseconds from underlying DateTime representation. Examples: >>> from datetime import datetime >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table( ... { ... "a": [ ... datetime(1978, 1, 1, 1, 1, 1, 0), ... datetime(2065, 1, 1, 10, 20, 30, 67000), ... ] ... } ... ) >>> df = nw.from_native(df_native) >>> df.with_columns( ... nw.col("a").dt.nanosecond().alias("nanosecond") ... ).to_native() pyarrow.Table a: timestamp[us] nanosecond: int64 ---- a: [[1978-01-01 01:01:01.000000,2065-01-01 10:20:30.067000]] nanosecond: [[0,67000000]] """ return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "dt.nanosecond")) def ordinal_day(self) -> ExprT: """Get ordinal day. Examples: >>> from datetime import datetime >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame( ... {"a": [datetime(2020, 1, 1), datetime(2020, 8, 3)]} ... ) >>> df = nw.from_native(df_native) >>> df.with_columns(a_ordinal_day=nw.col("a").dt.ordinal_day()) ┌───────────────────────────┐ | Narwhals DataFrame | |---------------------------| | a a_ordinal_day| |0 2020-01-01 1| |1 2020-08-03 216| └───────────────────────────┘ """ return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "dt.ordinal_day")) def weekday(self) -> ExprT: """Extract the week day from the underlying Date representation. Note that Monday = 1 and Sunday = 7. Examples: >>> from datetime import datetime >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame( ... {"a": [datetime(2020, 1, 1), datetime(2020, 8, 3)]} ... ) >>> df = nw.from_native(df_native) >>> df.with_columns(a_week_day=nw.col("a").dt.weekday()) ┌────────────────────────┐ | Narwhals DataFrame | |------------------------| | a a_week_day| |0 2020-01-01 3| |1 2020-08-03 1| └────────────────────────┘ """ return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "dt.weekday")) def total_minutes(self) -> ExprT: """Get total minutes. Notes: The function outputs the total minutes in the int dtype by default, however, pandas may change the dtype to float when there are missing values, consider using `fill_null()` and `cast` in this case. Examples: >>> from datetime import timedelta >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame( ... {"a": [timedelta(minutes=10), timedelta(minutes=20, seconds=40)]} ... ) >>> df = nw.from_native(df_native) >>> df.with_columns( ... a_total_minutes=nw.col("a").dt.total_minutes() ... ).to_native() shape: (2, 2) ┌──────────────┮─────────────────┐ │ a ┆ a_total_minutes │ │ --- ┆ --- │ │ duration[Ξs] ┆ i64 │ ╞══════════════╩═════════════════╡ │ 10m ┆ 10 │ │ 20m 40s ┆ 20 │ └──────────────â”ī─────────────────┘ """ return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "dt.total_minutes")) def total_seconds(self) -> ExprT: """Get total seconds. Notes: The function outputs the total seconds in the int dtype by default, however, pandas may change the dtype to float when there are missing values, consider using `fill_null()` and `cast` in this case. Examples: >>> from datetime import timedelta >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame( ... {"a": [timedelta(seconds=10), timedelta(seconds=20, milliseconds=40)]} ... ) >>> df = nw.from_native(df_native) >>> df.with_columns( ... a_total_seconds=nw.col("a").dt.total_seconds() ... ).to_native() shape: (2, 2) ┌──────────────┮─────────────────┐ │ a ┆ a_total_seconds │ │ --- ┆ --- │ │ duration[Ξs] ┆ i64 │ ╞══════════════╩═════════════════╡ │ 10s ┆ 10 │ │ 20s 40ms ┆ 20 │ └──────────────â”ī─────────────────┘ """ return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "dt.total_seconds")) def total_milliseconds(self) -> ExprT: """Get total milliseconds. Notes: The function outputs the total milliseconds in the int dtype by default, however, pandas may change the dtype to float when there are missing values, consider using `fill_null()` and `cast` in this case. Examples: >>> from datetime import timedelta >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame( ... { ... "a": [ ... timedelta(milliseconds=10), ... timedelta(milliseconds=20, microseconds=40), ... ] ... } ... ) >>> df = nw.from_native(df_native) >>> df.with_columns( ... a_total_milliseconds=nw.col("a").dt.total_milliseconds() ... ).to_native() shape: (2, 2) ┌──────────────┮──────────────────────┐ │ a ┆ a_total_milliseconds │ │ --- ┆ --- │ │ duration[Ξs] ┆ i64 │ ╞══════════════╩══════════════════════╡ │ 10ms ┆ 10 │ │ 20040Âĩs ┆ 20 │ └──────────────â”ī──────────────────────┘ """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "dt.total_milliseconds") ) def total_microseconds(self) -> ExprT: """Get total microseconds. Notes: The function outputs the total microseconds in the int dtype by default, however, pandas may change the dtype to float when there are missing values, consider using `fill_null()` and `cast` in this case. Examples: >>> from datetime import timedelta >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table( ... { ... "a": [ ... timedelta(microseconds=10), ... timedelta(milliseconds=1, microseconds=200), ... ] ... } ... ) >>> df = nw.from_native(df_native) >>> df.with_columns( ... a_total_microseconds=nw.col("a").dt.total_microseconds() ... ).to_native() pyarrow.Table a: duration[us] a_total_microseconds: int64 ---- a: [[10,1200]] a_total_microseconds: [[10,1200]] """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "dt.total_microseconds") ) def total_nanoseconds(self) -> ExprT: """Get total nanoseconds. Notes: The function outputs the total nanoseconds in the int dtype by default, however, pandas may change the dtype to float when there are missing values, consider using `fill_null()` and `cast` in this case. Examples: >>> from datetime import timedelta >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame( ... { ... "a": pd.to_datetime( ... [ ... "2024-01-01 00:00:00.000000001", ... "2024-01-01 00:00:00.000000002", ... ] ... ) ... } ... ) >>> df = nw.from_native(df_native) >>> df.with_columns( ... a_diff_total_nanoseconds=nw.col("a").diff().dt.total_nanoseconds() ... ).to_native() a a_diff_total_nanoseconds 0 2024-01-01 00:00:00.000000001 NaN 1 2024-01-01 00:00:00.000000002 1.0 """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "dt.total_nanoseconds") ) def to_string(self, format: str) -> ExprT: """Convert a Date/Time/Datetime column into a String column with the given format. Arguments: format: Format to format temporal column with. Notes: Unfortunately, different libraries interpret format directives a bit differently. - Chrono, the library used by Polars, uses `"%.f"` for fractional seconds, whereas pandas and Python stdlib use `".%f"`. - PyArrow interprets `"%S"` as "seconds, including fractional seconds" whereas most other tools interpret it as "just seconds, as 2 digits". --- Therefore, we make the following adjustments. - for pandas-like libraries, we replace `"%S.%f"` with `"%S%.f"`. - for PyArrow, we replace `"%S.%f"` with `"%S"`. --- Workarounds like these don't make us happy, and we try to avoid them as much as possible, but here we feel like it's the best compromise. If you just want to format a date/datetime Series as a local datetime string, and have it work as consistently as possible across libraries, we suggest using: - `"%Y-%m-%dT%H:%M:%S%.f"` for datetimes - `"%Y-%m-%d"` for dates --- Though note that, even then, different tools may return a different number of trailing zeros. Nonetheless, this is probably consistent enough for most applications. If you have an application where this is not enough, please open an issue and let us know. Examples: >>> from datetime import datetime >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame( ... {"a": [datetime(2020, 3, 1), datetime(2020, 5, 1)]} ... ) >>> df = nw.from_native(df_native) >>> df.select(nw.col("a").dt.to_string("%Y/%m/%d %H:%M:%S")) ┌───────────────────────┐ | Narwhals DataFrame | |-----------------------| |shape: (2, 1) | |┌─────────────────────┐| |│ a │| |│ --- │| |│ str │| |╞═════════════════════╡| |│ 2020/03/01 00:00:00 │| |│ 2020/05/01 00:00:00 │| |└─────────────────────┘| └───────────────────────┘ """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "dt.to_string", format=format) ) def replace_time_zone(self, time_zone: str | None) -> ExprT: """Replace time zone. Arguments: time_zone: Target time zone. Examples: >>> from datetime import datetime, timezone >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame( ... { ... "a": [ ... datetime(2024, 1, 1, tzinfo=timezone.utc), ... datetime(2024, 1, 2, tzinfo=timezone.utc), ... ] ... } ... ) >>> df = nw.from_native(df_native) >>> df.select(nw.col("a").dt.replace_time_zone("Asia/Kathmandu")).to_native() a 0 2024-01-01 00:00:00+05:45 1 2024-01-02 00:00:00+05:45 """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "dt.replace_time_zone", time_zone=time_zone) ) def convert_time_zone(self, time_zone: str) -> ExprT: """Convert to a new time zone. If converting from a time-zone-naive column, then conversion happens as if converting from UTC. Arguments: time_zone: Target time zone. Examples: >>> from datetime import datetime, timezone >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame( ... { ... "a": [ ... datetime(2024, 1, 1, tzinfo=timezone.utc), ... datetime(2024, 1, 2, tzinfo=timezone.utc), ... ] ... } ... ) >>> df = nw.from_native(df_native) >>> df.select(nw.col("a").dt.convert_time_zone("Asia/Kathmandu")).to_native() a 0 2024-01-01 05:45:00+05:45 1 2024-01-02 05:45:00+05:45 """ if time_zone is None: msg = "Target `time_zone` cannot be `None` in `convert_time_zone`. Please use `replace_time_zone(None)` if you want to remove the time zone." raise TypeError(msg) return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "dt.convert_time_zone", time_zone=time_zone) ) def timestamp(self, time_unit: TimeUnit = "us") -> ExprT: """Return a timestamp in the given time unit. Arguments: time_unit: One of - 'ns': nanosecond. - 'us': microsecond. - 'ms': millisecond. Examples: >>> from datetime import date >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"date": [date(2001, 1, 1), None]}) >>> df = nw.from_native(df_native) >>> df.with_columns(nw.col("date").dt.timestamp("ms").alias("timestamp_ms")) ┌─────────────────────────────┐ | Narwhals DataFrame | |-----------------------------| |shape: (2, 2) | |┌────────────┮──────────────┐| |│ date ┆ timestamp_ms │| |│ --- ┆ --- │| |│ date ┆ i64 │| |╞════════════╩══════════════╡| |│ 2001-01-01 ┆ 978307200000 │| |│ null ┆ null │| |└────────────â”ī──────────────┘| └─────────────────────────────┘ """ if time_unit not in {"ns", "us", "ms"}: msg = ( "invalid `time_unit`" f"\n\nExpected one of {{'ns', 'us', 'ms'}}, got {time_unit!r}." ) raise ValueError(msg) return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "dt.timestamp", time_unit=time_unit) ) def truncate(self, every: str) -> ExprT: """Divide the date/datetime range into buckets. Arguments: every: Length of bucket. Must be of form ``, where `multiple` is a positive integer and `unit` is one of - 'ns': nanosecond. - 'us': microsecond. - 'ms': millisecond. - 's': second. - 'm': minute. - 'h': hour. - 'd': day. - 'mo': month. - 'q': quarter. - 'y': year. Examples: >>> from datetime import datetime >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"datetime": [datetime(2021, 3, 1, 12, 34)]}) >>> df = nw.from_native(df_native) >>> df.with_columns( ... nw.col("datetime").dt.truncate("1h").alias("datetime_trunc") ... ) ┌─────────────────────────────────────────────┐ | Narwhals DataFrame | |---------------------------------------------| |shape: (1, 2) | |┌─────────────────────┮─────────────────────┐| |│ datetime ┆ datetime_trunc │| |│ --- ┆ --- │| |│ datetime[Ξs] ┆ datetime[Ξs] │| |╞═════════════════════╩═════════════════════╡| |│ 2021-03-01 12:34:00 ┆ 2021-03-01 12:00:00 │| |└─────────────────────â”ī─────────────────────┘| └─────────────────────────────────────────────┘ """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "dt.truncate", every=every) ) def offset_by(self, by: str) -> ExprT: """Offset this date by a relative time offset. Arguments: by: The offset. Must be of form ``, where `multiple` is a positive integer and `unit` is one of - 'ns': nanosecond. - 'us': microsecond. - 'ms': millisecond. - 's': second. - 'm': minute. - 'h': hour. - 'd': day. - 'mo': month. - 'q': quarter. - 'y': year. Examples: >>> from datetime import datetime >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"datetime": [datetime(2021, 3, 1, 12, 34)]}) >>> df = nw.from_native(df_native) >>> df.with_columns( ... nw.col("datetime").dt.offset_by("1h").alias("datetime_offset_by_1h") ... ) ┌───────────────────────────────────────────────┐ | Narwhals DataFrame | |-----------------------------------------------| |shape: (1, 2) | |┌─────────────────────┮───────────────────────┐| |│ datetime ┆ datetime_offset_by_1h │| |│ --- ┆ --- │| |│ datetime[Ξs] ┆ datetime[Ξs] │| |╞═════════════════════╩═══════════════════════╡| |│ 2021-03-01 12:34:00 ┆ 2021-03-01 13:34:00 │| |└─────────────────────â”ī───────────────────────┘| └───────────────────────────────────────────────┘ """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "dt.offset_by", by=by) ) narwhals-2.11.0/narwhals/expr_list.py0000644000000000000000000001477013615410400014545 0ustar00from __future__ import annotations from typing import TYPE_CHECKING, Generic, TypeVar from narwhals._expression_parsing import ExprKind, ExprNode if TYPE_CHECKING: from narwhals.expr import Expr from narwhals.typing import NonNestedLiteral ExprT = TypeVar("ExprT", bound="Expr") class ExprListNamespace(Generic[ExprT]): def __init__(self, expr: ExprT) -> None: self._expr = expr def len(self) -> ExprT: """Return the number of elements in each list. Null values count towards the total. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"a": [[1, 2], [3, 4, None], None, []]}) >>> df = nw.from_native(df_native) >>> df.with_columns(a_len=nw.col("a").list.len()) ┌────────────────────────┐ | Narwhals DataFrame | |------------------------| |shape: (4, 2) | |┌──────────────┮───────┐| |│ a ┆ a_len │| |│ --- ┆ --- │| |│ list[i64] ┆ u32 │| |╞══════════════╩═══════╡| |│ [1, 2] ┆ 2 │| |│ [3, 4, null] ┆ 3 │| |│ null ┆ null │| |│ [] ┆ 0 │| |└──────────────â”ī───────┘| └────────────────────────┘ """ return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "list.len")) def unique(self) -> ExprT: """Get the unique/distinct values in the list. Null values are included in the result. The order of unique values is not guaranteed. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"a": [[1, 1, 2], [3, 3, None], None, []]}) >>> df = nw.from_native(df_native) >>> df.with_columns(a_unique=nw.col("a").list.unique()) ┌────────────────────────────┐ | Narwhals DataFrame | |----------------------------| |shape: (4, 2) | |┌──────────────┮───────────┐| |│ a ┆ a_unique │| |│ --- ┆ --- │| |│ list[i64] ┆ list[i64] │| |╞══════════════╩═══════════╡| |│ [1, 1, 2] ┆ [1, 2] │| |│ [3, 3, null] ┆ [null, 3] │| |│ null ┆ null │| |│ [] ┆ [] │| |└──────────────â”ī───────────┘| └────────────────────────────┘ """ return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "list.unique")) def contains(self, item: NonNestedLiteral) -> ExprT: """Check if sublists contain the given item. Arguments: item: Item that will be checked for membership. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"a": [[1, 2], None, []]}) >>> df = nw.from_native(df_native) >>> df.with_columns(a_contains_1=nw.col("a").list.contains(1)) ┌────────────────────────────┐ | Narwhals DataFrame | |----------------------------| |shape: (3, 2) | |┌───────────┮──────────────┐| |│ a ┆ a_contains_1 │| |│ --- ┆ --- │| |│ list[i64] ┆ bool │| |╞═══════════╩══════════════╡| |│ [1, 2] ┆ true │| |│ null ┆ null │| |│ [] ┆ false │| |└───────────â”ī──────────────┘| └────────────────────────────┘ """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "list.contains", item=item) ) def get(self, index: int) -> ExprT: """Return the value by index in each list. Negative indices are not accepted. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"a": [[1, 2], [3, 4, None], [None, 5]]}) >>> df = nw.from_native(df_native) >>> df.with_columns(a_first=nw.col("a").list.get(0)) ┌──────────────────────────┐ | Narwhals DataFrame | |--------------------------| |shape: (3, 2) | |┌──────────────┮─────────┐| |│ a ┆ a_first │| |│ --- ┆ --- │| |│ list[i64] ┆ i64 │| |╞══════════════╩═════════╡| |│ [1, 2] ┆ 1 │| |│ [3, 4, null] ┆ 3 │| |│ [null, 5] ┆ null │| |└──────────────â”ī─────────┘| └──────────────────────────┘ """ if not isinstance(index, int): msg = ( f"Index must be of type 'int'. Got type '{type(index).__name__}' instead." ) raise TypeError(msg) if index < 0: msg = f"Index {index} is out of bounds: should be greater than or equal to 0." raise ValueError(msg) return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "list.get", index=index) ) narwhals-2.11.0/narwhals/expr_name.py0000644000000000000000000001131613615410400014503 0ustar00from __future__ import annotations from typing import TYPE_CHECKING, Callable, Generic, TypeVar from narwhals._expression_parsing import ExprKind, ExprNode if TYPE_CHECKING: from narwhals.expr import Expr ExprT = TypeVar("ExprT", bound="Expr") class ExprNameNamespace(Generic[ExprT]): def __init__(self, expr: ExprT) -> None: self._expr = expr def keep(self) -> ExprT: r"""Keep the original root name of the expression. Notes: For Polars versions prior to 1.32, this will undo any previous renaming operations on the expression. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"foo": [1, 2], "BAR": [4, 5]}) >>> df = nw.from_native(df_native) >>> df.select(nw.col("foo").alias("alias_for_foo").name.keep()).columns ['foo'] """ return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "name.keep")) def map(self, function: Callable[[str], str]) -> ExprT: r"""Rename the output of an expression by mapping a function over the root name. Arguments: function: Function that maps a root name to a new name. Notes: For Polars versions prior to 1.32, this will undo any previous renaming operations on the expression. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"foo": [1, 2], "BAR": [4, 5]}) >>> df = nw.from_native(df_native) >>> renaming_func = lambda s: s[::-1] # reverse column name >>> df.select(nw.col("foo", "BAR").name.map(renaming_func)).columns ['oof', 'RAB'] """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "name.map", function=function) ) def prefix(self, prefix: str) -> ExprT: r"""Add a prefix to the root column name of the expression. Arguments: prefix: Prefix to add to the root column name. Notes: For Polars versions prior to 1.32, this will undo any previous renaming operations on the expression. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"foo": [1, 2], "BAR": [4, 5]}) >>> df = nw.from_native(df_native) >>> df.select(nw.col("foo", "BAR").name.prefix("with_prefix")).columns ['with_prefixfoo', 'with_prefixBAR'] """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "name.prefix", prefix=prefix) ) def suffix(self, suffix: str) -> ExprT: r"""Add a suffix to the root column name of the expression. Arguments: suffix: Suffix to add to the root column name. Notes: For Polars versions prior to 1.32, this will undo any previous renaming operations on the expression. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"foo": [1, 2], "BAR": [4, 5]}) >>> df = nw.from_native(df_native) >>> df.select(nw.col("foo", "BAR").name.suffix("_with_suffix")).columns ['foo_with_suffix', 'BAR_with_suffix'] """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "name.suffix", suffix=suffix) ) def to_lowercase(self) -> ExprT: r"""Make the root column name lowercase. Notes: For Polars versions prior to 1.32, this will undo any previous renaming operations on the expression. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"foo": [1, 2], "BAR": [4, 5]}) >>> df = nw.from_native(df_native) >>> df.select(nw.col("foo", "BAR").name.to_lowercase()).columns ['foo', 'bar'] """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "name.to_lowercase") ) def to_uppercase(self) -> ExprT: r"""Make the root column name uppercase. Notes: For Polars versions prior to 1.32, this will undo any previous renaming operations on the expression. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"foo": [1, 2], "BAR": [4, 5]}) >>> df = nw.from_native(df_native) >>> df.select(nw.col("foo", "BAR").name.to_uppercase()).columns ['FOO', 'BAR'] """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "name.to_uppercase") ) narwhals-2.11.0/narwhals/expr_str.py0000644000000000000000000005313213615410400014375 0ustar00from __future__ import annotations from typing import TYPE_CHECKING, Generic, TypeVar from narwhals._expression_parsing import ExprKind, ExprNode if TYPE_CHECKING: from narwhals.expr import Expr from narwhals.typing import IntoExpr ExprT = TypeVar("ExprT", bound="Expr") class ExprStringNamespace(Generic[ExprT]): def __init__(self, expr: ExprT) -> None: self._expr = expr def len_chars(self) -> ExprT: r"""Return the length of each string as the number of characters. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"words": ["foo", "345", None]}) >>> df = nw.from_native(df_native) >>> df.with_columns(words_len=nw.col("words").str.len_chars()) ┌─────────────────────┐ | Narwhals DataFrame | |---------------------| |shape: (3, 2) | |┌───────┮───────────┐| |│ words ┆ words_len │| |│ --- ┆ --- │| |│ str ┆ u32 │| |╞═══════╩═══════════╡| |│ foo ┆ 3 │| |│ 345 ┆ 3 │| |│ null ┆ null │| |└───────â”ī───────────┘| └─────────────────────┘ """ return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "str.len_chars")) def replace( self, pattern: str, value: str | IntoExpr, *, literal: bool = False, n: int = 1 ) -> ExprT: r"""Replace first matching regex/literal substring with a new string value. Arguments: pattern: A valid regular expression pattern. value: String that will replace the matched substring. literal: Treat `pattern` as a literal string. n: Number of matches to replace. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"foo": ["123abc", "abc abc123"]}) >>> df = nw.from_native(df_native) >>> df.with_columns(replaced=nw.col("foo").str.replace("abc", "")) ┌──────────────────────┐ | Narwhals DataFrame | |----------------------| | foo replaced| |0 123abc 123| |1 abc abc123 abc123| └──────────────────────┘ """ return self._expr._append_node( ExprNode( ExprKind.ELEMENTWISE, "str.replace", value, pattern=pattern, literal=literal, n=n, str_as_lit=True, ) ) def replace_all( self, pattern: str, value: IntoExpr, *, literal: bool = False ) -> ExprT: r"""Replace all matching regex/literal substring with a new string value. Arguments: pattern: A valid regular expression pattern. value: String that will replace the matched substring. literal: Treat `pattern` as a literal string. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"foo": ["123abc", "abc abc123"]}) >>> df = nw.from_native(df_native) >>> df.with_columns(replaced=nw.col("foo").str.replace_all("abc", "")) ┌──────────────────────┐ | Narwhals DataFrame | |----------------------| | foo replaced| |0 123abc 123| |1 abc abc123 123| └──────────────────────┘ """ return self._expr._append_node( ExprNode( ExprKind.ELEMENTWISE, "str.replace_all", value, pattern=pattern, literal=literal, str_as_lit=True, ) ) def strip_chars(self, characters: str | None = None) -> ExprT: r"""Remove leading and trailing characters. Arguments: characters: The set of characters to be removed. All combinations of this set of characters will be stripped from the start and end of the string. If set to None (default), all leading and trailing whitespace is removed instead. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"fruits": ["apple", "\nmango"]}) >>> df = nw.from_native(df_native) >>> df.with_columns(stripped=nw.col("fruits").str.strip_chars()).to_dict( ... as_series=False ... ) {'fruits': ['apple', '\nmango'], 'stripped': ['apple', 'mango']} """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "str.strip_chars", characters=characters) ) def starts_with(self, prefix: str) -> ExprT: r"""Check if string values start with a substring. Arguments: prefix: prefix substring Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"fruits": ["apple", "mango", None]}) >>> df = nw.from_native(df_native) >>> df.with_columns(has_prefix=nw.col("fruits").str.starts_with("app")) ┌───────────────────┐ |Narwhals DataFrame | |-------------------| | fruits has_prefix| |0 apple True| |1 mango False| |2 None None| └───────────────────┘ """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "str.starts_with", prefix=prefix) ) def ends_with(self, suffix: str) -> ExprT: r"""Check if string values end with a substring. Arguments: suffix: suffix substring Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"fruits": ["apple", "mango", None]}) >>> df = nw.from_native(df_native) >>> df.with_columns(has_suffix=nw.col("fruits").str.ends_with("ngo")) ┌───────────────────┐ |Narwhals DataFrame | |-------------------| | fruits has_suffix| |0 apple False| |1 mango True| |2 None None| └───────────────────┘ """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "str.ends_with", suffix=suffix) ) def contains(self, pattern: str, *, literal: bool = False) -> ExprT: r"""Check if string contains a substring that matches a pattern. Arguments: pattern: A Character sequence or valid regular expression pattern. literal: If True, treats the pattern as a literal string. If False, assumes the pattern is a regular expression. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"pets": ["cat", "dog", "rabbit and parrot"]}) >>> df = nw.from_native(df_native) >>> df.with_columns( ... default_match=nw.col("pets").str.contains("cat|parrot"), ... case_insensitive_match=nw.col("pets").str.contains("cat|(?i)parrot"), ... ).to_native() pyarrow.Table pets: string default_match: bool case_insensitive_match: bool ---- pets: [["cat","dog","rabbit and parrot"]] default_match: [[true,false,true]] case_insensitive_match: [[true,false,true]] """ return self._expr._append_node( ExprNode( ExprKind.ELEMENTWISE, "str.contains", pattern=pattern, literal=literal ) ) def slice(self, offset: int, length: int | None = None) -> ExprT: r"""Create subslices of the string values of an expression. Arguments: offset: Start index. Negative indexing is supported. length: Length of the slice. If set to `None` (default), the slice is taken to the end of the string. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"s": ["pear", None, "papaya"]}) >>> df = nw.from_native(df_native) >>> df.with_columns(s_sliced=nw.col("s").str.slice(4, length=3)) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | s s_sliced| |0 pear | |1 None None| |2 papaya ya| └──────────────────┘ """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "str.slice", offset=offset, length=length) ) def split(self, by: str) -> ExprT: r"""Split the string values of an expression by a substring. Arguments: by: Substring to split by. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"s": ["foo bar", "foo_bar"]}) >>> df = nw.from_native(df_native) >>> df.with_columns(nw.col("s").str.split("_").alias("s_split")) ┌────────────────────────────┐ | Narwhals DataFrame | |----------------------------| |shape: (2, 2) | |┌─────────┮────────────────┐| |│ s ┆ s_split │| |│ --- ┆ --- │| |│ str ┆ list[str] │| |╞═════════╩════════════════╡| |│ foo bar ┆ ["foo bar"] │| |│ foo_bar ┆ ["foo", "bar"] │| |└─────────â”ī────────────────┘| └────────────────────────────┘ """ return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "str.split", by=by)) def head(self, n: int = 5) -> ExprT: r"""Take the first n elements of each string. Arguments: n: Number of elements to take. Negative indexing is **not** supported. Notes: If the length of the string has fewer than `n` characters, the full string is returned. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"lyrics": ["taata", "taatatata", "zukkyun"]}) >>> df = nw.from_native(df_native) >>> df.with_columns(lyrics_head=nw.col("lyrics").str.head()).to_native() pyarrow.Table lyrics: string lyrics_head: string ---- lyrics: [["taata","taatatata","zukkyun"]] lyrics_head: [["taata","taata","zukky"]] """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "str.slice", offset=0, length=n) ) def tail(self, n: int = 5) -> ExprT: r"""Take the last n elements of each string. Arguments: n: Number of elements to take. Negative indexing is **not** supported. Notes: If the length of the string has fewer than `n` characters, the full string is returned. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"lyrics": ["taata", "taatatata", "zukkyun"]}) >>> df = nw.from_native(df_native) >>> df.with_columns(lyrics_tail=nw.col("lyrics").str.tail()).to_native() pyarrow.Table lyrics: string lyrics_tail: string ---- lyrics: [["taata","taatatata","zukkyun"]] lyrics_tail: [["taata","atata","kkyun"]] """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "str.slice", offset=-n, length=None) ) def to_datetime(self, format: str | None = None) -> ExprT: """Convert to Datetime dtype. Notes: - pandas defaults to nanosecond time unit, Polars to microsecond. Prior to pandas 2.0, nanoseconds were the only time unit supported in pandas, with no ability to set any other one. The ability to set the time unit in pandas, if the version permits, will arrive. - timezone-aware strings are all converted to and parsed as UTC. Warning: As different backends auto-infer format in different ways, if `format=None` there is no guarantee that the result will be equal. Arguments: format: Format to use for conversion. If set to None (default), the format is inferred from the data. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"a": ["2020-01-01", "2020-01-02"]}) >>> df = nw.from_native(df_native) >>> df.select(nw.col("a").str.to_datetime(format="%Y-%m-%d")) ┌───────────────────────┐ | Narwhals DataFrame | |-----------------------| |shape: (2, 1) | |┌─────────────────────┐| |│ a │| |│ --- │| |│ datetime[Ξs] │| |╞═════════════════════╡| |│ 2020-01-01 00:00:00 │| |│ 2020-01-02 00:00:00 │| |└─────────────────────┘| └───────────────────────┘ """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "str.to_datetime", format=format) ) def to_date(self, format: str | None = None) -> ExprT: """Convert to date dtype. Warning: As different backends auto-infer format in different ways, if `format=None` there is no guarantee that the result will be equal. Arguments: format: Format to use for conversion. If set to None (default), the format is inferred from the data. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> df_native = pa.table({"a": ["2020-01-01", "2020-01-02"]}) >>> df = nw.from_native(df_native) >>> df.select(nw.col("a").str.to_date(format="%Y-%m-%d")) ┌────────────────────────────┐ | Narwhals DataFrame | |----------------------------| |pyarrow.Table | |a: date32[day] | |---- | |a: [[2020-01-01,2020-01-02]]| └────────────────────────────┘ """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "str.to_date", format=format) ) def to_uppercase(self) -> ExprT: r"""Transform string to uppercase variant. Notes: The PyArrow backend will convert 'ß' to '᚞' instead of 'SS'. For more info see [the related issue](https://github.com/apache/arrow/issues/34599). There may be other unicode-edge-case-related variations across implementations. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"fruits": ["apple", None]}) >>> df = nw.from_native(df_native) >>> df.with_columns(upper_col=nw.col("fruits").str.to_uppercase()) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | fruits upper_col| |0 apple APPLE| |1 None None| └──────────────────┘ """ return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "str.to_uppercase")) def to_lowercase(self) -> ExprT: r"""Transform string to lowercase variant. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"fruits": ["APPLE", None]}) >>> df = nw.from_native(df_native) >>> df.with_columns(lower_col=nw.col("fruits").str.to_lowercase()) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | fruits lower_col| |0 APPLE apple| |1 None None| └──────────────────┘ """ return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "str.to_lowercase")) def to_titlecase(self) -> ExprT: """Modify strings to their titlecase equivalent. Notes: This is a form of case transform where the first letter of each word is capitalized, with the rest of the word in lowercase. Word boundaries are defined by non-**alphabetic** characters, matching the behavior of [`str.title`](https://docs.python.org/3/library/stdtypes.html#str.title) Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame( ... { ... "quotes": [ ... "'e.t. phone home'", ... "you talkin' to me?", ... "to infinity,and BEYOND!", ... ] ... } ... ) >>> df = nw.from_native(df_native) >>> df.with_columns(quotes_title=nw.col("quotes").str.to_titlecase()) ┌─────────────────────────────────────────────────────┐ | Narwhals DataFrame | |-----------------------------------------------------| |shape: (3, 2) | |┌─────────────────────────┮─────────────────────────┐| |│ quotes ┆ quotes_title │| |│ --- ┆ --- │| |│ str ┆ str │| |╞═════════════════════════╩═════════════════════════╡| |│ 'e.t. phone home' ┆ 'E.T. Phone Home' │| |│ you talkin' to me? ┆ You Talkin' To Me? │| |│ to infinity,and BEYOND! ┆ To Infinity,And Beyond! │| |└─────────────────────────â”ī─────────────────────────┘| └─────────────────────────────────────────────────────┘ """ return self._expr._append_node(ExprNode(ExprKind.ELEMENTWISE, "str.to_titlecase")) def zfill(self, width: int) -> ExprT: """Transform string to zero-padded variant. Arguments: width: The desired length of the string after padding. If the length of the string is greater than `width`, no padding is applied. If `width` is less than 0, no padding is applied. Examples: >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame({"digits": ["+1", "-1", "1", None]}) >>> df = nw.from_native(df_native) >>> df.with_columns(zfill_col=nw.col("digits").str.zfill(3)) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | digits zfill_col| |0 +1 +01| |1 -1 -01| |2 1 001| |3 None None| └──────────────────┘ """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "str.zfill", width=width) ) narwhals-2.11.0/narwhals/expr_struct.py0000644000000000000000000000337613615410400015116 0ustar00from __future__ import annotations from typing import TYPE_CHECKING, Generic, TypeVar from narwhals._expression_parsing import ExprKind, ExprNode if TYPE_CHECKING: from narwhals.expr import Expr ExprT = TypeVar("ExprT", bound="Expr") class ExprStructNamespace(Generic[ExprT]): def __init__(self, expr: ExprT) -> None: self._expr = expr def field(self, name: str) -> ExprT: r"""Retrieve a Struct field as a new expression. Arguments: name: Name of the struct field to retrieve. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame( ... {"user": [{"id": "0", "name": "john"}, {"id": "1", "name": "jane"}]} ... ) >>> df = nw.from_native(df_native) >>> df.with_columns(name=nw.col("user").struct.field("name")) ┌───────────────────────┐ | Narwhals DataFrame | |-----------------------| |shape: (2, 2) | |┌──────────────┮──────┐| |│ user ┆ name │| |│ --- ┆ --- │| |│ struct[2] ┆ str │| |╞══════════════╩══════╡| |│ {"0","john"} ┆ john │| |│ {"1","jane"} ┆ jane │| |└──────────────â”ī──────┘| └───────────────────────┘ """ return self._expr._append_node( ExprNode(ExprKind.ELEMENTWISE, "struct.field", name=name) ) narwhals-2.11.0/narwhals/functions.py0000644000000000000000000017760413615410400014552 0ustar00from __future__ import annotations import platform import sys from collections.abc import Iterable, Mapping, Sequence from typing import TYPE_CHECKING, Any from narwhals._expression_parsing import ExprKind, ExprNode, is_expr, is_series from narwhals._utils import ( Implementation, Version, deprecate_native_namespace, flatten, is_eager_allowed, is_sequence_but_not_str, normalize_path, supports_arrow_c_stream, validate_laziness, ) from narwhals.dependencies import ( is_narwhals_series, is_numpy_array, is_numpy_array_2d, is_pyarrow_table, ) from narwhals.exceptions import InvalidOperationError from narwhals.expr import Expr from narwhals.translate import from_native, to_native if TYPE_CHECKING: from types import ModuleType from typing_extensions import TypeAlias, TypeIs from narwhals._native import NativeDataFrame, NativeLazyFrame, NativeSeries from narwhals._translate import IntoArrowTable from narwhals._typing import Backend, EagerAllowed, IntoBackend from narwhals.dataframe import DataFrame, LazyFrame from narwhals.dtypes import DType from narwhals.series import Series from narwhals.typing import ( ConcatMethod, FileSource, FrameT, IntoDType, IntoExpr, IntoSchema, NonNestedLiteral, _2DArray, ) _IntoSchema: TypeAlias = "IntoSchema | Sequence[str] | None" def concat(items: Iterable[FrameT], *, how: ConcatMethod = "vertical") -> FrameT: """Concatenate multiple DataFrames, LazyFrames into a single entity. Arguments: items: DataFrames, LazyFrames to concatenate. how: concatenating strategy - vertical: Concatenate vertically. Column names must match. - horizontal: Concatenate horizontally. If lengths don't match, then missing rows are filled with null values. This is only supported when all inputs are (eager) DataFrames. - diagonal: Finds a union between the column schemas and fills missing column values with null. Raises: TypeError: The items to concatenate should either all be eager, or all lazy Examples: Let's take an example of vertical concatenation: >>> import pandas as pd >>> import polars as pl >>> import pyarrow as pa >>> import narwhals as nw Let's look at one case a for vertical concatenation (pandas backed): >>> df_pd_1 = nw.from_native(pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})) >>> df_pd_2 = nw.from_native(pd.DataFrame({"a": [5, 2], "b": [1, 4]})) >>> nw.concat([df_pd_1, df_pd_2], how="vertical") ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 1 4 | | 1 2 5 | | 2 3 6 | | 0 5 1 | | 1 2 4 | └──────────────────┘ Let's look at one case a for horizontal concatenation (polars backed): >>> df_pl_1 = nw.from_native(pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})) >>> df_pl_2 = nw.from_native(pl.DataFrame({"c": [5, 2], "d": [1, 4]})) >>> nw.concat([df_pl_1, df_pl_2], how="horizontal") ┌───────────────────────────┐ | Narwhals DataFrame | |---------------------------| |shape: (3, 4) | |┌─────┮─────┮──────┮──────┐| |│ a ┆ b ┆ c ┆ d │| |│ --- ┆ --- ┆ --- ┆ --- │| |│ i64 ┆ i64 ┆ i64 ┆ i64 │| |╞═════╩═════╩══════╩══════╡| |│ 1 ┆ 4 ┆ 5 ┆ 1 │| |│ 2 ┆ 5 ┆ 2 ┆ 4 │| |│ 3 ┆ 6 ┆ null ┆ null │| |└─────â”ī─────â”ī──────â”ī──────┘| └───────────────────────────┘ Let's look at one case a for diagonal concatenation (pyarrow backed): >>> df_pa_1 = nw.from_native(pa.table({"a": [1, 2], "b": [3.5, 4.5]})) >>> df_pa_2 = nw.from_native(pa.table({"a": [3, 4], "z": ["x", "y"]})) >>> nw.concat([df_pa_1, df_pa_2], how="diagonal") ┌──────────────────────────┐ | Narwhals DataFrame | |--------------------------| |pyarrow.Table | |a: int64 | |b: double | |z: string | |---- | |a: [[1,2],[3,4]] | |b: [[3.5,4.5],[null,null]]| |z: [[null,null],["x","y"]]| └──────────────────────────┘ """ from narwhals.dependencies import is_narwhals_lazyframe if not items: msg = "No items to concatenate." raise ValueError(msg) items = tuple(items) validate_laziness(items) if how not in {"horizontal", "vertical", "diagonal"}: # pragma: no cover msg = "Only vertical, horizontal and diagonal concatenations are supported." raise NotImplementedError(msg) first_item = items[0] if is_narwhals_lazyframe(first_item) and how == "horizontal": msg = ( "Horizontal concatenation is not supported for LazyFrames.\n\n" "Hint: you may want to use `join` instead." ) raise InvalidOperationError(msg) plx = first_item.__narwhals_namespace__() return first_item._with_compliant( plx.concat([df._compliant_frame for df in items], how=how) ) def new_series( name: str, values: Any, dtype: IntoDType | None = None, *, backend: IntoBackend[EagerAllowed], ) -> Series[Any]: """Instantiate Narwhals Series from iterable (e.g. list or array). Arguments: name: Name of resulting Series. values: Values of make Series from. dtype: (Narwhals) dtype. If not provided, the native library may auto-infer it from `values`. backend: specifies which eager backend instantiate to. `backend` can be specified in various ways - As `Implementation.` with `BACKEND` being `PANDAS`, `PYARROW`, `POLARS`, `MODIN` or `CUDF`. - As a string: `"pandas"`, `"pyarrow"`, `"polars"`, `"modin"` or `"cudf"`. - Directly as a module `pandas`, `pyarrow`, `polars`, `modin` or `cudf`. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> values = [4, 1, 2, 3] >>> nw.new_series(name="a", values=values, dtype=nw.Int32, backend=pd) ┌─────────────────────┐ | Narwhals Series | |---------------------| |0 4 | |1 1 | |2 2 | |3 3 | |Name: a, dtype: int32| └─────────────────────┘ """ return _new_series_impl(name, values, dtype, backend=backend) def _new_series_impl( name: str, values: Any, dtype: IntoDType | None = None, *, backend: IntoBackend[EagerAllowed], ) -> Series[Any]: implementation = Implementation.from_backend(backend) if is_eager_allowed(implementation): ns = Version.MAIN.namespace.from_backend(implementation).compliant series = ns._series.from_iterable(values, name=name, context=ns, dtype=dtype) return series.to_narwhals() if implementation is Implementation.UNKNOWN: # pragma: no cover _native_namespace = implementation.to_native_namespace() try: native_series: NativeSeries = _native_namespace.new_series( name, values, dtype ) return from_native(native_series, series_only=True).alias(name) except AttributeError as e: msg = "Unknown namespace is expected to implement `new_series` constructor." raise AttributeError(msg) from e msg = ( f"{implementation} support in Narwhals is lazy-only, but `new_series` is an eager-only function.\n\n" "Hint: you may want to use an eager backend and then call `.lazy`, e.g.:\n\n" f" nw.new_series('a', [1,2,3], backend='pyarrow').to_frame().lazy('{implementation}')" ) raise ValueError(msg) @deprecate_native_namespace(warn_version="1.26.0") def from_dict( data: Mapping[str, Any], schema: IntoSchema | Mapping[str, DType | None] | None = None, *, backend: IntoBackend[EagerAllowed] | None = None, native_namespace: ModuleType | None = None, # noqa: ARG001 ) -> DataFrame[Any]: """Instantiate DataFrame from dictionary. Indexes (if present, for pandas-like backends) are aligned following the [left-hand-rule](../concepts/pandas_index.md/). Notes: For pandas-like dataframes, conversion to schema is applied after dataframe creation. Arguments: data: Dictionary to create DataFrame from. schema: The DataFrame schema as Schema or dict of {name: type}. If not specified, the schema will be inferred by the native library. If any `dtype` is `None`, the data type for that column will be inferred by the native library. backend: specifies which eager backend instantiate to. Only necessary if inputs are not Narwhals Series. `backend` can be specified in various ways - As `Implementation.` with `BACKEND` being `PANDAS`, `PYARROW`, `POLARS`, `MODIN` or `CUDF`. - As a string: `"pandas"`, `"pyarrow"`, `"polars"`, `"modin"` or `"cudf"`. - Directly as a module `pandas`, `pyarrow`, `polars`, `modin` or `cudf`. native_namespace: deprecated, same as `backend`. Examples: >>> import pandas as pd >>> import narwhals as nw >>> data = {"c": [5, 2], "d": [1, 4]} >>> nw.from_dict(data, backend="pandas") ┌──────────────────┐ |Narwhals DataFrame| |------------------| | c d | | 0 5 1 | | 1 2 4 | └──────────────────┘ """ if backend is None: data, backend = _from_dict_no_backend(data) if schema and data and (diff := set(schema.keys()).symmetric_difference(data.keys())): msg = f"Keys in `schema` and `data` are expected to match, found unmatched keys: {diff}" raise InvalidOperationError(msg) implementation = Implementation.from_backend(backend) if is_eager_allowed(implementation): ns = Version.MAIN.namespace.from_backend(implementation).compliant return ns._dataframe.from_dict(data, schema=schema, context=ns).to_narwhals() if implementation is Implementation.UNKNOWN: # pragma: no cover _native_namespace = implementation.to_native_namespace() try: # implementation is UNKNOWN, Narwhals extension using this feature should # implement `from_dict` function in the top-level namespace. native_frame: NativeDataFrame = _native_namespace.from_dict( data, schema=schema ) except AttributeError as e: msg = "Unknown namespace is expected to implement `from_dict` function." raise AttributeError(msg) from e return from_native(native_frame, eager_only=True) msg = ( f"{implementation} support in Narwhals is lazy-only, but `from_dict` is an eager-only function.\n\n" "Hint: you may want to use an eager backend and then call `.lazy`, e.g.:\n\n" f" nw.from_dict({{'a': [1, 2]}}, backend='pyarrow').lazy('{implementation}')" ) raise ValueError(msg) def _from_dict_no_backend( data: Mapping[str, Series[Any] | Any], / ) -> tuple[dict[str, Series[Any] | Any], ModuleType]: for val in data.values(): if is_narwhals_series(val): native_namespace = val.__native_namespace__() break else: msg = "Calling `from_dict` without `backend` is only supported if all input values are already Narwhals Series" raise TypeError(msg) data = {key: to_native(value, pass_through=True) for key, value in data.items()} return data, native_namespace def from_dicts( data: Sequence[Mapping[str, Any]], schema: IntoSchema | Mapping[str, DType | None] | None = None, *, backend: IntoBackend[EagerAllowed], ) -> DataFrame[Any]: """Instantiate DataFrame from a sequence of dictionaries representing rows. Notes: For pandas-like dataframes, conversion to schema is applied after dataframe creation. Arguments: data: Sequence with dictionaries mapping column name to value. schema: The DataFrame schema as Schema or dict of {name: type}. If not specified, the schema will be inferred by the native library. If any `dtype` is `None`, the data type for that column will be inferred by the native library. backend: Specifies which eager backend instantiate to. `backend` can be specified in various ways - As `Implementation.` with `BACKEND` being `PANDAS`, `PYARROW`, `POLARS`, `MODIN` or `CUDF`. - As a string: `"pandas"`, `"pyarrow"`, `"polars"`, `"modin"` or `"cudf"`. - Directly as a module `pandas`, `pyarrow`, `polars`, `modin` or `cudf`. Tip: If you expect non-uniform keys in `data`, consider passing `schema` for more consistent results, as **inference varies between backends**: - pandas uses all rows - polars uses the first 100 rows - pyarrow uses only the first row Examples: >>> import polars as pl >>> import narwhals as nw >>> data = [ ... {"item": "apple", "weight": 80, "price": 0.60}, ... {"item": "egg", "weight": 55, "price": 0.40}, ... ] >>> nw.DataFrame.from_dicts(data, backend="polars") ┌──────────────────────────┐ | Narwhals DataFrame | |--------------------------| |shape: (2, 3) | |┌───────┮────────┮───────┐| |│ item ┆ weight ┆ price │| |│ --- ┆ --- ┆ --- │| |│ str ┆ i64 ┆ f64 │| |╞═══════╩════════╩═══════╡| |│ apple ┆ 80 ┆ 0.6 │| |│ egg ┆ 55 ┆ 0.4 │| |└───────â”ī────────â”ī───────┘| └──────────────────────────┘ """ return Version.MAIN.dataframe.from_dicts(data, schema, backend=backend) def from_numpy( data: _2DArray, schema: IntoSchema | Sequence[str] | None = None, *, backend: IntoBackend[EagerAllowed], ) -> DataFrame[Any]: """Construct a DataFrame from a NumPy ndarray. Notes: Only row orientation is currently supported. For pandas-like dataframes, conversion to schema is applied after dataframe creation. Arguments: data: Two-dimensional data represented as a NumPy ndarray. schema: The DataFrame schema as Schema, dict of {name: type}, or a sequence of str. backend: specifies which eager backend instantiate to. `backend` can be specified in various ways - As `Implementation.` with `BACKEND` being `PANDAS`, `PYARROW`, `POLARS`, `MODIN` or `CUDF`. - As a string: `"pandas"`, `"pyarrow"`, `"polars"`, `"modin"` or `"cudf"`. - Directly as a module `pandas`, `pyarrow`, `polars`, `modin` or `cudf`. Examples: >>> import numpy as np >>> import pyarrow as pa >>> import narwhals as nw >>> >>> arr = np.array([[5, 2, 1], [1, 4, 3]]) >>> schema = {"c": nw.Int16(), "d": nw.Float32(), "e": nw.Int8()} >>> nw.from_numpy(arr, schema=schema, backend="pyarrow") ┌──────────────────┐ |Narwhals DataFrame| |------------------| | pyarrow.Table | | c: int16 | | d: float | | e: int8 | | ---- | | c: [[5,1]] | | d: [[2,4]] | | e: [[1,3]] | └──────────────────┘ """ if not is_numpy_array_2d(data): msg = "`from_numpy` only accepts 2D numpy arrays" raise ValueError(msg) if not _is_into_schema(schema): msg = ( "`schema` is expected to be one of the following types: " "IntoSchema | Sequence[str]. " f"Got {type(schema)}." ) raise TypeError(msg) implementation = Implementation.from_backend(backend) if is_eager_allowed(implementation): ns = Version.MAIN.namespace.from_backend(implementation).compliant return ns.from_numpy(data, schema).to_narwhals() if implementation is Implementation.UNKNOWN: # pragma: no cover _native_namespace = implementation.to_native_namespace() try: # implementation is UNKNOWN, Narwhals extension using this feature should # implement `from_numpy` function in the top-level namespace. native_frame: NativeDataFrame = _native_namespace.from_numpy( data, schema=schema ) except AttributeError as e: msg = "Unknown namespace is expected to implement `from_numpy` function." raise AttributeError(msg) from e return from_native(native_frame, eager_only=True) msg = ( f"{implementation} support in Narwhals is lazy-only, but `from_numpy` is an eager-only function.\n\n" "Hint: you may want to use an eager backend and then call `.lazy`, e.g.:\n\n" f" nw.from_numpy(arr, backend='pyarrow').lazy('{implementation}')" ) raise ValueError(msg) def _is_into_schema(obj: Any) -> TypeIs[_IntoSchema]: from narwhals.schema import Schema return ( obj is None or isinstance(obj, (Mapping, Schema)) or is_sequence_but_not_str(obj) ) def from_arrow( native_frame: IntoArrowTable, *, backend: IntoBackend[EagerAllowed] ) -> DataFrame[Any]: # pragma: no cover """Construct a DataFrame from an object which supports the PyCapsule Interface. Arguments: native_frame: Object which implements `__arrow_c_stream__`. backend: specifies which eager backend instantiate to. `backend` can be specified in various ways - As `Implementation.` with `BACKEND` being `PANDAS`, `PYARROW`, `POLARS`, `MODIN` or `CUDF`. - As a string: `"pandas"`, `"pyarrow"`, `"polars"`, `"modin"` or `"cudf"`. - Directly as a module `pandas`, `pyarrow`, `polars`, `modin` or `cudf`. Examples: >>> import pandas as pd >>> import polars as pl >>> import narwhals as nw >>> >>> df_native = pd.DataFrame({"a": [1, 2], "b": [4.2, 5.1]}) >>> nw.from_arrow(df_native, backend="polars") ┌──────────────────┐ |Narwhals DataFrame| |------------------| | shape: (2, 2) | | ┌─────┮─────┐ | | │ a ┆ b │ | | │ --- ┆ --- │ | | │ i64 ┆ f64 │ | | ╞═════╩═════╡ | | │ 1 ┆ 4.2 │ | | │ 2 ┆ 5.1 │ | | └─────â”ī─────┘ | └──────────────────┘ """ if not (supports_arrow_c_stream(native_frame) or is_pyarrow_table(native_frame)): msg = f"Given object of type {type(native_frame)} does not support PyCapsule interface" raise TypeError(msg) implementation = Implementation.from_backend(backend) if is_eager_allowed(implementation): ns = Version.MAIN.namespace.from_backend(implementation).compliant return ns._dataframe.from_arrow(native_frame, context=ns).to_narwhals() if implementation is Implementation.UNKNOWN: # pragma: no cover _native_namespace = implementation.to_native_namespace() try: # implementation is UNKNOWN, Narwhals extension using this feature should # implement PyCapsule support native: NativeDataFrame = _native_namespace.DataFrame(native_frame) except AttributeError as e: msg = "Unknown namespace is expected to implement `DataFrame` class which accepts object which supports PyCapsule Interface." raise AttributeError(msg) from e return from_native(native, eager_only=True) msg = ( f"{implementation} support in Narwhals is lazy-only, but `from_arrow` is an eager-only function.\n\n" "Hint: you may want to use an eager backend and then call `.lazy`, e.g.:\n\n" f" nw.from_arrow(df, backend='pyarrow').lazy('{implementation}')" ) raise ValueError(msg) def _get_sys_info() -> dict[str, str]: """System information. Returns system and Python version information Copied from sklearn Returns: Dictionary with system info. """ python = sys.version.replace("\n", " ") blob = ( ("python", python), ("executable", sys.executable), ("machine", platform.platform()), ) return dict(blob) def _get_deps_info() -> dict[str, str]: """Overview of the installed version of main dependencies. This function does not import the modules to collect the version numbers but instead relies on standard Python package metadata. Returns version information on relevant Python libraries This function and show_versions were copied from sklearn and adapted Returns: Mapping from dependency to version. """ from importlib.metadata import distributions extra_names = ("narwhals", "numpy") member_names = Implementation._member_names_ exclude = {"PYSPARK_CONNECT", "UNKNOWN"} target_names = tuple( name.lower() for name in (*extra_names, *member_names) if name not in exclude ) result = dict.fromkeys(target_names, "") # Initialize with empty strings for dist in distributions(): dist_name, dist_version = dist.name.lower(), dist.version if dist_name in result: # exact match result[dist_name] = dist_version else: # prefix match for target in target_names: if not result[target] and dist_name.startswith(target): result[target] = dist_version break return result def show_versions() -> None: """Print useful debugging information. Examples: >>> from narwhals import show_versions >>> show_versions() # doctest: +SKIP """ sys_info = _get_sys_info() deps_info = _get_deps_info() print("\nSystem:") # noqa: T201 for k, stat in sys_info.items(): print(f"{k:>10}: {stat}") # noqa: T201 print("\nPython dependencies:") # noqa: T201 for k, stat in deps_info.items(): print(f"{k:>13}: {stat}") # noqa: T201 def read_csv( source: FileSource, *, backend: IntoBackend[EagerAllowed], **kwargs: Any ) -> DataFrame[Any]: """Read a CSV file into a DataFrame. Arguments: source: Path to a file. backend: The eager backend for DataFrame creation. `backend` can be specified in various ways - As `Implementation.` with `BACKEND` being `PANDAS`, `PYARROW`, `POLARS`, `MODIN` or `CUDF`. - As a string: `"pandas"`, `"pyarrow"`, `"polars"`, `"modin"` or `"cudf"`. - Directly as a module `pandas`, `pyarrow`, `polars`, `modin` or `cudf`. kwargs: Extra keyword arguments which are passed to the native CSV reader. For example, you could use `nw.read_csv('file.csv', backend='pandas', engine='pyarrow')`. Examples: >>> import narwhals as nw >>> nw.read_csv("file.csv", backend="pandas") # doctest:+SKIP ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 1 4 | | 1 2 5 | └──────────────────┘ """ impl = Implementation.from_backend(backend) native_namespace = impl.to_native_namespace() native_frame: NativeDataFrame if impl in { Implementation.POLARS, Implementation.PANDAS, Implementation.MODIN, Implementation.CUDF, }: native_frame = native_namespace.read_csv(normalize_path(source), **kwargs) elif impl is Implementation.PYARROW: from pyarrow import csv # ignore-banned-import native_frame = csv.read_csv(source, **kwargs) elif impl in { Implementation.PYSPARK, Implementation.DASK, Implementation.DUCKDB, Implementation.IBIS, Implementation.SQLFRAME, Implementation.PYSPARK_CONNECT, }: msg = ( f"Expected eager backend, found {impl}.\n\n" f"Hint: use nw.scan_csv(source={source}, backend={backend})" ) raise ValueError(msg) else: # pragma: no cover try: # implementation is UNKNOWN, Narwhals extension using this feature should # implement `read_csv` function in the top-level namespace. native_frame = native_namespace.read_csv(source=source, **kwargs) except AttributeError as e: msg = "Unknown namespace is expected to implement `read_csv` function." raise AttributeError(msg) from e return from_native(native_frame, eager_only=True) def scan_csv( source: FileSource, *, backend: IntoBackend[Backend], **kwargs: Any ) -> LazyFrame[Any]: """Lazily read from a CSV file. For the libraries that do not support lazy dataframes, the function reads a csv file eagerly and then converts the resulting dataframe to a lazyframe. Arguments: source: Path to a file. backend: The eager backend for DataFrame creation. `backend` can be specified in various ways - As `Implementation.` with `BACKEND` being `PANDAS`, `PYARROW`, `POLARS`, `MODIN` or `CUDF`. - As a string: `"pandas"`, `"pyarrow"`, `"polars"`, `"modin"` or `"cudf"`. - Directly as a module `pandas`, `pyarrow`, `polars`, `modin` or `cudf`. kwargs: Extra keyword arguments which are passed to the native CSV reader. For example, you could use `nw.scan_csv('file.csv', backend=pd, engine='pyarrow')`. Examples: >>> import duckdb >>> import narwhals as nw >>> >>> nw.scan_csv("file.csv", backend="duckdb").to_native() # doctest:+SKIP ┌─────────┮───────┐ │ a │ b │ │ varchar │ int32 │ ├─────────┾───────â”Ī │ x │ 1 │ │ y │ 2 │ │ z │ 3 │ └─────────â”ī───────┘ """ implementation = Implementation.from_backend(backend) native_namespace = implementation.to_native_namespace() native_frame: NativeDataFrame | NativeLazyFrame source = normalize_path(source) if implementation is Implementation.POLARS: native_frame = native_namespace.scan_csv(source, **kwargs) elif implementation in { Implementation.PANDAS, Implementation.MODIN, Implementation.CUDF, Implementation.DASK, Implementation.DUCKDB, Implementation.IBIS, }: native_frame = native_namespace.read_csv(source, **kwargs) elif implementation is Implementation.PYARROW: from pyarrow import csv # ignore-banned-import native_frame = csv.read_csv(source, **kwargs) elif implementation.is_spark_like(): if (session := kwargs.pop("session", None)) is None: msg = "Spark like backends require a session object to be passed in `kwargs`." raise ValueError(msg) csv_reader = session.read.format("csv") native_frame = ( csv_reader.load(source) if ( implementation is Implementation.SQLFRAME and implementation._backend_version() < (3, 27, 0) ) else csv_reader.options(**kwargs).load(source) ) else: # pragma: no cover try: # implementation is UNKNOWN, Narwhals extension using this feature should # implement `scan_csv` function in the top-level namespace. native_frame = native_namespace.scan_csv(source=source, **kwargs) except AttributeError as e: msg = "Unknown namespace is expected to implement `scan_csv` function." raise AttributeError(msg) from e return from_native(native_frame).lazy() def read_parquet( source: FileSource, *, backend: IntoBackend[EagerAllowed], **kwargs: Any ) -> DataFrame[Any]: """Read into a DataFrame from a parquet file. Arguments: source: Path to a file. backend: The eager backend for DataFrame creation. `backend` can be specified in various ways - As `Implementation.` with `BACKEND` being `PANDAS`, `PYARROW`, `POLARS`, `MODIN` or `CUDF`. - As a string: `"pandas"`, `"pyarrow"`, `"polars"`, `"modin"` or `"cudf"`. - Directly as a module `pandas`, `pyarrow`, `polars`, `modin` or `cudf`. kwargs: Extra keyword arguments which are passed to the native parquet reader. For example, you could use `nw.read_parquet('file.parquet', backend=pd, engine='pyarrow')`. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> nw.read_parquet("file.parquet", backend="pyarrow") # doctest:+SKIP ┌──────────────────┐ |Narwhals DataFrame| |------------------| |pyarrow.Table | |a: int64 | |c: double | |---- | |a: [[1,2]] | |c: [[0.2,0.1]] | └──────────────────┘ """ impl = Implementation.from_backend(backend) native_namespace = impl.to_native_namespace() native_frame: NativeDataFrame if impl in { Implementation.POLARS, Implementation.PANDAS, Implementation.MODIN, Implementation.CUDF, }: source = normalize_path(source) native_frame = native_namespace.read_parquet(source, **kwargs) elif impl is Implementation.PYARROW: import pyarrow.parquet as pq # ignore-banned-import native_frame = pq.read_table(source, **kwargs) # type: ignore[arg-type] elif impl in { Implementation.PYSPARK, Implementation.DASK, Implementation.DUCKDB, Implementation.IBIS, Implementation.SQLFRAME, Implementation.PYSPARK_CONNECT, }: msg = ( f"Expected eager backend, found {impl}.\n\n" f"Hint: use nw.scan_parquet(source={source}, backend={backend})" ) raise ValueError(msg) else: # pragma: no cover try: # implementation is UNKNOWN, Narwhals extension using this feature should # implement `read_parquet` function in the top-level namespace. native_frame = native_namespace.read_parquet(source=source, **kwargs) except AttributeError as e: msg = "Unknown namespace is expected to implement `read_parquet` function." raise AttributeError(msg) from e return from_native(native_frame, eager_only=True) def scan_parquet( source: FileSource, *, backend: IntoBackend[Backend], **kwargs: Any ) -> LazyFrame[Any]: """Lazily read from a parquet file. For the libraries that do not support lazy dataframes, the function reads a parquet file eagerly and then converts the resulting dataframe to a lazyframe. Note: Spark like backends require a session object to be passed in `kwargs`. For instance: ```py import narwhals as nw from sqlframe.duckdb import DuckDBSession nw.scan_parquet(source, backend="sqlframe", session=DuckDBSession()) ``` Arguments: source: Path to a file. backend: The eager backend for DataFrame creation. `backend` can be specified in various ways - As `Implementation.` with `BACKEND` being `PANDAS`, `PYARROW`, `POLARS`, `MODIN`, `CUDF`, `PYSPARK` or `SQLFRAME`. - As a string: `"pandas"`, `"pyarrow"`, `"polars"`, `"modin"`, `"cudf"`, `"pyspark"` or `"sqlframe"`. - Directly as a module `pandas`, `pyarrow`, `polars`, `modin`, `cudf`, `pyspark.sql` or `sqlframe`. kwargs: Extra keyword arguments which are passed to the native parquet reader. For example, you could use `nw.scan_parquet('file.parquet', backend=pd, engine='pyarrow')`. Examples: >>> import dask.dataframe as dd >>> from sqlframe.duckdb import DuckDBSession >>> import narwhals as nw >>> >>> nw.scan_parquet("file.parquet", backend="dask").collect() # doctest:+SKIP ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 1 4 | | 1 2 5 | └──────────────────┘ >>> nw.scan_parquet( ... "file.parquet", backend="sqlframe", session=DuckDBSession() ... ).collect() # doctest:+SKIP ┌──────────────────┐ |Narwhals DataFrame| |------------------| | pyarrow.Table | | a: int64 | | b: int64 | | ---- | | a: [[1,2]] | | b: [[4,5]] | └──────────────────┘ """ implementation = Implementation.from_backend(backend) native_namespace = implementation.to_native_namespace() native_frame: NativeDataFrame | NativeLazyFrame source = normalize_path(source) if implementation is Implementation.POLARS: native_frame = native_namespace.scan_parquet(source, **kwargs) elif implementation in { Implementation.PANDAS, Implementation.MODIN, Implementation.CUDF, Implementation.DASK, Implementation.DUCKDB, Implementation.IBIS, }: native_frame = native_namespace.read_parquet(source, **kwargs) elif implementation is Implementation.PYARROW: import pyarrow.parquet as pq # ignore-banned-import native_frame = pq.read_table(source, **kwargs) elif implementation.is_spark_like(): if (session := kwargs.pop("session", None)) is None: msg = "Spark like backends require a session object to be passed in `kwargs`." raise ValueError(msg) pq_reader = session.read.format("parquet") native_frame = ( pq_reader.load(source) if ( implementation is Implementation.SQLFRAME and implementation._backend_version() < (3, 27, 0) ) else pq_reader.options(**kwargs).load(source) ) else: # pragma: no cover try: # implementation is UNKNOWN, Narwhals extension using this feature should # implement `scan_parquet` function in the top-level namespace. native_frame = native_namespace.scan_parquet(source=source, **kwargs) except AttributeError as e: msg = "Unknown namespace is expected to implement `scan_parquet` function." raise AttributeError(msg) from e return from_native(native_frame).lazy() def col(*names: str | Iterable[str]) -> Expr: """Creates an expression that references one or more columns by their name(s). Arguments: names: Name(s) of the columns to use. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> df_native = pl.DataFrame({"a": [1, 2], "b": [3, 4], "c": ["x", "z"]}) >>> nw.from_native(df_native).select(nw.col("a", "b") * nw.col("b")) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | shape: (2, 2) | | ┌─────┮─────┐ | | │ a ┆ b │ | | │ --- ┆ --- │ | | │ i64 ┆ i64 │ | | ╞═════╩═════╡ | | │ 3 ┆ 9 │ | | │ 8 ┆ 16 │ | | └─────â”ī─────┘ | └──────────────────┘ """ flat_names = flatten(names) return Expr(ExprNode(ExprKind.COL, "col", names=flat_names)) def exclude(*names: str | Iterable[str]) -> Expr: """Creates an expression that excludes columns by their name(s). Arguments: names: Name(s) of the columns to exclude. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> df_native = pl.DataFrame({"a": [1, 2], "b": [3, 4], "c": ["x", "z"]}) >>> nw.from_native(df_native).select(nw.exclude("c", "a")) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | shape: (2, 1) | | ┌─────┐ | | │ b │ | | │ --- │ | | │ i64 │ | | ╞═════╡ | | │ 3 │ | | │ 4 │ | | └─────┘ | └──────────────────┘ """ return Expr(ExprNode(ExprKind.EXCLUDE, "exclude", names=frozenset(flatten(names)))) def nth(*indices: int | Sequence[int]) -> Expr: """Creates an expression that references one or more columns by their index(es). Notes: `nth` is not supported for Polars version<1.0.0. Please use [`narwhals.col`][] instead. Arguments: indices: One or more indices representing the columns to retrieve. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> df_native = pa.table({"a": [1, 2], "b": [3, 4], "c": [0.123, 3.14]}) >>> nw.from_native(df_native).select(nw.nth(0, 2) * 2) ┌──────────────────┐ |Narwhals DataFrame| |------------------| |pyarrow.Table | |a: int64 | |c: double | |---- | |a: [[2,4]] | |c: [[0.246,6.28]] | └──────────────────┘ """ flat_indices = flatten(indices) return Expr(ExprNode(ExprKind.NTH, "nth", indices=flat_indices)) # Add underscore so it doesn't conflict with builtin `all` def all_() -> Expr: """Instantiate an expression representing all columns. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> df_native = pd.DataFrame({"a": [1, 2], "b": [3.14, 0.123]}) >>> nw.from_native(df_native).select(nw.all() * 2) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 2 6.280 | | 1 4 0.246 | └──────────────────┘ """ return Expr(ExprNode(ExprKind.ALL, "all")) # Add underscore so it doesn't conflict with builtin `len` def len_() -> Expr: """Return the number of rows. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> df_native = pl.DataFrame({"a": [1, 2], "b": [5, None]}) >>> nw.from_native(df_native).select(nw.len()) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | shape: (1, 1) | | ┌─────┐ | | │ len │ | | │ --- │ | | │ u32 │ | | ╞═════╡ | | │ 2 │ | | └─────┘ | └──────────────────┘ """ return Expr(ExprNode(ExprKind.AGGREGATION, "len")) def sum(*columns: str) -> Expr: """Sum all values. Note: Syntactic sugar for ``nw.col(columns).sum()`` Arguments: columns: Name(s) of the columns to use in the aggregation function Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> df_native = pd.DataFrame({"a": [1, 2], "b": [-1.4, 6.2]}) >>> nw.from_native(df_native).select(nw.sum("a", "b")) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 3 4.8 | └──────────────────┘ """ return col(*columns).sum() def mean(*columns: str) -> Expr: """Get the mean value. Note: Syntactic sugar for ``nw.col(columns).mean()`` Arguments: columns: Name(s) of the columns to use in the aggregation function Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> df_native = pa.table({"a": [1, 8, 3], "b": [3.14, 6.28, 42.1]}) >>> nw.from_native(df_native).select(nw.mean("a", "b")) ┌─────────────────────────┐ | Narwhals DataFrame | |-------------------------| |pyarrow.Table | |a: double | |b: double | |---- | |a: [[4]] | |b: [[17.173333333333336]]| └─────────────────────────┘ """ return col(*columns).mean() def median(*columns: str) -> Expr: """Get the median value. Notes: - Syntactic sugar for ``nw.col(columns).median()`` - Results might slightly differ across backends due to differences in the underlying algorithms used to compute the median. Arguments: columns: Name(s) of the columns to use in the aggregation function Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> df_native = pl.DataFrame({"a": [4, 5, 2]}) >>> nw.from_native(df_native).select(nw.median("a")) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | shape: (1, 1) | | ┌─────┐ | | │ a │ | | │ --- │ | | │ f64 │ | | ╞═════╡ | | │ 4.0 │ | | └─────┘ | └──────────────────┘ """ return col(*columns).median() def min(*columns: str) -> Expr: """Return the minimum value. Note: Syntactic sugar for ``nw.col(columns).min()``. Arguments: columns: Name(s) of the columns to use in the aggregation function. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> df_native = pa.table({"a": [1, 2], "b": [5, 10]}) >>> nw.from_native(df_native).select(nw.min("a", "b")) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | pyarrow.Table | | a: int64 | | b: int64 | | ---- | | a: [[1]] | | b: [[5]] | └──────────────────┘ """ return col(*columns).min() def max(*columns: str) -> Expr: """Return the maximum value. Note: Syntactic sugar for ``nw.col(columns).max()``. Arguments: columns: Name(s) of the columns to use in the aggregation function. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> df_native = pd.DataFrame({"a": [1, 2], "b": [5, 10]}) >>> nw.from_native(df_native).select(nw.max("a", "b")) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 2 10 | └──────────────────┘ """ return col(*columns).max() def _expr_with_horizontal_op(name: str, *exprs: IntoExpr, **kwargs: Any) -> Expr: if not exprs: msg = f"At least one expression must be passed to `{name}`" raise ValueError(msg) return Expr( ExprNode(ExprKind.ELEMENTWISE, name, *exprs, **kwargs, allow_multi_output=True) ) def sum_horizontal(*exprs: IntoExpr | Iterable[IntoExpr]) -> Expr: """Sum all values horizontally across columns. Warning: Unlike Polars, we support horizontal sum over numeric columns only. Arguments: exprs: Name(s) of the columns to use in the aggregation function. Accepts expression input. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> df_native = pl.DataFrame({"a": [1, 2, 3], "b": [5, 10, None]}) >>> nw.from_native(df_native).with_columns(sum=nw.sum_horizontal("a", "b")) ┌────────────────────┐ | Narwhals DataFrame | |--------------------| |shape: (3, 3) | |┌─────┮──────┮─────┐| |│ a ┆ b ┆ sum │| |│ --- ┆ --- ┆ --- │| |│ i64 ┆ i64 ┆ i64 │| |╞═════╩══════╩═════╡| |│ 1 ┆ 5 ┆ 6 │| |│ 2 ┆ 10 ┆ 12 │| |│ 3 ┆ null ┆ 3 │| |└─────â”ī──────â”ī─────┘| └────────────────────┘ """ return _expr_with_horizontal_op("sum_horizontal", *flatten(exprs)) def min_horizontal(*exprs: IntoExpr | Iterable[IntoExpr]) -> Expr: """Get the minimum value horizontally across columns. Notes: We support `min_horizontal` over numeric columns only. Arguments: exprs: Name(s) of the columns to use in the aggregation function. Accepts expression input. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> df_native = pa.table({"a": [1, 8, 3], "b": [4, 5, None]}) >>> nw.from_native(df_native).with_columns(h_min=nw.min_horizontal("a", "b")) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | pyarrow.Table | | a: int64 | | b: int64 | | h_min: int64 | | ---- | | a: [[1,8,3]] | | b: [[4,5,null]] | | h_min: [[1,5,3]] | └──────────────────┘ """ return _expr_with_horizontal_op("min_horizontal", *flatten(exprs)) def max_horizontal(*exprs: IntoExpr | Iterable[IntoExpr]) -> Expr: """Get the maximum value horizontally across columns. Notes: We support `max_horizontal` over numeric columns only. Arguments: exprs: Name(s) of the columns to use in the aggregation function. Accepts expression input. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> df_native = pl.DataFrame({"a": [1, 8, 3], "b": [4, 5, None]}) >>> nw.from_native(df_native).with_columns(h_max=nw.max_horizontal("a", "b")) ┌──────────────────────┐ | Narwhals DataFrame | |----------------------| |shape: (3, 3) | |┌─────┮──────┮───────┐| |│ a ┆ b ┆ h_max │| |│ --- ┆ --- ┆ --- │| |│ i64 ┆ i64 ┆ i64 │| |╞═════╩══════╩═══════╡| |│ 1 ┆ 4 ┆ 4 │| |│ 8 ┆ 5 ┆ 8 │| |│ 3 ┆ null ┆ 3 │| |└─────â”ī──────â”ī───────┘| └──────────────────────┘ """ return _expr_with_horizontal_op("max_horizontal", *flatten(exprs)) class When: def __init__(self, *predicates: IntoExpr | Iterable[IntoExpr]) -> None: self._predicate = all_horizontal(*flatten(predicates), ignore_nulls=False) def then(self, value: IntoExpr | NonNestedLiteral) -> Then: return Then( ExprNode( ExprKind.ELEMENTWISE, "when_then", self._predicate, value, allow_multi_output=False, ) ) class Then(Expr): def otherwise(self, value: IntoExpr | NonNestedLiteral) -> Expr: node = self._nodes[0] return Expr(ExprNode(ExprKind.ELEMENTWISE, "when_then", *node.exprs, value)) def when(*predicates: IntoExpr | Iterable[IntoExpr]) -> When: """Start a `when-then-otherwise` expression. Expression similar to an `if-else` statement in Python. Always initiated by a `pl.when().then()`, and optionally followed by a `.otherwise()` can be appended at the end. If not appended, and the condition is not `True`, `None` will be returned. Info: Chaining multiple `.when().then()` statements is currently not supported. See [Narwhals#668](https://github.com/narwhals-dev/narwhals/issues/668). Arguments: predicates: Condition(s) that must be met in order to apply the subsequent statement. Accepts one or more boolean expressions, which are implicitly combined with `&`. String input is parsed as a column name. Returns: A "when" object, which `.then` can be called on. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> data = {"a": [1, 2, 3], "b": [5, 10, 15]} >>> df_native = pd.DataFrame(data) >>> nw.from_native(df_native).with_columns( ... nw.when(nw.col("a") < 3).then(5).otherwise(6).alias("a_when") ... ) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b a_when | | 0 1 5 5 | | 1 2 10 5 | | 2 3 15 6 | └──────────────────┘ """ return When(*predicates) def all_horizontal(*exprs: IntoExpr | Iterable[IntoExpr], ignore_nulls: bool) -> Expr: r"""Compute the bitwise AND horizontally across columns. Arguments: exprs: Name(s) of the columns to use in the aggregation function. Accepts expression input. ignore_nulls: Whether to ignore nulls: - If `True`, null values are ignored. If there are no elements, the result is `True`. - If `False`, Kleene logic is followed. Note that this is not allowed for pandas with classical NumPy dtypes when null values are present. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> data = { ... "a": [False, False, True, True, False, None], ... "b": [False, True, True, None, None, None], ... } >>> df_native = pa.table(data) >>> nw.from_native(df_native).select( ... "a", "b", all=nw.all_horizontal("a", "b", ignore_nulls=False) ... ) ┌─────────────────────────────────────────┐ | Narwhals DataFrame | |-----------------------------------------| |pyarrow.Table | |a: bool | |b: bool | |all: bool | |---- | |a: [[false,false,true,true,false,null]] | |b: [[false,true,true,null,null,null]] | |all: [[false,false,true,null,false,null]]| └─────────────────────────────────────────┘ """ return _expr_with_horizontal_op( "all_horizontal", *flatten(exprs), ignore_nulls=ignore_nulls ) def lit(value: NonNestedLiteral, dtype: IntoDType | None = None) -> Expr: """Return an expression representing a literal value. Arguments: value: The value to use as literal. dtype: The data type of the literal value. If not provided, the data type will be inferred by the native library. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> df_native = pd.DataFrame({"a": [1, 2]}) >>> nw.from_native(df_native).with_columns(nw.lit(3)) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a literal | | 0 1 3 | | 1 2 3 | └──────────────────┘ """ if is_numpy_array(value): msg = ( "numpy arrays are not supported as literal values. " "Consider using `with_columns` to create a new column from the array." ) raise ValueError(msg) if isinstance(value, (list, tuple)): msg = f"Nested datatypes are not supported yet. Got {value}" raise NotImplementedError(msg) return Expr(ExprNode(ExprKind.LITERAL, "lit", value=value, dtype=dtype)) def any_horizontal(*exprs: IntoExpr | Iterable[IntoExpr], ignore_nulls: bool) -> Expr: r"""Compute the bitwise OR horizontally across columns. Arguments: exprs: Name(s) of the columns to use in the aggregation function. Accepts expression input. ignore_nulls: Whether to ignore nulls: - If `True`, null values are ignored. If there are no elements, the result is `False`. - If `False`, Kleene logic is followed. Note that this is not allowed for pandas with classical NumPy dtypes when null values are present. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> data = { ... "a": [False, False, True, True, False, None], ... "b": [False, True, True, None, None, None], ... } >>> df_native = pl.DataFrame(data) >>> nw.from_native(df_native).select( ... "a", "b", any=nw.any_horizontal("a", "b", ignore_nulls=False) ... ) ┌─────────────────────────┐ | Narwhals DataFrame | |-------------------------| |shape: (6, 3) | |┌───────┮───────┮───────┐| |│ a ┆ b ┆ any │| |│ --- ┆ --- ┆ --- │| |│ bool ┆ bool ┆ bool │| |╞═══════╩═══════╩═══════╡| |│ false ┆ false ┆ false │| |│ false ┆ true ┆ true │| |│ true ┆ true ┆ true │| |│ true ┆ null ┆ true │| |│ false ┆ null ┆ null │| |│ null ┆ null ┆ null │| |└───────â”ī───────â”ī───────┘| └─────────────────────────┘ """ return _expr_with_horizontal_op( "any_horizontal", *flatten(exprs), ignore_nulls=ignore_nulls ) def mean_horizontal(*exprs: IntoExpr | Iterable[IntoExpr]) -> Expr: """Compute the mean of all values horizontally across columns. Arguments: exprs: Name(s) of the columns to use in the aggregation function. Accepts expression input. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> data = {"a": [1, 8, 3], "b": [4, 5, None], "c": ["x", "y", "z"]} >>> df_native = pa.table(data) We define a dataframe-agnostic function that computes the horizontal mean of "a" and "b" columns: >>> nw.from_native(df_native).select(nw.mean_horizontal("a", "b")) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | pyarrow.Table | | a: double | | ---- | | a: [[2.5,6.5,3]] | └──────────────────┘ """ return _expr_with_horizontal_op("mean_horizontal", *flatten(exprs)) def concat_str( exprs: IntoExpr | Iterable[IntoExpr], *more_exprs: IntoExpr, separator: str = "", ignore_nulls: bool = False, ) -> Expr: r"""Horizontally concatenate columns into a single string column. Arguments: exprs: Columns to concatenate into a single string column. Accepts expression input. Strings are parsed as column names, other non-expression inputs are parsed as literals. Non-`String` columns are cast to `String`. *more_exprs: Additional columns to concatenate into a single string column, specified as positional arguments. separator: String that will be used to separate the values of each column. ignore_nulls: Ignore null values (default is `False`). If set to `False`, null values will be propagated and if the row contains any null values, the output is null. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> data = { ... "a": [1, 2, 3], ... "b": ["dogs", "cats", None], ... "c": ["play", "swim", "walk"], ... } >>> df_native = pd.DataFrame(data) >>> ( ... nw.from_native(df_native).select( ... nw.concat_str( ... [nw.col("a") * 2, nw.col("b"), nw.col("c")], separator=" " ... ).alias("full_sentence") ... ) ... ) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | full_sentence | | 0 2 dogs play | | 1 4 cats swim | | 2 None | └──────────────────┘ """ flat_exprs = flatten([*flatten([exprs]), *more_exprs]) return _expr_with_horizontal_op( "concat_str", *flat_exprs, separator=separator, ignore_nulls=ignore_nulls ) def coalesce( exprs: IntoExpr | Iterable[IntoExpr], *more_exprs: IntoExpr | NonNestedLiteral ) -> Expr: """Folds the columns from left to right, keeping the first non-null value. Arguments: exprs: Columns to coalesce, must be a str, nw.Expr, or nw.Series where strings are parsed as column names and both nw.Expr/nw.Series are passed through as-is. Scalar values must be wrapped in `nw.lit`. *more_exprs: Additional columns to coalesce, specified as positional arguments. Raises: TypeError: If any of the inputs are not a str, nw.Expr, or nw.Series. Examples: >>> import polars as pl >>> import narwhals as nw >>> data = [ ... (1, 5, None), ... (None, 6, None), ... (None, None, 9), ... (4, 8, 10), ... (None, None, None), ... ] >>> df = pl.DataFrame(data, schema=["a", "b", "c"], orient="row") >>> nw.from_native(df).select(nw.coalesce("a", "b", "c", nw.lit(-1))) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | shape: (5, 1) | | ┌─────┐ | | │ a │ | | │ --- │ | | │ i64 │ | | ╞═════╡ | | │ 1 │ | | │ 6 │ | | │ 9 │ | | │ 4 │ | | │ -1 │ | | └─────┘ | └──────────────────┘ """ flat_exprs = flatten([*flatten([exprs]), *more_exprs]) non_exprs = [ expr for expr in flat_exprs if not (isinstance(expr, str) or is_expr(expr) or is_series(expr)) ] if non_exprs: msg = ( f"All arguments to `coalesce` must be of type {str!r}, Expr, or Series." "\nGot the following invalid arguments (type, value):" f"\n {', '.join(repr((type(e), e)) for e in non_exprs)}" ) raise TypeError(msg) return Expr( ExprNode(ExprKind.ELEMENTWISE, "coalesce", *flat_exprs, allow_multi_output=True) ) def format(f_string: str, *args: IntoExpr) -> Expr: """Format expressions as a string. Arguments: f_string: A string that with placeholders. args: Expression(s) that fill the placeholders. Examples: >>> import duckdb >>> import narwhals as nw >>> rel = duckdb.sql("select * from values ('a', 1), ('b', 2), ('c', 3) df(a, b)") >>> df = nw.from_native(rel) >>> df.with_columns(formatted=nw.format("foo_{}_bar_{}", nw.col("a"), "b")) ┌─────────────────────────────────┐ | Narwhals LazyFrame | |---------------------------------| |┌─────────┮───────┮─────────────┐| |│ a │ b │ formatted │| |│ varchar │ int32 │ varchar │| |├─────────┾───────┾─────────────â”Ī| |│ a │ 1 │ foo_a_bar_1 │| |│ b │ 2 │ foo_b_bar_2 │| |│ c │ 3 │ foo_c_bar_3 │| |└─────────â”ī───────â”ī─────────────┘| └─────────────────────────────────┘ """ if (n_placeholders := f_string.count("{}")) != len(args): msg = f"number of placeholders should equal the number of arguments. Expected {n_placeholders} arguments, got {len(args)}." raise ValueError(msg) exprs = [] it = iter(args) for i, s in enumerate(f_string.split("{}")): if i > 0: exprs.append(next(it)) if len(s) > 0: exprs.append(lit(s)) return concat_str(exprs, separator="") narwhals-2.11.0/narwhals/group_by.py0000644000000000000000000001475013615410400014360 0ustar00from __future__ import annotations from typing import TYPE_CHECKING, Any, Generic, TypeVar from narwhals._expression_parsing import is_scalar_like from narwhals._utils import tupleify from narwhals.exceptions import InvalidOperationError from narwhals.typing import DataFrameT if TYPE_CHECKING: from collections.abc import Iterable, Iterator, Sequence from narwhals._compliant.typing import CompliantExprAny from narwhals.dataframe import LazyFrame from narwhals.expr import Expr LazyFrameT = TypeVar("LazyFrameT", bound="LazyFrame[Any]") class GroupBy(Generic[DataFrameT]): def __init__( self, df: DataFrameT, keys: Sequence[str] | Sequence[CompliantExprAny], /, *, drop_null_keys: bool, ) -> None: self._df: DataFrameT = df self._keys = keys self._grouped = self._df._compliant_frame.group_by( self._keys, drop_null_keys=drop_null_keys ) def agg(self, *aggs: Expr | Iterable[Expr], **named_aggs: Expr) -> DataFrameT: """Compute aggregations for each group of a group by operation. Arguments: aggs: Aggregations to compute for each group of the group by operation, specified as positional arguments. named_aggs: Additional aggregations, specified as keyword arguments. Examples: Group by one column or by multiple columns and call `agg` to compute the grouped sum of another column. >>> import pandas as pd >>> import narwhals as nw >>> df_native = pd.DataFrame( ... { ... "a": ["a", "b", "a", "b", "c"], ... "b": [1, 2, 1, 3, 3], ... "c": [5, 4, 3, 2, 1], ... } ... ) >>> df = nw.from_native(df_native) >>> >>> df.group_by("a").agg(nw.col("b").sum()).sort("a") ┌──────────────────┐ |Narwhals DataFrame| |------------------| | a b | | 0 a 2 | | 1 b 5 | | 2 c 3 | └──────────────────┘ >>> >>> df.group_by("a", "b").agg(nw.col("c").sum()).sort("a", "b").to_native() a b c 0 a 1 8 1 b 2 4 2 b 3 2 3 c 3 1 """ compliant_aggs = self._df._flatten_and_extract(*aggs, **named_aggs) if not all(is_scalar_like(x) for x in compliant_aggs): msg = ( "Found expression which does not aggregate.\n\n" "All expressions passed to GroupBy.agg must aggregate.\n" "For example, `df.group_by('a').agg(nw.col('b').sum())` is valid,\n" "but `df.group_by('a').agg(nw.col('b'))` is not." ) raise InvalidOperationError(msg) return self._df._with_compliant(self._grouped.agg(*compliant_aggs)) def __iter__(self) -> Iterator[tuple[Any, DataFrameT]]: yield from ( (tupleify(key), self._df._with_compliant(df)) for (key, df) in self._grouped.__iter__() ) class LazyGroupBy(Generic[LazyFrameT]): def __init__( self, df: LazyFrameT, keys: Sequence[str] | Sequence[CompliantExprAny], /, *, drop_null_keys: bool, ) -> None: self._df: LazyFrameT = df self._keys = keys self._grouped = self._df._compliant_frame.group_by( self._keys, drop_null_keys=drop_null_keys ) def agg(self, *aggs: Expr | Iterable[Expr], **named_aggs: Expr) -> LazyFrameT: """Compute aggregations for each group of a group by operation. Arguments: aggs: Aggregations to compute for each group of the group by operation, specified as positional arguments. named_aggs: Additional aggregations, specified as keyword arguments. Examples: Group by one column or by multiple columns and call `agg` to compute the grouped sum of another column. >>> import polars as pl >>> import narwhals as nw >>> from narwhals.typing import IntoFrameT >>> lf_native = pl.LazyFrame( ... { ... "a": ["a", "b", "a", "b", "c"], ... "b": [1, 2, 1, 3, 3], ... "c": [5, 4, 3, 2, 1], ... } ... ) >>> lf = nw.from_native(lf_native) >>> >>> nw.to_native(lf.group_by("a").agg(nw.col("b").sum()).sort("a")).collect() shape: (3, 2) ┌─────┮─────┐ │ a ┆ b │ │ --- ┆ --- │ │ str ┆ i64 │ ╞═════╩═════╡ │ a ┆ 2 │ │ b ┆ 5 │ │ c ┆ 3 │ └─────â”ī─────┘ >>> >>> lf.group_by("a", "b").agg(nw.sum("c")).sort("a", "b").collect() ┌───────────────────┐ |Narwhals DataFrame | |-------------------| |shape: (4, 3) | |┌─────┮─────┮─────┐| |│ a ┆ b ┆ c │| |│ --- ┆ --- ┆ --- │| |│ str ┆ i64 ┆ i64 │| |╞═════╩═════╩═════╡| |│ a ┆ 1 ┆ 8 │| |│ b ┆ 2 ┆ 4 │| |│ b ┆ 3 ┆ 2 │| |│ c ┆ 3 ┆ 1 │| |└─────â”ī─────â”ī─────┘| └───────────────────┘ """ compliant_aggs = self._df._flatten_and_extract(*aggs, **named_aggs) if not all(is_scalar_like(x) for x in compliant_aggs): msg = ( "Found expression which does not aggregate.\n\n" "All expressions passed to GroupBy.agg must aggregate.\n" "For example, `df.group_by('a').agg(nw.col('b').sum())` is valid,\n" "but `df.group_by('a').agg(nw.col('b'))` is not." ) raise InvalidOperationError(msg) return self._df._with_compliant(self._grouped.agg(*compliant_aggs)) narwhals-2.11.0/narwhals/plugins.py0000644000000000000000000000737113615410400014214 0ustar00from __future__ import annotations import sys from functools import cache from typing import TYPE_CHECKING, Any, Protocol, cast from narwhals._compliant import CompliantNamespace from narwhals._typing_compat import TypeVar if TYPE_CHECKING: from collections.abc import Iterator from importlib.metadata import EntryPoints from typing_extensions import LiteralString, TypeAlias from narwhals._compliant.typing import ( CompliantDataFrameAny, CompliantFrameAny, CompliantLazyFrameAny, CompliantSeriesAny, ) from narwhals.utils import Version __all__ = ["Plugin", "from_native"] CompliantAny: TypeAlias = ( "CompliantDataFrameAny | CompliantLazyFrameAny | CompliantSeriesAny" ) """A statically-unknown, Compliant object originating from a plugin.""" FrameT = TypeVar( "FrameT", bound="CompliantFrameAny", default="CompliantDataFrameAny | CompliantLazyFrameAny", ) FromNativeR_co = TypeVar( "FromNativeR_co", bound=CompliantAny, covariant=True, default=CompliantAny ) @cache def _discover_entrypoints() -> EntryPoints: from importlib.metadata import entry_points as eps group = "narwhals.plugins" if sys.version_info < (3, 10): return cast("EntryPoints", eps().get(group, ())) return eps(group=group) class PluginNamespace(CompliantNamespace[FrameT, Any], Protocol[FrameT, FromNativeR_co]): def from_native(self, data: Any, /) -> FromNativeR_co: ... class Plugin(Protocol[FrameT, FromNativeR_co]): NATIVE_PACKAGE: LiteralString def __narwhals_namespace__( self, version: Version ) -> PluginNamespace[FrameT, FromNativeR_co]: ... def is_native(self, native_object: object, /) -> bool: ... @cache def _might_be(cls: type, type_: str) -> bool: # pragma: no cover try: return any(type_ in o.__module__.split(".") for o in cls.mro()) except TypeError: return False def _is_native_plugin(native_object: Any, plugin: Plugin) -> bool: pkg = plugin.NATIVE_PACKAGE return ( sys.modules.get(pkg) is not None and _might_be(type(native_object), pkg) # type: ignore[arg-type] and plugin.is_native(native_object) ) def _iter_from_native(native_object: Any, version: Version) -> Iterator[CompliantAny]: for entry_point in _discover_entrypoints(): plugin: Plugin = entry_point.load() if _is_native_plugin(native_object, plugin): compliant_namespace = plugin.__narwhals_namespace__(version=version) yield compliant_namespace.from_native(native_object) def from_native(native_object: Any, version: Version) -> CompliantAny | None: """Attempt to convert `native_object` to a Compliant object, using any available plugin(s). Arguments: native_object: Raw object from user. version: Narwhals API version. Returns: If the following conditions are met - at least 1 plugin is installed - at least 1 installed plugin supports `type(native_object)` Then for the **first matching plugin**, the result of the call below. This *should* be an object accepted by a Narwhals Dataframe, Lazyframe, or Series: plugin: Plugin plugin.__narwhals_namespace__(version).from_native(native_object) In all other cases, `None` is returned instead. """ return next(_iter_from_native(native_object, version), None) def _show_suggestions(native_object_type: type) -> str | None: if _might_be(native_object_type, "daft"): # pragma: no cover return ( "Hint: it looks like you passed a `daft.DataFrame` but don't have `narwhals-daft` installed.\n" "Please refer to https://github.com/narwhals-dev/narwhals-daft for installation instructions." ) return None narwhals-2.11.0/narwhals/py.typed0000644000000000000000000000000013615410400013636 0ustar00narwhals-2.11.0/narwhals/schema.py0000644000000000000000000003012113615410400013760 0ustar00"""Schema. Adapted from Polars implementation at: https://github.com/pola-rs/polars/blob/main/py-polars/polars/schema.py. """ from __future__ import annotations from collections import OrderedDict from collections.abc import Mapping from functools import partial from typing import TYPE_CHECKING, cast from narwhals._utils import Implementation, Version, qualified_type_name, zip_strict from narwhals.dependencies import ( get_cudf, is_cudf_dtype, is_pandas_like_dtype, is_polars_data_type, is_polars_schema, is_pyarrow_data_type, is_pyarrow_schema, ) if TYPE_CHECKING: from collections.abc import Iterable from typing import Any, ClassVar import polars as pl import pyarrow as pa from typing_extensions import Self from narwhals.dtypes import DType from narwhals.typing import ( DTypeBackend, IntoArrowSchema, IntoPandasSchema, IntoPolarsSchema, ) __all__ = ["Schema"] class Schema(OrderedDict[str, "DType"]): """Ordered mapping of column names to their data type. Arguments: schema: The schema definition given by column names and their associated *instantiated* Narwhals data type. Accepts a mapping or an iterable of tuples. Examples: Define a schema by passing *instantiated* data types. >>> import narwhals as nw >>> schema = nw.Schema({"foo": nw.Int8(), "bar": nw.String()}) >>> schema Schema({'foo': Int8, 'bar': String}) Access the data type associated with a specific column name. >>> schema["foo"] Int8 Access various schema properties using the `names`, `dtypes`, and `len` methods. >>> schema.names() ['foo', 'bar'] >>> schema.dtypes() [Int8, String] >>> schema.len() 2 """ _version: ClassVar[Version] = Version.MAIN def __init__( self, schema: Mapping[str, DType] | Iterable[tuple[str, DType]] | None = None ) -> None: schema = schema or {} super().__init__(schema) def names(self) -> list[str]: """Get the column names of the schema.""" return list(self.keys()) def dtypes(self) -> list[DType]: """Get the data types of the schema.""" return list(self.values()) def len(self) -> int: """Get the number of columns in the schema.""" return len(self) @classmethod def from_arrow(cls, schema: IntoArrowSchema, /) -> Self: """Construct a Schema from a pyarrow Schema. Arguments: schema: A pyarrow Schema or mapping of column names to pyarrow data types. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> mapping = { ... "a": pa.timestamp("us", "UTC"), ... "b": pa.date32(), ... "c": pa.string(), ... "d": pa.uint8(), ... } >>> native = pa.schema(mapping) >>> >>> nw.Schema.from_arrow(native) Schema({'a': Datetime(time_unit='us', time_zone='UTC'), 'b': Date, 'c': String, 'd': UInt8}) >>> nw.Schema.from_arrow(mapping) == nw.Schema.from_arrow(native) True """ if isinstance(schema, Mapping): if not schema: return cls() import pyarrow as pa # ignore-banned-import schema = pa.schema(schema) from narwhals._arrow.utils import native_to_narwhals_dtype return cls( (field.name, native_to_narwhals_dtype(field.type, cls._version)) for field in schema ) @classmethod def from_pandas_like(cls, schema: IntoPandasSchema, /) -> Self: """Construct a Schema from a pandas-like schema representation. Arguments: schema: A mapping of column names to pandas-like data types. Examples: >>> import numpy as np >>> import pandas as pd >>> import pyarrow as pa >>> import narwhals as nw >>> >>> data = {"a": [1], "b": ["a"], "c": [False], "d": [9.2]} >>> native = pd.DataFrame(data).convert_dtypes().dtypes.to_dict() >>> >>> nw.Schema.from_pandas_like(native) Schema({'a': Int64, 'b': String, 'c': Boolean, 'd': Float64}) >>> >>> mapping = { ... "a": pd.DatetimeTZDtype("us", "UTC"), ... "b": pd.ArrowDtype(pa.date32()), ... "c": pd.StringDtype("python"), ... "d": np.dtype("uint8"), ... } >>> >>> nw.Schema.from_pandas_like(mapping) Schema({'a': Datetime(time_unit='us', time_zone='UTC'), 'b': Date, 'c': String, 'd': UInt8}) """ if not schema: return cls() impl = ( Implementation.CUDF if get_cudf() and any(is_cudf_dtype(dtype) for dtype in schema.values()) else Implementation.PANDAS ) return cls._from_pandas_like(schema, impl) @classmethod def from_native( cls, schema: IntoArrowSchema | IntoPolarsSchema | IntoPandasSchema, / ) -> Self: """Construct a Schema from a native schema representation. Arguments: schema: A native schema object, or mapping of column names to *instantiated* native data types. Examples: >>> import datetime as dt >>> import pyarrow as pa >>> import narwhals as nw >>> >>> data = {"a": [1], "b": ["a"], "c": [dt.time(1, 2, 3)], "d": [[2]]} >>> native = pa.table(data).schema >>> >>> nw.Schema.from_native(native) Schema({'a': Int64, 'b': String, 'c': Time, 'd': List(Int64)}) """ if is_pyarrow_schema(schema): return cls.from_arrow(schema) if is_polars_schema(schema): return cls.from_polars(schema) if isinstance(schema, Mapping): return cls._from_native_mapping(schema) if schema else cls() msg = ( f"Expected an arrow, polars, or pandas schema, but got " f"{qualified_type_name(schema)!r}\n\n{schema!r}" ) raise TypeError(msg) @classmethod def from_polars(cls, schema: IntoPolarsSchema, /) -> Self: """Construct a Schema from a polars Schema. Arguments: schema: A polars Schema or mapping of column names to *instantiated* polars data types. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> mapping = { ... "a": pl.Datetime(time_zone="UTC"), ... "b": pl.Date(), ... "c": pl.String(), ... "d": pl.UInt8(), ... } >>> native = pl.Schema(mapping) >>> >>> nw.Schema.from_polars(native) Schema({'a': Datetime(time_unit='us', time_zone='UTC'), 'b': Date, 'c': String, 'd': UInt8}) >>> nw.Schema.from_polars(mapping) == nw.Schema.from_polars(native) True """ if not schema: return cls() from narwhals._polars.utils import native_to_narwhals_dtype return cls( (name, native_to_narwhals_dtype(dtype, cls._version)) for name, dtype in schema.items() ) def to_arrow(self) -> pa.Schema: """Convert Schema to a pyarrow Schema. Examples: >>> import narwhals as nw >>> schema = nw.Schema({"a": nw.Int64(), "b": nw.Datetime("ns")}) >>> schema.to_arrow() a: int64 b: timestamp[ns] """ import pyarrow as pa # ignore-banned-import from narwhals._arrow.utils import narwhals_to_native_dtype return pa.schema( (name, narwhals_to_native_dtype(dtype, self._version)) for name, dtype in self.items() ) def to_pandas( self, dtype_backend: DTypeBackend | Iterable[DTypeBackend] = None ) -> dict[str, Any]: """Convert Schema to an ordered mapping of column names to their pandas data type. Arguments: dtype_backend: Backend(s) used for the native types. When providing more than one, the length of the iterable must be equal to the length of the schema. Examples: >>> import narwhals as nw >>> schema = nw.Schema({"a": nw.Int64(), "b": nw.Datetime("ns")}) >>> schema.to_pandas() {'a': 'int64', 'b': 'datetime64[ns]'} >>> schema.to_pandas("pyarrow") {'a': 'Int64[pyarrow]', 'b': 'timestamp[ns][pyarrow]'} """ from narwhals._pandas_like.utils import narwhals_to_native_dtype to_native_dtype = partial( narwhals_to_native_dtype, implementation=Implementation.PANDAS, version=self._version, ) if dtype_backend is None or isinstance(dtype_backend, str): return { name: to_native_dtype(dtype=dtype, dtype_backend=dtype_backend) for name, dtype in self.items() } backends = tuple(dtype_backend) if len(backends) != len(self): from itertools import chain, islice, repeat n_user, n_actual = len(backends), len(self) suggestion = tuple( islice(chain.from_iterable(islice(repeat(backends), n_actual)), n_actual) ) msg = ( f"Provided {n_user!r} `dtype_backend`(s), but schema contains {n_actual!r} field(s).\n" "Hint: instead of\n" f" schema.to_pandas({backends})\n" "you may want to use\n" f" schema.to_pandas({backends[0]})\n" f"or\n" f" schema.to_pandas({suggestion})" ) raise ValueError(msg) return { name: to_native_dtype(dtype=dtype, dtype_backend=backend) for name, dtype, backend in zip_strict(self.keys(), self.values(), backends) } def to_polars(self) -> pl.Schema: """Convert Schema to a polars Schema. Examples: >>> import narwhals as nw >>> schema = nw.Schema({"a": nw.Int64(), "b": nw.Datetime("ns")}) >>> schema.to_polars() Schema({'a': Int64, 'b': Datetime(time_unit='ns', time_zone=None)}) """ import polars as pl # ignore-banned-import from narwhals._polars.utils import narwhals_to_native_dtype pl_version = Implementation.POLARS._backend_version() schema = ( (name, narwhals_to_native_dtype(dtype, self._version)) for name, dtype in self.items() ) return ( pl.Schema(schema) if pl_version >= (1, 0, 0) else cast("pl.Schema", dict(schema)) ) @classmethod def _from_native_mapping( cls, native: Mapping[str, pa.DataType] | Mapping[str, pl.DataType] | IntoPandasSchema, /, ) -> Self: first_item = next(iter(native.items())) first_key, first_dtype = first_item if is_polars_data_type(first_dtype): return cls.from_polars(cast("IntoPolarsSchema", native)) if is_pandas_like_dtype(first_dtype): return cls.from_pandas_like(cast("IntoPandasSchema", native)) if is_pyarrow_data_type(first_dtype): return cls.from_arrow(cast("IntoArrowSchema", native)) msg = ( f"Expected an arrow, polars, or pandas dtype, but found " f"`{first_key}: {qualified_type_name(first_dtype)}`\n\n{native!r}" ) raise TypeError(msg) @classmethod def _from_pandas_like( cls, schema: IntoPandasSchema, implementation: Implementation, / ) -> Self: from narwhals._pandas_like.utils import native_to_narwhals_dtype impl = implementation return cls( (name, native_to_narwhals_dtype(dtype, cls._version, impl, allow_object=True)) for name, dtype in schema.items() ) narwhals-2.11.0/narwhals/selectors.py0000644000000000000000000002365313615410400014537 0ustar00from __future__ import annotations from typing import TYPE_CHECKING, Any, NoReturn from narwhals._expression_parsing import ExprKind, ExprNode from narwhals._utils import flatten from narwhals.expr import Expr if TYPE_CHECKING: from collections.abc import Iterable from datetime import timezone from narwhals.dtypes import DType from narwhals.typing import TimeUnit class Selector(Expr): def _to_expr(self) -> Expr: return Expr(*self._nodes) def __add__(self, other: Any) -> Expr: # type: ignore[override] if isinstance(other, Selector): msg = "unsupported operand type(s) for op: ('Selector' + 'Selector')" raise TypeError(msg) return self._to_expr()._append_node( ExprNode(ExprKind.ELEMENTWISE, "__add__", other, str_as_lit=True) ) def __or__(self, other: Any) -> Expr: # type: ignore[override] if isinstance(other, Selector): return self._append_node( ExprNode( ExprKind.ELEMENTWISE, "__or__", other, str_as_lit=True, allow_multi_output=True, ) ) return self._to_expr()._append_node( ExprNode(ExprKind.ELEMENTWISE, "__or__", other, str_as_lit=True) ) def __and__(self, other: Any) -> Expr: # type: ignore[override] if isinstance(other, Selector): return self._append_node( ExprNode( ExprKind.ELEMENTWISE, "__and__", other, str_as_lit=True, allow_multi_output=True, ) ) return self._to_expr()._append_node( ExprNode(ExprKind.ELEMENTWISE, "__and__", other, str_as_lit=True) ) def __rsub__(self, other: Any) -> NoReturn: raise NotImplementedError def __rand__(self, other: Any) -> NoReturn: raise NotImplementedError def __ror__(self, other: Any) -> NoReturn: raise NotImplementedError def by_dtype(*dtypes: DType | type[DType] | Iterable[DType | type[DType]]) -> Selector: """Select columns based on their dtype. Arguments: dtypes: one or data types to select Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> import narwhals.selectors as ncs >>> df_native = pa.table({"a": [1, 2], "b": ["x", "y"], "c": [4.1, 2.3]}) >>> df = nw.from_native(df_native) Let's select int64 and float64 dtypes and multiply each value by 2: >>> df.select(ncs.by_dtype(nw.Int64, nw.Float64) * 2).to_native() pyarrow.Table a: int64 c: double ---- a: [[2,4]] c: [[8.2,4.6]] """ flattened = flatten(dtypes) return Selector(ExprNode(ExprKind.SELECTOR, "selectors.by_dtype", dtypes=flattened)) def matches(pattern: str) -> Selector: """Select all columns that match the given regex pattern. Arguments: pattern: A valid regular expression pattern. Examples: >>> import pandas as pd >>> import narwhals as nw >>> import narwhals.selectors as ncs >>> df_native = pd.DataFrame( ... {"bar": [123, 456], "baz": [2.0, 5.5], "zap": [0, 1]} ... ) >>> df = nw.from_native(df_native) Let's select column names containing an 'a', preceded by a character that is not 'z': >>> df.select(ncs.matches("[^z]a")).to_native() bar baz 0 123 2.0 1 456 5.5 """ return Selector(ExprNode(ExprKind.SELECTOR, "selectors.matches", pattern=pattern)) def numeric() -> Selector: """Select numeric columns. Examples: >>> import polars as pl >>> import narwhals as nw >>> import narwhals.selectors as ncs >>> df_native = pl.DataFrame({"a": [1, 2], "b": ["x", "y"], "c": [4.1, 2.3]}) >>> df = nw.from_native(df_native) Let's select numeric dtypes and multiply each value by 2: >>> df.select(ncs.numeric() * 2).to_native() shape: (2, 2) ┌─────┮─────┐ │ a ┆ c │ │ --- ┆ --- │ │ i64 ┆ f64 │ ╞═════╩═════╡ │ 2 ┆ 8.2 │ │ 4 ┆ 4.6 │ └─────â”ī─────┘ """ return Selector(ExprNode(ExprKind.SELECTOR, "selectors.numeric")) def boolean() -> Selector: """Select boolean columns. Examples: >>> import polars as pl >>> import narwhals as nw >>> import narwhals.selectors as ncs >>> df_native = pl.DataFrame({"a": [1, 2], "b": ["x", "y"], "c": [False, True]}) >>> df = nw.from_native(df_native) Let's select boolean dtypes: >>> df.select(ncs.boolean()) ┌──────────────────┐ |Narwhals DataFrame| |------------------| | shape: (2, 1) | | ┌───────┐ | | │ c │ | | │ --- │ | | │ bool │ | | ╞═══════╡ | | │ false │ | | │ true │ | | └───────┘ | └──────────────────┘ """ return Selector(ExprNode(ExprKind.SELECTOR, "selectors.boolean")) def string() -> Selector: """Select string columns. Examples: >>> import polars as pl >>> import narwhals as nw >>> import narwhals.selectors as ncs >>> df_native = pl.DataFrame({"a": [1, 2], "b": ["x", "y"], "c": [False, True]}) >>> df = nw.from_native(df_native) Let's select string dtypes: >>> df.select(ncs.string()).to_native() shape: (2, 1) ┌─────┐ │ b │ │ --- │ │ str │ ╞═════╡ │ x │ │ y │ └─────┘ """ return Selector(ExprNode(ExprKind.SELECTOR, "selectors.string")) def categorical() -> Selector: """Select categorical columns. Examples: >>> import polars as pl >>> import narwhals as nw >>> import narwhals.selectors as ncs >>> df_native = pl.DataFrame({"a": [1, 2], "b": ["x", "y"], "c": [False, True]}) Let's convert column "b" to categorical, and then select categorical dtypes: >>> df = nw.from_native(df_native).with_columns( ... b=nw.col("b").cast(nw.Categorical()) ... ) >>> df.select(ncs.categorical()).to_native() shape: (2, 1) ┌─────┐ │ b │ │ --- │ │ cat │ ╞═════╡ │ x │ │ y │ └─────┘ """ return Selector(ExprNode(ExprKind.SELECTOR, "selectors.categorical")) def all() -> Selector: """Select all columns. Examples: >>> import pandas as pd >>> import narwhals as nw >>> import narwhals.selectors as ncs >>> df_native = pd.DataFrame({"a": [1, 2], "b": ["x", "y"], "c": [False, True]}) >>> df = nw.from_native(df_native) Let's select all dtypes: >>> df.select(ncs.all()).to_native() a b c 0 1 x False 1 2 y True """ return Selector(ExprNode(ExprKind.SELECTOR, "selectors.all")) def datetime( time_unit: TimeUnit | Iterable[TimeUnit] | None = None, time_zone: str | timezone | Iterable[str | timezone | None] | None = ("*", None), ) -> Selector: """Select all datetime columns, optionally filtering by time unit/zone. Arguments: time_unit: One (or more) of the allowed timeunit precision strings, "ms", "us", "ns" and "s". Omit to select columns with any valid timeunit. time_zone: Specify which timezone(s) to select * One or more timezone strings, as defined in zoneinfo (to see valid options run `import zoneinfo; zoneinfo.available_timezones()` for a full list). * Set `None` to select Datetime columns that do not have a timezone. * Set `"*"` to select Datetime columns that have *any* timezone. Examples: >>> from datetime import datetime, timezone >>> import pyarrow as pa >>> import narwhals as nw >>> import narwhals.selectors as ncs >>> >>> utc_tz = timezone.utc >>> data = { ... "tstamp_utc": [ ... datetime(2023, 4, 10, 12, 14, 16, 999000, tzinfo=utc_tz), ... datetime(2025, 8, 25, 14, 18, 22, 666000, tzinfo=utc_tz), ... ], ... "tstamp": [ ... datetime(2000, 11, 20, 18, 12, 16, 600000), ... datetime(2020, 10, 30, 10, 20, 25, 123000), ... ], ... "numeric": [3.14, 6.28], ... } >>> df_native = pa.table(data) >>> df_nw = nw.from_native(df_native) >>> df_nw.select(ncs.datetime()).to_native() pyarrow.Table tstamp_utc: timestamp[us, tz=UTC] tstamp: timestamp[us] ---- tstamp_utc: [[2023-04-10 12:14:16.999000Z,2025-08-25 14:18:22.666000Z]] tstamp: [[2000-11-20 18:12:16.600000,2020-10-30 10:20:25.123000]] Select only datetime columns that have any time_zone specification: >>> df_nw.select(ncs.datetime(time_zone="*")).to_native() pyarrow.Table tstamp_utc: timestamp[us, tz=UTC] ---- tstamp_utc: [[2023-04-10 12:14:16.999000Z,2025-08-25 14:18:22.666000Z]] """ return Selector( ExprNode( ExprKind.SELECTOR, "selectors.datetime", time_unit=time_unit, time_zone=time_zone, ) ) __all__ = [ "all", "boolean", "by_dtype", "categorical", "datetime", "matches", "numeric", "string", ] narwhals-2.11.0/narwhals/series.py0000644000000000000000000030250713615410400014024 0ustar00from __future__ import annotations import math from collections.abc import Iterable, Iterator, Mapping, Sequence from typing import ( TYPE_CHECKING, Any, Callable, ClassVar, Generic, Literal, cast, overload, ) from narwhals._expression_parsing import ExprKind, ExprNode from narwhals._utils import ( Implementation, Version, _Implementation, _validate_rolling_arguments, ensure_type, generate_repr, is_compliant_series, is_eager_allowed, is_index_selector, no_default, qualified_type_name, supports_arrow_c_stream, ) from narwhals.dependencies import is_numpy_array, is_numpy_array_1d, is_numpy_scalar from narwhals.dtypes import _validate_dtype, _validate_into_dtype from narwhals.exceptions import ComputeError, InvalidOperationError from narwhals.expr import Expr from narwhals.functions import col from narwhals.series_cat import SeriesCatNamespace from narwhals.series_dt import SeriesDateTimeNamespace from narwhals.series_list import SeriesListNamespace from narwhals.series_str import SeriesStringNamespace from narwhals.series_struct import SeriesStructNamespace from narwhals.translate import to_native from narwhals.typing import IntoSeriesT if TYPE_CHECKING: from types import ModuleType from typing import NoReturn import pandas as pd import polars as pl import pyarrow as pa from typing_extensions import Self from narwhals._compliant import CompliantSeries from narwhals._typing import EagerAllowed, IntoBackend, NoDefault from narwhals.dataframe import DataFrame, MultiIndexSelector from narwhals.dtypes import DType from narwhals.typing import ( ClosedInterval, FillNullStrategy, IntoDType, ModeKeepStrategy, NonNestedLiteral, NumericLiteral, PythonLiteral, RankMethod, RollingInterpolationMethod, SingleIndexSelector, TemporalLiteral, _1DArray, ) class Series(Generic[IntoSeriesT]): """Narwhals Series, backed by a native series. Warning: This class is not meant to be instantiated directly - instead: - If the native object is a series from one of the supported backend (e.g. pandas.Series, polars.Series, pyarrow.ChunkedArray), you can use [`narwhals.from_native`][]: ```py narwhals.from_native(native_series, allow_series=True) narwhals.from_native(native_series, series_only=True) ``` - If the object is a generic sequence (e.g. a list or a tuple of values), you can create a series via [`narwhals.new_series`][], e.g.: ```py narwhals.new_series(name="price", values=[10.5, 9.4, 1.2], backend="pandas") ``` """ _version: ClassVar[Version] = Version.MAIN @property def _compliant(self) -> CompliantSeries[IntoSeriesT]: return self._compliant_series @property def _dataframe(self) -> type[DataFrame[Any]]: from narwhals.dataframe import DataFrame return DataFrame def _to_expr(self) -> Expr: return Expr( ExprNode(ExprKind.SERIES, "_expr._from_series", series=self._compliant) ) def __init__( self, series: Any, *, level: Literal["full", "lazy", "interchange"] ) -> None: self._level: Literal["full", "lazy", "interchange"] = level if is_compliant_series(series): self._compliant_series: CompliantSeries[IntoSeriesT] = ( series.__narwhals_series__() ) else: # pragma: no cover msg = f"Expected Polars Series or an object which implements `__narwhals_series__`, got: {type(series)}." raise AssertionError(msg) @classmethod def from_numpy( cls, name: str, values: _1DArray, dtype: IntoDType | None = None, *, backend: IntoBackend[EagerAllowed], ) -> Series[Any]: """Construct a Series from a NumPy ndarray. Arguments: name: Name of resulting Series. values: One-dimensional data represented as a NumPy ndarray. dtype: (Narwhals) dtype. If not provided, the native library may auto-infer it from `values`. backend: specifies which eager backend instantiate to. `backend` can be specified in various ways - As `Implementation.` with `BACKEND` being `PANDAS`, `PYARROW`, `POLARS`, `MODIN` or `CUDF`. - As a string: `"pandas"`, `"pyarrow"`, `"polars"`, `"modin"` or `"cudf"`. - Directly as a module `pandas`, `pyarrow`, `polars`, `modin` or `cudf`. Examples: >>> import numpy as np >>> import polars as pl >>> import narwhals as nw >>> >>> arr = np.arange(5, 10) >>> nw.Series.from_numpy("arr", arr, dtype=nw.Int8, backend="polars") ┌──────────────────┐ | Narwhals Series | |------------------| |shape: (5,) | |Series: 'arr' [i8]| |[ | | 5 | | 6 | | 7 | | 8 | | 9 | |] | └──────────────────┘ """ if not is_numpy_array_1d(values): msg = "`from_numpy` only accepts 1D numpy arrays" raise ValueError(msg) if dtype: _validate_into_dtype(dtype) implementation = Implementation.from_backend(backend) if is_eager_allowed(implementation): ns = cls._version.namespace.from_backend(implementation).compliant compliant = ns.from_numpy(values).alias(name) if dtype: return cls(compliant.cast(dtype), level="full") return cls(compliant, level="full") msg = ( f"{implementation} support in Narwhals is lazy-only, but `Series.from_numpy` is an eager-only function.\n\n" "Hint: you may want to use an eager backend and then call `.lazy`, e.g.:\n\n" f" nw.Series.from_numpy(arr, backend='pyarrow').to_frame().lazy('{implementation}')" ) raise ValueError(msg) @classmethod def from_iterable( cls, name: str, values: Iterable[Any], dtype: IntoDType | None = None, *, backend: IntoBackend[EagerAllowed], ) -> Series[Any]: """Construct a Series from an iterable. Arguments: name: Name of resulting Series. values: One-dimensional data represented as an iterable. dtype: (Narwhals) dtype. If not provided, the native library may auto-infer it from `values`. backend: specifies which eager backend instantiate to. `backend` can be specified in various ways - As `Implementation.` with `BACKEND` being `PANDAS`, `PYARROW`, `POLARS`, `MODIN` or `CUDF`. - As a string: `"pandas"`, `"pyarrow"`, `"polars"`, `"modin"` or `"cudf"`. - Directly as a module `pandas`, `pyarrow`, `polars`, `modin` or `cudf`. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> values = [4, 1, 3, 2] >>> nw.Series.from_iterable("a", values, dtype=nw.UInt32, backend="pandas") ┌──────────────────────┐ | Narwhals Series | |----------------------| |0 4 | |1 1 | |2 3 | |3 2 | |Name: a, dtype: uint32| └──────────────────────┘ """ if is_numpy_array(values): return cls.from_numpy(name, values, dtype, backend=backend) if dtype: _validate_into_dtype(dtype) if not isinstance(values, Iterable): msg = f"Expected values to be an iterable, got: {qualified_type_name(values)!r}." raise TypeError(msg) implementation = Implementation.from_backend(backend) if is_eager_allowed(implementation): ns = cls._version.namespace.from_backend(implementation).compliant compliant = ns._series.from_iterable( values, context=ns, name=name, dtype=dtype ) return cls(compliant, level="full") msg = ( f"{implementation} support in Narwhals is lazy-only, but `Series.from_iterable` is an eager-only function.\n\n" "Hint: you may want to use an eager backend and then call `.lazy`, e.g.:\n\n" f" nw.Series.from_iterable('a', [1,2,3], backend='pyarrow').to_frame().lazy('{implementation}')" ) raise ValueError(msg) implementation: _Implementation = _Implementation() """Return [`narwhals.Implementation`][] of native Series. This can be useful when you need to use special-casing for features outside of Narwhals' scope - for example, when dealing with pandas' Period Dtype. Examples: >>> import narwhals as nw >>> import pandas as pd >>> s_native = pd.Series([1, 2, 3]) >>> s = nw.from_native(s_native, series_only=True) >>> s.implementation >>> s.implementation.is_pandas() True >>> s.implementation.is_pandas_like() True >>> s.implementation.is_polars() False """ def __bool__(self) -> NoReturn: msg = ( "the truth value of a Series is ambiguous" "\n\n" "Here are some things you might want to try:\n" "- instead of `if s`, use `if not s.is_empty()`\n" "- instead of `s1 and s2`, use `s1 & s2`\n" "- instead of `s1 or s2`, use `s1 | s2`\n" "- instead of `s in [y, z]`, use `s.is_in([y, z])`\n" ) raise TypeError(msg) def __array__(self, dtype: Any = None, copy: bool | None = None) -> _1DArray: # noqa: FBT001 return self._compliant_series.__array__(dtype=dtype, copy=copy) @overload def __getitem__(self, idx: SingleIndexSelector) -> Any: ... @overload def __getitem__(self, idx: MultiIndexSelector) -> Self: ... def __getitem__(self, idx: SingleIndexSelector | MultiIndexSelector) -> Any | Self: """Retrieve elements from the object using integer indexing or slicing. Arguments: idx: The index, slice, or sequence of indices to retrieve. - If `idx` is an integer, a single element is returned. - If `idx` is a slice, a sequence of integers, or another Series (with integer values) a subset of the Series is returned. Returns: A single element if `idx` is an integer, else a subset of the Series. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[1, 2, 3]]) >>> nw.from_native(s_native, series_only=True)[0] 1 >>> nw.from_native(s_native, series_only=True)[ ... :2 ... ].to_native() # doctest:+ELLIPSIS [ [ 1, 2 ] ] """ if isinstance(idx, int) or ( is_numpy_scalar(idx) and idx.dtype.kind in {"i", "u"} ): idx = int(idx) if not isinstance(idx, int) else idx return self._compliant_series.item(idx) if isinstance(idx, self.to_native().__class__): idx = self._with_compliant(self._compliant_series._with_native(idx)) if not is_index_selector(idx): msg = ( f"Unexpected type for `Series.__getitem__`: {type(idx)}.\n\n" "Hints:\n" "- use `s.item` to select a single item.\n" "- Use `s[indices]` to select rows positionally.\n" "- Use `s.filter(mask)` to filter rows based on a boolean mask." ) raise TypeError(msg) if isinstance(idx, Series): return self._with_compliant(self._compliant_series[idx._compliant_series]) assert not isinstance(idx, int) # noqa: S101 # help mypy return self._with_compliant(self._compliant_series[idx]) def __native_namespace__(self) -> ModuleType: return self._compliant_series.__native_namespace__() def __arrow_c_stream__(self, requested_schema: object | None = None) -> object: """Export a Series via the Arrow PyCapsule Interface. Narwhals doesn't implement anything itself here: - if the underlying series implements the interface, it'll return that - else, it'll call `to_arrow` and then defer to PyArrow's implementation See [PyCapsule Interface](https://arrow.apache.org/docs/dev/format/CDataInterface/PyCapsuleInterface.html) for more. """ native_series = self._compliant_series.native if supports_arrow_c_stream(native_series): return native_series.__arrow_c_stream__(requested_schema=requested_schema) try: pa_version = Implementation.PYARROW._backend_version() except ModuleNotFoundError as exc: msg = f"'pyarrow>=16.0.0' is required for `Series.__arrow_c_stream__` for object of type {type(native_series)}" raise ModuleNotFoundError(msg) from exc if pa_version < (16, 0): # pragma: no cover msg = f"'pyarrow>=16.0.0' is required for `Series.__arrow_c_stream__` for object of type {type(native_series)}" raise ModuleNotFoundError(msg) from narwhals._arrow.utils import chunked_array ca = chunked_array(self.to_arrow()) return ca.__arrow_c_stream__(requested_schema=requested_schema) def to_native(self) -> IntoSeriesT: """Convert Narwhals series to native series. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series([1, 2]) >>> s = nw.from_native(s_native, series_only=True) >>> s.to_native() # doctest: +NORMALIZE_WHITESPACE shape: (2,) Series: '' [i64] [ 1 2 ] """ return self._compliant_series.native def scatter(self, indices: int | Sequence[int], values: Any) -> Self: """Set value(s) at given position(s). Arguments: indices: Position(s) to set items at. values: Values to set. Note: This method always returns a new Series, without modifying the original one. Using this function in a for-loop is an anti-pattern, we recommend building up your positions and values beforehand and doing an update in one go. For example, instead of ```python for i in [1, 3, 2]: value = some_function(i) s = s.scatter(i, value) ``` prefer ```python positions = [1, 3, 2] values = [some_function(x) for x in positions] s = s.scatter(positions, values) ``` Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> df_native = pa.table({"a": [1, 2, 3], "b": [4, 5, 6]}) >>> df_nw = nw.from_native(df_native) >>> df_nw.with_columns(df_nw["a"].scatter([0, 1], [999, 888])).to_native() pyarrow.Table a: int64 b: int64 ---- a: [[999,888,3]] b: [[4,5,6]] """ return self._with_compliant( self._compliant_series.scatter(indices, self._extract_native(values)) ) @property def shape(self) -> tuple[int]: """Get the shape of the Series. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series([1, 2, 3]) >>> nw.from_native(s_native, series_only=True).shape (3,) """ return (self._compliant_series.len(),) def _extract_native(self, arg: Any) -> Any: from narwhals.series import Series if isinstance(arg, Series): return arg._compliant_series return arg def _with_compliant(self, series: Any) -> Self: return self.__class__(series, level=self._level) def pipe(self, function: Callable[[Any], Self], *args: Any, **kwargs: Any) -> Self: """Pipe function call. Examples: >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series([1, 2, 3]) >>> s = nw.from_native(s_native, series_only=True) >>> s.pipe(lambda x: x + 2).to_native() # doctest: +NORMALIZE_WHITESPACE shape: (3,) Series: '' [i64] [ 3 4 5 ] """ return function(self, *args, **kwargs) def __repr__(self) -> str: # pragma: no cover return generate_repr("Narwhals Series", self.to_native().__repr__()) def __len__(self) -> int: return len(self._compliant_series) def len(self) -> int: r"""Return the number of elements in the Series. Null values count towards the total. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[1, 2, None]]) >>> nw.from_native(s_native, series_only=True).len() 3 """ return len(self._compliant_series) @property def dtype(self) -> DType: """Get the data type of the Series. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series([1, 2, 3]) >>> nw.from_native(s_native, series_only=True).dtype Int64 """ return self._compliant_series.dtype @property def name(self) -> str: """Get the name of the Series. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series("foo", [1, 2, 3]) >>> nw.from_native(s_native, series_only=True).name 'foo' """ return self._compliant_series.name def ewm_mean( self, *, com: float | None = None, span: float | None = None, half_life: float | None = None, alpha: float | None = None, adjust: bool = True, min_samples: int = 1, ignore_nulls: bool = False, ) -> Self: r"""Compute exponentially-weighted moving average. Arguments: com: Specify decay in terms of center of mass, $\gamma$, with
$\alpha = \frac{1}{1+\gamma}\forall\gamma\geq0$ span: Specify decay in terms of span, $\theta$, with
$\alpha = \frac{2}{\theta + 1} \forall \theta \geq 1$ half_life: Specify decay in terms of half-life, $\tau$, with
$\alpha = 1 - \exp \left\{ \frac{ -\ln(2) }{ \tau } \right\} \forall \tau > 0$ alpha: Specify smoothing factor alpha directly, $0 < \alpha \leq 1$. adjust: Divide by decaying adjustment factor in beginning periods to account for imbalance in relative weightings - When `adjust=True` (the default) the EW function is calculated using weights $w_i = (1 - \alpha)^i$ - When `adjust=False` the EW function is calculated recursively by $$ y_0=x_0 $$ $$ y_t = (1 - \alpha)y_{t - 1} + \alpha x_t $$ min_samples: Minimum number of observations in window required to have a value (otherwise result is null). ignore_nulls: Ignore missing values when calculating weights. - When `ignore_nulls=False` (default), weights are based on absolute positions. For example, the weights of $x_0$ and $x_2$ used in calculating the final weighted average of $[x_0, None, x_2]$ are $(1-\alpha)^2$ and $1$ if `adjust=True`, and $(1-\alpha)^2$ and $\alpha$ if `adjust=False`. - When `ignore_nulls=True`, weights are based on relative positions. For example, the weights of $x_0$ and $x_2$ used in calculating the final weighted average of $[x_0, None, x_2]$ are $1-\alpha$ and $1$ if `adjust=True`, and $1-\alpha$ and $\alpha$ if `adjust=False`. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series(name="a", data=[1, 2, 3]) >>> nw.from_native(s_native, series_only=True).ewm_mean( ... com=1, ignore_nulls=False ... ).to_native() 0 1.000000 1 1.666667 2 2.428571 Name: a, dtype: float64 """ return self._with_compliant( self._compliant_series.ewm_mean( com=com, span=span, half_life=half_life, alpha=alpha, adjust=adjust, min_samples=min_samples, ignore_nulls=ignore_nulls, ) ) def cast(self, dtype: IntoDType) -> Self: """Cast between data types. Arguments: dtype: Data type that the object will be cast into. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[True, False, True]]) >>> nw.from_native(s_native, series_only=True).cast(nw.Int64).to_native() [ [ 1, 0, 1 ] ] """ _validate_dtype(dtype) return self._with_compliant(self._compliant_series.cast(dtype)) def to_frame(self) -> DataFrame[Any]: """Convert to dataframe. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series("a", [1, 2]) >>> nw.from_native(s_native, series_only=True).to_frame().to_native() shape: (2, 1) ┌─────┐ │ a │ │ --- │ │ i64 │ ╞═════╡ │ 1 │ │ 2 │ └─────┘ """ return self._dataframe(self._compliant_series.to_frame(), level=self._level) def to_list(self) -> list[Any]: """Convert to list. Notes: This function converts to Python scalars. It's typically more efficient to keep your data in the format native to your original dataframe, so we recommend only calling this when you absolutely need to. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[1, 2, 3]]) >>> nw.from_native(s_native, series_only=True).to_list() [1, 2, 3] """ return self._compliant_series.to_list() def mean(self) -> float: """Reduce this Series to the mean value. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series([1.2, 4.2]) >>> nw.from_native(s_native, series_only=True).mean() np.float64(2.7) """ return self._compliant_series.mean() def median(self) -> float: """Reduce this Series to the median value. Notes: Results might slightly differ across backends due to differences in the underlying algorithms used to compute the median. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[5, 3, 8]]) >>> nw.from_native(s_native, series_only=True).median() 5.0 """ return self._compliant_series.median() def skew(self) -> float | None: """Calculate the sample skewness of the Series. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series([1, 1, 2, 10, 100]) >>> nw.from_native(s_native, series_only=True).skew() 1.4724267269058975 Notes: The skewness is a measure of the asymmetry of the probability distribution. A perfectly symmetric distribution has a skewness of 0. """ return self._compliant_series.skew() def kurtosis(self) -> float | None: """Compute the kurtosis (Fisher's definition) without bias correction. Kurtosis is the fourth central moment divided by the square of the variance. The Fisher's definition is used where 3.0 is subtracted from the result to give 0.0 for a normal distribution. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series([1, 1, 2, 10, 100]) >>> nw.from_native(s_native, series_only=True).kurtosis() 0.2106571340718002 """ return self._compliant_series.kurtosis() def count(self) -> int: """Returns the number of non-null elements in the Series. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[1, 2, None]]) >>> nw.from_native(s_native, series_only=True).count() 2 """ return self._compliant_series.count() def any(self) -> bool: """Return whether any of the values in the Series are True. If there are no non-null elements, the result is `False`. Notes: Only works on Series of data type Boolean. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series([False, True, False]) >>> nw.from_native(s_native, series_only=True).any() np.True_ """ return self._compliant_series.any() def all(self) -> bool: """Return whether all values in the Series are True. If there are no non-null elements, the result is `True`. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[False, True, False]]) >>> nw.from_native(s_native, series_only=True).all() False """ return self._compliant_series.all() def min(self) -> Any: """Get the minimal value in this Series. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series([1, 2, 3]) >>> nw.from_native(s_native, series_only=True).min() 1 """ return self._compliant_series.min() def max(self) -> Any: """Get the maximum value in this Series. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series([1, 2, 3]) >>> nw.from_native(s_native, series_only=True).max() np.int64(3) """ return self._compliant_series.max() def arg_min(self) -> int: """Returns the index of the minimum value. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[1, 2, 3]]) >>> nw.from_native(s_native, series_only=True).arg_min() 0 """ return self._compliant_series.arg_min() def arg_max(self) -> int: """Returns the index of the maximum value. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series([1, 2, 3]) >>> nw.from_native(s_native, series_only=True).arg_max() 2 """ return self._compliant_series.arg_max() def sum(self) -> float: """Reduce this Series to the sum value. If there are no non-null elements, the result is zero. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[1, 2, 3]]) >>> nw.from_native(s_native, series_only=True).sum() 6 """ return self._compliant_series.sum() def std(self, *, ddof: int = 1) -> float: """Get the standard deviation of this Series. Arguments: ddof: "Delta Degrees of Freedom": the divisor used in the calculation is N - ddof, where N represents the number of elements. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series([1, 2, 3]) >>> nw.from_native(s_native, series_only=True).std() 1.0 """ return self._compliant_series.std(ddof=ddof) def var(self, *, ddof: int = 1) -> float: """Get the variance of this Series. Arguments: ddof: "Delta Degrees of Freedom": the divisor used in the calculation is N - ddof, where N represents the number of elements. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[1, 2, 3]]) >>> nw.from_native(s_native, series_only=True).var() 1.0 """ return self._compliant_series.var(ddof=ddof) def clip( self, lower_bound: Self | NumericLiteral | TemporalLiteral | None = None, upper_bound: Self | NumericLiteral | TemporalLiteral | None = None, ) -> Self: r"""Clip values in the Series. Arguments: lower_bound: Lower bound value. upper_bound: Upper bound value. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series([-1, 1, -3, 3, -5, 5]) >>> nw.from_native(s_native, series_only=True).clip(-1, 3).to_native() 0 -1 1 1 2 -1 3 3 4 -1 5 3 dtype: int64 """ if lower_bound is None: return self._with_compliant( self._compliant_series.clip_upper( upper_bound=self._extract_native(upper_bound) ) ) if upper_bound is None: return self._with_compliant( self._compliant_series.clip_lower( lower_bound=self._extract_native(lower_bound) ) ) return self._with_compliant( self._compliant_series.clip( lower_bound=self._extract_native(lower_bound), upper_bound=self._extract_native(upper_bound), ) ) def first(self) -> PythonLiteral: """Get the first element of the Series. Returns: A scalar value or `None` if the Series is empty. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series([1, 2, 3]) >>> s_nw = nw.from_native(s_native, series_only=True) >>> s_nw.first() 1 >>> s_nw.filter(s_nw > 5).first() is None True """ return self._compliant_series.first() def last(self) -> PythonLiteral: """Get the last element of the Series. Returns: A scalar value or `None` if the Series is empty. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[1, 2, 3]]) >>> s_nw = nw.from_native(s_native, series_only=True) >>> s_nw.last() 3 >>> s_nw.filter(s_nw > 5).last() is None True """ return self._compliant_series.last() def is_in(self, other: Any) -> Self: """Check if the elements of this Series are in the other sequence. Arguments: other: Sequence of primitive type. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[1, 2, 3]]) >>> s = nw.from_native(s_native, series_only=True) >>> s.is_in([3, 2, 8]).to_native() # doctest: +ELLIPSIS [ [ false, true, true ] ] """ return self._with_compliant( self._compliant_series.is_in(to_native(other, pass_through=True)) ) def arg_true(self) -> Self: """Find elements where boolean Series is True. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series([1, None, None, 2]) >>> nw.from_native( ... s_native, series_only=True ... ).is_null().arg_true().to_native() # doctest: +NORMALIZE_WHITESPACE shape: (2,) Series: '' [u32] [ 1 2 ] """ return self._with_compliant(self._compliant_series.arg_true()) def drop_nulls(self) -> Self: """Drop null values. Notes: pandas handles null values differently from Polars and PyArrow. See [null_handling](../concepts/null_handling.md/) for reference. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series([2, 4, None, 3, 5]) >>> nw.from_native(s_native, series_only=True).drop_nulls().to_native() 0 2.0 1 4.0 3 3.0 4 5.0 dtype: float64 """ return self._with_compliant(self._compliant_series.drop_nulls()) def abs(self) -> Self: """Calculate the absolute value of each element. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[2, -4, 3]]) >>> nw.from_native( ... s_native, series_only=True ... ).abs().to_native() # doctest: +ELLIPSIS [ [ 2, 4, 3 ] ] """ return self._with_compliant(self._compliant_series.abs()) def cum_sum(self, *, reverse: bool = False) -> Self: """Calculate the cumulative sum. Arguments: reverse: reverse the operation Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series([2, 4, 3]) >>> nw.from_native(s_native, series_only=True).cum_sum().to_native() 0 2 1 6 2 9 dtype: int64 """ return self._with_compliant(self._compliant_series.cum_sum(reverse=reverse)) def unique(self, *, maintain_order: bool = False) -> Self: """Returns unique values of the series. Arguments: maintain_order: Keep the same order as the original series. This may be more expensive to compute. Settings this to `True` blocks the possibility to run on the streaming engine for Polars. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series([2, 4, 4, 6]) >>> s = nw.from_native(s_native, series_only=True) >>> s.unique( ... maintain_order=True ... ).to_native() # doctest: +NORMALIZE_WHITESPACE shape: (3,) Series: '' [i64] [ 2 4 6 ] """ return self._with_compliant( self._compliant_series.unique(maintain_order=maintain_order) ) def diff(self) -> Self: """Calculate the difference with the previous element, for each element. Notes: pandas may change the dtype here, for example when introducing missing values in an integer column. To ensure, that the dtype doesn't change, you may want to use `fill_null` and `cast`. For example, to calculate the diff and fill missing values with `0` in a Int64 column, you could do: s.diff().fill_null(0).cast(nw.Int64) Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[2, 4, 3]]) >>> nw.from_native( ... s_native, series_only=True ... ).diff().to_native() # doctest: +ELLIPSIS [ [ null, 2, -1 ] ] """ return self._with_compliant(self._compliant_series.diff()) def shift(self, n: int) -> Self: """Shift values by `n` positions. Arguments: n: Number of indices to shift forward. If a negative value is passed, values are shifted in the opposite direction instead. Notes: pandas may change the dtype here, for example when introducing missing values in an integer column. To ensure, that the dtype doesn't change, you may want to use `fill_null` and `cast`. For example, to shift and fill missing values with `0` in a Int64 column, you could do: s.shift(1).fill_null(0).cast(nw.Int64) Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series([2, 4, 3]) >>> nw.from_native(s_native, series_only=True).shift(1).to_native() 0 NaN 1 2.0 2 4.0 dtype: float64 """ ensure_type(n, int, param_name="n") return self._with_compliant(self._compliant_series.shift(n)) def sample( self, n: int | None = None, *, fraction: float | None = None, with_replacement: bool = False, seed: int | None = None, ) -> Self: """Sample randomly from this Series. Arguments: n: Number of items to return. Cannot be used with fraction. fraction: Fraction of items to return. Cannot be used with n. with_replacement: Allow values to be sampled more than once. seed: Seed for the random number generator. If set to None (default), a random seed is generated for each sample operation. Notes: The `sample` method returns a Series with a specified number of randomly selected items chosen from this Series. The results are not consistent across libraries. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series([1, 2, 3, 4]) >>> s = nw.from_native(s_native, series_only=True) >>> s.sample( ... fraction=1.0, with_replacement=True ... ).to_native() # doctest: +SKIP shape: (4,) Series: '' [i64] [ 1 4 3 4 ] """ return self._with_compliant( self._compliant_series.sample( n=n, fraction=fraction, with_replacement=with_replacement, seed=seed ) ) def alias(self, name: str) -> Self: """Rename the Series. Notes: This method is very cheap, but does not guarantee that data will be copied. For example: ```python s1: nw.Series s2 = s1.alias("foo") arr = s2.to_numpy() arr[0] = 999 ``` may (depending on the backend, and on the version) result in `s1`'s data being modified. We recommend: - if you need to alias an object and don't need the original one around any more, just use `alias` without worrying about it. - if you were expecting `alias` to copy data, then explicitly call `.clone` before calling `alias`. Arguments: name: The new name. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series([1, 2, 3], name="foo") >>> nw.from_native(s_native, series_only=True).alias("bar").to_native() 0 1 1 2 2 3 Name: bar, dtype: int64 """ return self._with_compliant(self._compliant_series.alias(name=name)) def rename(self, name: str) -> Self: """Rename the Series. Alias for `Series.alias()`. Notes: This method is very cheap, but does not guarantee that data will be copied. For example: ```python s1: nw.Series s2 = s1.rename("foo") arr = s2.to_numpy() arr[0] = 999 ``` may (depending on the backend, and on the version) result in `s1`'s data being modified. We recommend: - if you need to rename an object and don't need the original one around any more, just use `rename` without worrying about it. - if you were expecting `rename` to copy data, then explicitly call `.clone` before calling `rename`. Arguments: name: The new name. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series("foo", [1, 2, 3]) >>> s = nw.from_native(s_native, series_only=True) >>> s.rename("bar").to_native() # doctest: +NORMALIZE_WHITESPACE shape: (3,) Series: 'bar' [i64] [ 1 2 3 ] """ return self.alias(name=name) def replace_strict( self, old: Sequence[Any] | Mapping[Any, Any], new: Sequence[Any] | None = None, *, default: Any | NoDefault = no_default, return_dtype: IntoDType | None = None, ) -> Self: """Replace all values by different values. Arguments: old: Sequence of values to replace. It also accepts a mapping of values to their replacement as syntactic sugar for `replace_strict(old=list(mapping.keys()), new=list(mapping.values()))`. new: Sequence of values to replace by. Length must match the length of `old`. default: Set values that were not replaced to this value. If no default is specified, (default), an error is raised if any values were not replaced. Accepts expression input. Non-expression inputs are parsed as literals. return_dtype: The data type of the resulting expression. If set to `None` (default), the data type is determined automatically based on the other inputs. Raises: InvalidOperationError: If any non-null values in the original column were not replaced, and no default was specified. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series([3, 0, 1, 2], name="a") >>> nw.from_native(s_native, series_only=True).replace_strict( ... [0, 1, 2, 3], ["zero", "one", "two", "three"], return_dtype=nw.String ... ).to_native() 0 three 1 zero 2 one 3 two Name: a, dtype: object Replace values and set a default for values not in the mapping: >>> s_native = pd.Series([1, 2, 3, 4], name="a") >>> s_default = pd.Series(["beluga", "narwhal", "orca", "vaquita"]) >>> nw.from_native(s_native, series_only=True).replace_strict( ... {1: "one", 2: "two"}, ... default=nw.from_native(s_default, series_only=True), ... return_dtype=nw.String, ... ).to_native() 0 one 1 two 2 orca 3 vaquita Name: a, dtype: object """ if new is None: if not isinstance(old, Mapping): msg = "`new` argument is required if `old` argument is not a Mapping type" raise TypeError(msg) new = list(old.values()) old = list(old.keys()) assert isinstance(new, Sequence) # noqa: S101, help mypy assert isinstance(old, Sequence) # noqa: S101, help mypy return self._with_compliant( self._compliant_series.replace_strict( old=old, new=new, default=self._extract_native(default), return_dtype=return_dtype, ) ) def sort(self, *, descending: bool = False, nulls_last: bool = False) -> Self: """Sort this Series. Place null values first. Arguments: descending: Sort in descending order. nulls_last: Place null values last instead of first. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series([5, None, 1, 2]) >>> s = nw.from_native(s_native, series_only=True) >>> s.sort(descending=True).to_native() # doctest: +NORMALIZE_WHITESPACE shape: (4,) Series: '' [i64] [ null 5 2 1 ] """ return self._with_compliant( self._compliant_series.sort(descending=descending, nulls_last=nulls_last) ) def is_null(self) -> Self: """Returns a boolean Series indicating which values are null. Notes: pandas handles null values differently from Polars and PyArrow. See [null_handling](../concepts/null_handling.md/) for reference. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[1, 2, None]]) >>> nw.from_native( ... s_native, series_only=True ... ).is_null().to_native() # doctest:+ELLIPSIS [ [ false, false, true ] ] """ return self._with_compliant(self._compliant_series.is_null()) def is_nan(self) -> Self: """Returns a boolean Series indicating which values are NaN. Notes: pandas handles null values differently from Polars and PyArrow. See [null_handling](../concepts/null_handling.md/) for reference. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series([0.0, None, 2.0], dtype="Float64") >>> nw.from_native(s_native, series_only=True).is_nan().to_native() 0 False 1 2 False dtype: boolean """ return self._with_compliant(self._compliant_series.is_nan()) def fill_null( self, value: Self | NonNestedLiteral = None, strategy: FillNullStrategy | None = None, limit: int | None = None, ) -> Self: """Fill null values using the specified value. Arguments: value: Value used to fill null values. strategy: Strategy used to fill null values. limit: Number of consecutive null values to fill when using the 'forward' or 'backward' strategy. Notes: - pandas handles null values differently from other libraries. See [null_handling](../concepts/null_handling.md/) for reference. - For pandas Series of `object` dtype, `fill_null` will not automatically change the Series' dtype as pandas used to do. Explicitly call `cast` if you want the dtype to change. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series([1, 2, None]) >>> >>> nw.from_native(s_native, series_only=True).fill_null(5).to_native() 0 1.0 1 2.0 2 5.0 dtype: float64 Or using a strategy: >>> nw.from_native(s_native, series_only=True).fill_null( ... strategy="forward", limit=1 ... ).to_native() 0 1.0 1 2.0 2 2.0 dtype: float64 """ if value is not None and strategy is not None: msg = "cannot specify both `value` and `strategy`" raise ValueError(msg) if value is None and strategy is None: msg = "must specify either a fill `value` or `strategy`" raise ValueError(msg) if strategy is not None and strategy not in {"forward", "backward"}: msg = f"strategy not supported: {strategy}" raise ValueError(msg) return self._with_compliant( self._compliant_series.fill_null( value=self._extract_native(value), strategy=strategy, limit=limit ) ) def fill_nan(self, value: float | None) -> Self: """Fill floating point NaN values with given value. Arguments: value: Value used to fill NaN values. Notes: This function only fills `'NaN'` values, not null ones, except for pandas which doesn't distinguish between them. See [null_handling](../concepts/null_handling.md/) for reference. Examples: >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series([1.0, 2.0, float("nan"), None]) >>> result = nw.from_native(s_native, series_only=True).fill_nan(0) >>> result.to_native() # doctest: +NORMALIZE_WHITESPACE shape: (4,) Series: '' [f64] [ 1.0 2.0 0.0 null ] """ return self._with_compliant( self._compliant_series.fill_nan(value=self._extract_native(value)) ) def is_between( self, lower_bound: Any | Self, upper_bound: Any | Self, closed: ClosedInterval = "both", ) -> Self: """Get a boolean mask of the values that are between the given lower/upper bounds. Arguments: lower_bound: Lower bound value. upper_bound: Upper bound value. closed: Define which sides of the interval are closed (inclusive). Options are {"left", "right", "none", "both"}. Notes: If the value of the `lower_bound` is greater than that of the `upper_bound`, then the values will be False, as no value can satisfy the condition. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[1, 2, 3, 4, 5]]) >>> s = nw.from_native(s_native, series_only=True) >>> s.is_between(2, 4, "right").to_native() # doctest: +ELLIPSIS [ [ false, false, true, true, false ] ] """ return self._with_compliant( self._compliant_series.is_between( self._extract_native(lower_bound), self._extract_native(upper_bound), closed=closed, ) ) def n_unique(self) -> int: """Count the number of unique values. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series([1, 2, 2, 3]) >>> nw.from_native(s_native, series_only=True).n_unique() 3 """ return self._compliant_series.n_unique() def to_numpy(self) -> _1DArray: """Convert to numpy. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series([1, 2, 3], name="a") >>> nw.from_native(s_native, series_only=True).to_numpy() array([1, 2, 3]...) """ return self._compliant_series.to_numpy(None, copy=None) def to_pandas(self) -> pd.Series[Any]: """Convert to pandas Series. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series("a", [1, 2, 3]) >>> nw.from_native(s_native, series_only=True).to_pandas() 0 1 1 2 2 3 Name: a, dtype: int64 """ return self._compliant_series.to_pandas() def to_polars(self) -> pl.Series: """Convert to polars Series. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[1, 2, 3]]) >>> nw.from_native( ... s_native, series_only=True ... ).to_polars() # doctest: +NORMALIZE_WHITESPACE shape: (3,) Series: '' [i64] [ 1 2 3 ] """ return self._compliant_series.to_polars() def __add__(self, other: object) -> Self: return self._with_compliant( self._compliant_series.__add__(self._extract_native(other)) ) def __radd__(self, other: object) -> Self: return self._with_compliant( self._compliant_series.__radd__(self._extract_native(other)) ) def __sub__(self, other: object) -> Self: return self._with_compliant( self._compliant_series.__sub__(self._extract_native(other)) ) def __rsub__(self, other: object) -> Self: return self._with_compliant( self._compliant_series.__rsub__(self._extract_native(other)) ) def __mul__(self, other: object) -> Self: return self._with_compliant( self._compliant_series.__mul__(self._extract_native(other)) ) def __rmul__(self, other: object) -> Self: return self._with_compliant( self._compliant_series.__rmul__(self._extract_native(other)) ) def __truediv__(self, other: object) -> Self: return self._with_compliant( self._compliant_series.__truediv__(self._extract_native(other)) ) def __rtruediv__(self, other: object) -> Self: return self._with_compliant( self._compliant_series.__rtruediv__(self._extract_native(other)) ) def __floordiv__(self, other: object) -> Self: return self._with_compliant( self._compliant_series.__floordiv__(self._extract_native(other)) ) def __rfloordiv__(self, other: object) -> Self: return self._with_compliant( self._compliant_series.__rfloordiv__(self._extract_native(other)) ) def __pow__(self, other: object) -> Self: return self._with_compliant( self._compliant_series.__pow__(self._extract_native(other)) ) def __rpow__(self, other: object) -> Self: return self._with_compliant( self._compliant_series.__rpow__(self._extract_native(other)) ) def __mod__(self, other: object) -> Self: return self._with_compliant( self._compliant_series.__mod__(self._extract_native(other)) ) def __rmod__(self, other: object) -> Self: return self._with_compliant( self._compliant_series.__rmod__(self._extract_native(other)) ) def __eq__(self, other: object) -> Self: # type: ignore[override] return self._with_compliant( self._compliant_series.__eq__(self._extract_native(other)) ) def __ne__(self, other: object) -> Self: # type: ignore[override] return self._with_compliant( self._compliant_series.__ne__(self._extract_native(other)) ) def __gt__(self, other: Any) -> Self: return self._with_compliant( self._compliant_series.__gt__(self._extract_native(other)) ) def __ge__(self, other: Any) -> Self: return self._with_compliant( self._compliant_series.__ge__(self._extract_native(other)) ) def __lt__(self, other: Any) -> Self: return self._with_compliant( self._compliant_series.__lt__(self._extract_native(other)) ) def __le__(self, other: Any) -> Self: return self._with_compliant( self._compliant_series.__le__(self._extract_native(other)) ) def __and__(self, other: Any) -> Self: return self._with_compliant( self._compliant_series.__and__(self._extract_native(other)) ) def __rand__(self, other: Any) -> Self: return self._with_compliant( self._compliant_series.__rand__(self._extract_native(other)) ) def __or__(self, other: Any) -> Self: return self._with_compliant( self._compliant_series.__or__(self._extract_native(other)) ) def __ror__(self, other: Any) -> Self: return self._with_compliant( self._compliant_series.__ror__(self._extract_native(other)) ) # unary def __invert__(self) -> Self: return self._with_compliant(self._compliant_series.__invert__()) def filter(self, predicate: Any) -> Self: """Filter elements in the Series based on a condition. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series([4, 10, 15, 34, 50]) >>> s_nw = nw.from_native(s_native, series_only=True) >>> s_nw.filter(s_nw > 10).to_native() 2 15 3 34 4 50 dtype: int64 """ return self._with_compliant( self._compliant_series.filter(self._extract_native(predicate)) ) # --- descriptive --- def is_duplicated(self) -> Self: r"""Get a mask of all duplicated rows in the Series. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[1, 2, 3, 1]]) >>> nw.from_native( ... s_native, series_only=True ... ).is_duplicated().to_native() # doctest: +ELLIPSIS [ [ true, false, false, true ] ] """ return self._with_compliant(self._compliant_series.is_duplicated()) def is_empty(self) -> bool: r"""Check if the series is empty. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series([1, 2, 3]) >>> s_nw = nw.from_native(s_native, series_only=True) >>> s_nw.is_empty() False >>> s_nw.filter(s_nw > 10).is_empty() True """ return self._compliant_series.is_empty() def is_unique(self) -> Self: r"""Get a mask of all unique rows in the Series. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series([1, 2, 3, 1]) >>> nw.from_native(s_native, series_only=True).is_unique().to_native() 0 False 1 True 2 True 3 False dtype: bool """ return self._with_compliant(self._compliant_series.is_unique()) def null_count(self) -> int: r"""Count the number of null values. Notes: pandas handles null values differently from Polars and PyArrow. See [null_handling](../concepts/null_handling.md/) for reference. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[1, None, None]]) >>> nw.from_native(s_native, series_only=True).null_count() 2 """ return self._compliant_series.null_count() def is_first_distinct(self) -> Self: r"""Return a boolean mask indicating the first occurrence of each distinct value. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series([1, 1, 2, 3, 2]) >>> nw.from_native( ... s_native, series_only=True ... ).is_first_distinct().to_native() # doctest: +NORMALIZE_WHITESPACE shape: (5,) Series: '' [bool] [ true false true true false ] """ return self._with_compliant(self._compliant_series.is_first_distinct()) def is_last_distinct(self) -> Self: r"""Return a boolean mask indicating the last occurrence of each distinct value. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series([1, 1, 2, 3, 2]) >>> nw.from_native(s_native, series_only=True).is_last_distinct().to_native() 0 False 1 True 2 False 3 True 4 True dtype: bool """ return self._with_compliant(self._compliant_series.is_last_distinct()) def is_sorted(self, *, descending: bool = False) -> bool: r"""Check if the Series is sorted. Arguments: descending: Check if the Series is sorted in descending order. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[3, 2, 1]]) >>> s_nw = nw.from_native(s_native, series_only=True) >>> s_nw.is_sorted(descending=False) False >>> s_nw.is_sorted(descending=True) True """ return self._compliant_series.is_sorted(descending=descending) def value_counts( self, *, sort: bool = False, parallel: bool = False, name: str | None = None, normalize: bool = False, ) -> DataFrame[Any]: r"""Count the occurrences of unique values. Arguments: sort: Sort the output by count in descending order. If set to False (default), the order of the output is random. parallel: Execute the computation in parallel. Used for Polars only. name: Give the resulting count column a specific name; if `normalize` is True defaults to "proportion", otherwise defaults to "count". normalize: If true gives relative frequencies of the unique values - The original values as first column - Either count or proportion as second column, depending on normalize parameter. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series([1, 1, 2, 3, 2], name="s") >>> nw.from_native(s_native, series_only=True).value_counts( ... sort=True ... ).to_native() s count 0 1 2 1 2 2 2 3 1 """ return self._dataframe( self._compliant_series.value_counts( sort=sort, parallel=parallel, name=name, normalize=normalize ), level=self._level, ) def quantile( self, quantile: float, interpolation: RollingInterpolationMethod ) -> float: """Get quantile value of the series. Note: pandas and Polars may have implementation differences for a given interpolation method. Arguments: quantile: Quantile between 0.0 and 1.0. interpolation: Interpolation method. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series(list(range(50))) >>> s_nw = nw.from_native(s_native, series_only=True) >>> [ ... s_nw.quantile(quantile=q, interpolation="nearest") ... for q in (0.1, 0.25, 0.5, 0.75, 0.9) ... ] [5.0, 12.0, 25.0, 37.0, 44.0] """ return self._compliant_series.quantile( quantile=quantile, interpolation=interpolation ) def zip_with(self, mask: Self, other: Self) -> Self: """Take values from self or other based on the given mask. Where mask evaluates true, take values from self. Where mask evaluates false, take values from other. Arguments: mask: Boolean Series other: Series of same type. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> data_native = pa.chunked_array([[1, 2, 3, 4, 5]]) >>> other_native = pa.chunked_array([[5, 4, 3, 2, 1]]) >>> mask_native = pa.chunked_array([[True, False, True, False, True]]) >>> >>> data_nw = nw.from_native(data_native, series_only=True) >>> other_nw = nw.from_native(other_native, series_only=True) >>> mask_nw = nw.from_native(mask_native, series_only=True) >>> >>> data_nw.zip_with(mask_nw, other_nw).to_native() # doctest: +ELLIPSIS [ [ 1, 4, 3, 2, 5 ] ] """ return self._with_compliant( self._compliant_series.zip_with( self._extract_native(mask), self._extract_native(other) ) ) def item(self, index: int | None = None) -> Any: r"""Return the Series as a scalar, or return the element at the given index. If no index is provided, this is equivalent to `s[0]`, with a check that the shape is (1,). With an index, this is equivalent to `s[index]`. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> nw.from_native(pl.Series("a", [1]), series_only=True).item() 1 >>> nw.from_native(pl.Series("a", [9, 8, 7]), series_only=True).item(-1) 7 """ return self._compliant_series.item(index=index) def head(self, n: int = 10) -> Self: r"""Get the first `n` rows. Arguments: n: Number of rows to return. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series(list(range(10))) >>> nw.from_native(s_native, series_only=True).head(3).to_native() 0 0 1 1 2 2 dtype: int64 """ return self._with_compliant(self._compliant_series.head(n)) def tail(self, n: int = 10) -> Self: r"""Get the last `n` rows. Arguments: n: Number of rows to return. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([list(range(10))]) >>> s = nw.from_native(s_native, series_only=True) >>> s.tail(3).to_native() # doctest: +ELLIPSIS [ [ 7, 8, 9 ] ] """ return self._with_compliant(self._compliant_series.tail(n)) def round(self, decimals: int = 0) -> Self: r"""Round underlying floating point data by `decimals` digits. Arguments: decimals: Number of decimals to round by. Notes: For values exactly halfway between rounded decimal values pandas behaves differently than Polars and Arrow. pandas rounds to the nearest even value (e.g. -0.5 and 0.5 round to 0.0, 1.5 and 2.5 round to 2.0, 3.5 and 4.5 to 4.0, etc..). Polars and Arrow round away from 0 (e.g. -0.5 to -1.0, 0.5 to 1.0, 1.5 to 2.0, 2.5 to 3.0, etc..). Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series([1.12345, 2.56789, 3.901234]) >>> s = nw.from_native(s_native, series_only=True) >>> s.round(1).to_native() # doctest: +NORMALIZE_WHITESPACE shape: (3,) Series: '' [f64] [ 1.1 2.6 3.9 ] """ return self._with_compliant(self._compliant_series.round(decimals)) def floor(self) -> Self: r"""Compute the numerical floor. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> s_native = pa.chunked_array([[1.1, 4.3, -1.3]]) >>> s = nw.from_native(s_native, series_only=True) >>> s.floor().to_native() # doctest:+ELLIPSIS [ [ 1, 4, -2 ] ] """ return self._with_compliant(self._compliant_series.floor()) def ceil(self) -> Self: r"""Compute the numerical ceiling. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> s_native = pa.chunked_array([[1.1, 4.3, -1.3]]) >>> s = nw.from_native(s_native, series_only=True) >>> s.ceil().to_native() # doctest:+ELLIPSIS [ [ 2, 5, -1 ] ] """ return self._with_compliant(self._compliant_series.ceil()) def to_dummies( self, *, separator: str = "_", drop_first: bool = False ) -> DataFrame[Any]: r"""Get dummy/indicator variables. Arguments: separator: Separator/delimiter used when generating column names. drop_first: Remove the first category from the variable being encoded. Notes: pandas and Polars handle null values differently. Polars distinguishes between NaN and Null, whereas pandas doesn't. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series([1, 2, 3], name="a") >>> s_nw = nw.from_native(s_native, series_only=True) >>> s_nw.to_dummies(drop_first=False).to_native() a_1 a_2 a_3 0 1 0 0 1 0 1 0 2 0 0 1 >>> s_nw.to_dummies(drop_first=True).to_native() a_2 a_3 0 0 0 1 1 0 2 0 1 """ return self._dataframe( self._compliant_series.to_dummies(separator=separator, drop_first=drop_first), level=self._level, ) def gather_every(self, n: int, offset: int = 0) -> Self: r"""Take every nth value in the Series and return as new Series. Arguments: n: Gather every *n*-th row. offset: Starting index. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[1, 2, 3, 4]]) >>> nw.from_native(s_native, series_only=True).gather_every( ... n=2, offset=1 ... ).to_native() # doctest:+ELLIPSIS [ [ 2, 4 ] ] """ return self._with_compliant( self._compliant_series.gather_every(n=n, offset=offset) ) def to_arrow(self) -> pa.Array[Any]: r"""Convert to arrow. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series([1, 2, 3, 4]) >>> nw.from_native( ... s_native, series_only=True ... ).to_arrow() # doctest:+NORMALIZE_WHITESPACE [ 1, 2, 3, 4 ] """ return self._compliant_series.to_arrow() @overload def mode(self, *, keep: Literal["all"] = "all") -> Self: ... @overload def mode(self, *, keep: Literal["any"]) -> NonNestedLiteral: ... def mode(self, *, keep: ModeKeepStrategy = "all") -> Self | NonNestedLiteral: r"""Compute the most occurring value(s). Can return multiple values. Note: For `keep="any"` a scalar is returned, while for `keep="all"` a Series in returned even in the case of unimodal values. Arguments: keep: Whether to keep all modes or any mode found. Remark that `keep='any'` is not deterministic for multimodal values. Examples: >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series([1, 1, 2, 2, 3]) >>> nw.from_native(s_native, series_only=True).mode().sort().to_native() 0 1 1 2 dtype: int64 """ _supported_keep_values = ("all", "any") if keep not in _supported_keep_values: # pragma: no cover msg = f"`keep` must be one of {_supported_keep_values}, found '{keep}'" raise ValueError(msg) result = self._with_compliant(self._compliant_series.mode(keep=keep)) return result.item(0) if keep == "any" else result def is_finite(self) -> Self: """Returns a boolean Series indicating which values are finite. Warning: Different backend handle null values differently. `is_finite` will return False for NaN and Null's in the Dask and pandas non-nullable backend, while for Polars, PyArrow and pandas nullable backends null values are kept as such. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[float("nan"), float("inf"), 2.0, None]]) >>> nw.from_native( ... s_native, series_only=True ... ).is_finite().to_native() # doctest: +ELLIPSIS [ [ false, false, true, null ] ] """ return self._with_compliant(self._compliant_series.is_finite()) def cum_count(self, *, reverse: bool = False) -> Self: r"""Return the cumulative count of the non-null values in the series. Arguments: reverse: reverse the operation Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series(["x", "k", None, "d"]) >>> nw.from_native(s_native, series_only=True).cum_count( ... reverse=True ... ).to_native() # doctest:+NORMALIZE_WHITESPACE shape: (4,) Series: '' [u32] [ 3 2 1 1 ] """ return self._with_compliant(self._compliant_series.cum_count(reverse=reverse)) def cum_min(self, *, reverse: bool = False) -> Self: r"""Return the cumulative min of the non-null values in the series. Arguments: reverse: reverse the operation Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series([3, 1, None, 2]) >>> nw.from_native(s_native, series_only=True).cum_min().to_native() 0 3.0 1 1.0 2 NaN 3 1.0 dtype: float64 """ return self._with_compliant(self._compliant_series.cum_min(reverse=reverse)) def cum_max(self, *, reverse: bool = False) -> Self: r"""Return the cumulative max of the non-null values in the series. Arguments: reverse: reverse the operation Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[1, 3, None, 2]]) >>> nw.from_native( ... s_native, series_only=True ... ).cum_max().to_native() # doctest:+ELLIPSIS [ [ 1, 3, null, 3 ] ] """ return self._with_compliant(self._compliant_series.cum_max(reverse=reverse)) def cum_prod(self, *, reverse: bool = False) -> Self: r"""Return the cumulative product of the non-null values in the series. Arguments: reverse: reverse the operation Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series([1, 3, None, 2]) >>> nw.from_native( ... s_native, series_only=True ... ).cum_prod().to_native() # doctest:+NORMALIZE_WHITESPACE shape: (4,) Series: '' [i64] [ 1 3 null 6 ] """ return self._with_compliant(self._compliant_series.cum_prod(reverse=reverse)) def rolling_sum( self, window_size: int, *, min_samples: int | None = None, center: bool = False ) -> Self: """Apply a rolling sum (moving sum) over the values. A window of length `window_size` will traverse the values. The resulting values will be aggregated to their sum. The window at a given row will include the row itself and the `window_size - 1` elements before it. Arguments: window_size: The length of the window in number of elements. It must be a strictly positive integer. min_samples: The number of values in the window that should be non-null before computing a result. If set to `None` (default), it will be set equal to `window_size`. If provided, it must be a strictly positive integer, and less than or equal to `window_size` center: Set the labels at the center of the window. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series([1.0, 2.0, 3.0, 4.0]) >>> nw.from_native(s_native, series_only=True).rolling_sum( ... window_size=2 ... ).to_native() 0 NaN 1 3.0 2 5.0 3 7.0 dtype: float64 """ window_size, min_samples_int = _validate_rolling_arguments( window_size=window_size, min_samples=min_samples ) if len(self) == 0: # pragma: no cover return self return self._with_compliant( self._compliant_series.rolling_sum( window_size=window_size, min_samples=min_samples_int, center=center ) ) def rolling_mean( self, window_size: int, *, min_samples: int | None = None, center: bool = False ) -> Self: """Apply a rolling mean (moving mean) over the values. A window of length `window_size` will traverse the values. The resulting values will be aggregated to their mean. The window at a given row will include the row itself and the `window_size - 1` elements before it. Arguments: window_size: The length of the window in number of elements. It must be a strictly positive integer. min_samples: The number of values in the window that should be non-null before computing a result. If set to `None` (default), it will be set equal to `window_size`. If provided, it must be a strictly positive integer, and less than or equal to `window_size` center: Set the labels at the center of the window. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[1.0, 2.0, 3.0, 4.0]]) >>> nw.from_native(s_native, series_only=True).rolling_mean( ... window_size=2 ... ).to_native() # doctest:+ELLIPSIS [ [ null, 1.5, 2.5, 3.5 ] ] """ window_size, min_samples = _validate_rolling_arguments( window_size=window_size, min_samples=min_samples ) if len(self) == 0: # pragma: no cover return self return self._with_compliant( self._compliant_series.rolling_mean( window_size=window_size, min_samples=min_samples, center=center ) ) def rolling_var( self, window_size: int, *, min_samples: int | None = None, center: bool = False, ddof: int = 1, ) -> Self: """Apply a rolling variance (moving variance) over the values. A window of length `window_size` will traverse the values. The resulting values will be aggregated to their variance. The window at a given row will include the row itself and the `window_size - 1` elements before it. Arguments: window_size: The length of the window in number of elements. It must be a strictly positive integer. min_samples: The number of values in the window that should be non-null before computing a result. If set to `None` (default), it will be set equal to `window_size`. If provided, it must be a strictly positive integer, and less than or equal to `window_size`. center: Set the labels at the center of the window. ddof: Delta Degrees of Freedom; the divisor for a length N window is N - ddof. Examples: >>> import polars as pl >>> import narwhals as nw >>> >>> s_native = pl.Series([1.0, 3.0, 1.0, 4.0]) >>> nw.from_native(s_native, series_only=True).rolling_var( ... window_size=2, min_samples=1 ... ).to_native() # doctest:+NORMALIZE_WHITESPACE shape: (4,) Series: '' [f64] [ null 2.0 2.0 4.5 ] """ window_size, min_samples = _validate_rolling_arguments( window_size=window_size, min_samples=min_samples ) if len(self) == 0: # pragma: no cover return self return self._with_compliant( self._compliant_series.rolling_var( window_size=window_size, min_samples=min_samples, center=center, ddof=ddof ) ) def rolling_std( self, window_size: int, *, min_samples: int | None = None, center: bool = False, ddof: int = 1, ) -> Self: """Apply a rolling standard deviation (moving standard deviation) over the values. A window of length `window_size` will traverse the values. The resulting values will be aggregated to their standard deviation. The window at a given row will include the row itself and the `window_size - 1` elements before it. Arguments: window_size: The length of the window in number of elements. It must be a strictly positive integer. min_samples: The number of values in the window that should be non-null before computing a result. If set to `None` (default), it will be set equal to `window_size`. If provided, it must be a strictly positive integer, and less than or equal to `window_size`. center: Set the labels at the center of the window. ddof: Delta Degrees of Freedom; the divisor for a length N window is N - ddof. Examples: >>> import pandas as pd >>> import narwhals as nw >>> >>> s_native = pd.Series([1.0, 3.0, 1.0, 4.0]) >>> nw.from_native(s_native, series_only=True).rolling_std( ... window_size=2, min_samples=1 ... ).to_native() 0 NaN 1 1.414214 2 1.414214 3 2.121320 dtype: float64 """ window_size, min_samples = _validate_rolling_arguments( window_size=window_size, min_samples=min_samples ) if len(self) == 0: # pragma: no cover return self return self._with_compliant( self._compliant_series.rolling_std( window_size=window_size, min_samples=min_samples, center=center, ddof=ddof ) ) def __iter__(self) -> Iterator[Any]: yield from self._compliant_series.__iter__() def __contains__(self, other: Any) -> bool: return self._compliant_series.__contains__(other) def rank(self, method: RankMethod = "average", *, descending: bool = False) -> Self: """Assign ranks to data, dealing with ties appropriately. Notes: The resulting dtype may differ between backends. Arguments: method: The method used to assign ranks to tied elements. The following methods are available (default is 'average') - *"average"*: The average of the ranks that would have been assigned to all the tied values is assigned to each value. - *"min"*: The minimum of the ranks that would have been assigned to all the tied values is assigned to each value. (This is also referred to as "competition" ranking.) - *"max"*: The maximum of the ranks that would have been assigned to all the tied values is assigned to each value. - *"dense"*: Like "min", but the rank of the next highest element is assigned the rank immediately after those assigned to the tied elements. - *"ordinal"*: All values are given a distinct rank, corresponding to the order that the values occur in the Series. descending: Rank in descending order. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> s_native = pa.chunked_array([[3, 6, 1, 1, 6]]) >>> nw.from_native(s_native, series_only=True).rank( ... method="dense" ... ).to_native() # doctest:+ELLIPSIS [ [ 2, 3, 1, 1, 3 ] ] """ supported_rank_methods = {"average", "min", "max", "dense", "ordinal"} if method not in supported_rank_methods: msg = ( "Ranking method must be one of {'average', 'min', 'max', 'dense', 'ordinal'}. " f"Found '{method}'" ) raise ValueError(msg) return self._with_compliant( self._compliant_series.rank(method=method, descending=descending) ) def hist( self, bins: list[float] | None = None, *, bin_count: int | None = None, include_breakpoint: bool = True, ) -> DataFrame[Any]: """Bin values into buckets and count their occurrences. Warning: This functionality is considered **unstable**. It may be changed at any point without it being considered a breaking change. Arguments: bins: A monotonically increasing sequence of values. bin_count: If no bins provided, this will be used to determine the distance of the bins. include_breakpoint: Include a column that shows the intervals as categories. Examples: >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series([1, 3, 8, 8, 2, 1, 3], name="a") >>> nw.from_native(s_native, series_only=True).hist(bin_count=4) ┌────────────────────┐ | Narwhals DataFrame | |--------------------| | breakpoint count| |0 2.75 3| |1 4.50 2| |2 6.25 0| |3 8.00 2| └────────────────────┘ """ if bins is not None: if any(bins[i - 1] >= bins[i] for i in range(1, len(bins))): msg = "bins must increase monotonically" raise ComputeError(msg) if bin_count is not None: msg = f"can only provide one of `bin_count` or `bins`, got: {bin_count=}, {bins=}" raise ComputeError(msg) result = self._compliant_series.hist_from_bins( bins=bins, include_breakpoint=include_breakpoint ) else: # polars (v1.20) sets bin=10 if neither are provided. default = 10 bin_count = default if bin_count is None else bin_count result = self._compliant_series.hist_from_bin_count( bin_count=bin_count, include_breakpoint=include_breakpoint ) return self._dataframe(result, level=self._level) def log(self, base: float = math.e) -> Self: r"""Compute the logarithm to a given base. Arguments: base: Given base, defaults to `e` Examples: >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series([1, 2, 4], name="a") >>> s = nw.from_native(s_native, series_only=True) >>> s.log(base=2) ┌───────────────────────┐ | Narwhals Series | |-----------------------| |0 0.0 | |1 1.0 | |2 2.0 | |Name: a, dtype: float64| └───────────────────────┘ """ return self._with_compliant(self._compliant_series.log(base=base)) def exp(self) -> Self: r"""Compute the exponent. Examples: >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series([-1, 0, 1], name="a") >>> s = nw.from_native(s_native, series_only=True) >>> s.exp() ┌───────────────────────┐ | Narwhals Series | |-----------------------| |0 0.367879 | |1 1.000000 | |2 2.718282 | |Name: a, dtype: float64| └───────────────────────┘ """ return self._with_compliant(self._compliant_series.exp()) def sqrt(self) -> Self: r"""Compute the square root. Examples: >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series([1, 4, 9], name="a") >>> s = nw.from_native(s_native, series_only=True) >>> s.sqrt() ┌───────────────────────┐ | Narwhals Series | |-----------------------| |0 1.0 | |1 2.0 | |2 3.0 | |Name: a, dtype: float64| └───────────────────────┘ """ return self._with_compliant(self._compliant_series.sqrt()) def is_close( self, other: Self | NumericLiteral, *, abs_tol: float = 0.0, rel_tol: float = 1e-09, nans_equal: bool = False, ) -> Self: r"""Get a boolean mask of the values being close to the other values. Two values `a` and `b` are considered close if the following condition holds: $$ |a-b| \le max \{ \text{rel\_tol} \cdot max \{ |a|, |b| \}, \text{abs\_tol} \} $$ Arguments: other: Values to compare with. abs_tol: Absolute tolerance. This is the maximum allowed absolute difference between two values. Must be non-negative. rel_tol: Relative tolerance. This is the maximum allowed difference between two values, relative to the larger absolute value. Must be in the range [0, 1). nans_equal: Whether NaN values should be considered equal. Notes: The implementation of this method is symmetric and mirrors the behavior of `math.isclose`. Specifically note that this behavior is different to `numpy.isclose`. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> >>> data = [1.0, float("inf"), 1.41, None, float("nan")] >>> s_native = pa.chunked_array([data]) >>> s = nw.from_native(s_native, series_only=True) >>> s.is_close(1.4, abs_tol=0.1).to_native() # doctest:+ELLIPSIS [ [ false, false, true, null, false ] ] """ if not self.dtype.is_numeric(): msg = ( f"`is_close` operation not supported for dtype `{self.dtype}`\n\n" "Hint: `is_close` is only supported for numeric types" ) raise InvalidOperationError(msg) # Creating a temporary name if series is unnamed (possible in pandas-like scenario) # as `select`-ing a column named `None` would break otherwise orig_name = self.name name_is_none = orig_name is None tmp_name = "__nw_is_close__" if name_is_none else orig_name expr = col(tmp_name).is_close( other, abs_tol=abs_tol, rel_tol=rel_tol, nans_equal=nans_equal ) series = self.rename(tmp_name) if name_is_none else self result = series.to_frame().select(expr).get_column(tmp_name) result = result.rename(orig_name) if name_is_none else result return cast("Self", result) @property def str(self) -> SeriesStringNamespace[Self]: return SeriesStringNamespace(self) @property def dt(self) -> SeriesDateTimeNamespace[Self]: return SeriesDateTimeNamespace(self) @property def cat(self) -> SeriesCatNamespace[Self]: return SeriesCatNamespace(self) @property def list(self) -> SeriesListNamespace[Self]: return SeriesListNamespace(self) @property def struct(self) -> SeriesStructNamespace[Self]: return SeriesStructNamespace(self) narwhals-2.11.0/narwhals/series_cat.py0000644000000000000000000000150213615410400014642 0ustar00from __future__ import annotations from typing import Generic from narwhals.typing import SeriesT class SeriesCatNamespace(Generic[SeriesT]): def __init__(self, series: SeriesT) -> None: self._narwhals_series = series def get_categories(self) -> SeriesT: """Get unique categories from column. Examples: >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series(["apple", "mango", "mango"], dtype="category") >>> s = nw.from_native(s_native, series_only=True) >>> s.cat.get_categories().to_native() 0 apple 1 mango dtype: object """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.cat.get_categories() ) narwhals-2.11.0/narwhals/series_dt.py0000644000000000000000000005501313615410400014510 0ustar00from __future__ import annotations from typing import TYPE_CHECKING, Generic from narwhals.typing import SeriesT if TYPE_CHECKING: from narwhals.typing import TimeUnit class SeriesDateTimeNamespace(Generic[SeriesT]): def __init__(self, series: SeriesT) -> None: self._narwhals_series = series def date(self) -> SeriesT: """Get the date in a datetime series. Raises: NotImplementedError: If pandas default backend is being used. Examples: >>> from datetime import datetime >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series( ... [datetime(2012, 1, 7, 10, 20), datetime(2023, 3, 10, 11, 32)] ... ).convert_dtypes(dtype_backend="pyarrow") >>> s = nw.from_native(s_native, series_only=True) >>> s.dt.date().to_native() 0 2012-01-07 1 2023-03-10 dtype: date32[day][pyarrow] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.date() ) def year(self) -> SeriesT: """Get the year in a datetime series. Examples: >>> from datetime import datetime >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series([datetime(2012, 1, 7), datetime(2023, 3, 10)]) >>> s = nw.from_native(s_native, series_only=True) >>> s.dt.year().to_native() # doctest: +NORMALIZE_WHITESPACE shape: (2,) Series: '' [i32] [ 2012 2023 ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.year() ) def month(self) -> SeriesT: """Gets the month in a datetime series. Examples: >>> from datetime import datetime >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series([datetime(2012, 1, 7), datetime(2023, 3, 10)]) >>> s = nw.from_native(s_native, series_only=True) >>> s.dt.month().to_native() # doctest: +NORMALIZE_WHITESPACE shape: (2,) Series: '' [i8] [ 1 3 ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.month() ) def day(self) -> SeriesT: """Extracts the day in a datetime series. Examples: >>> from datetime import datetime >>> import pyarrow as pa >>> import narwhals as nw >>> s_native = pa.chunked_array( ... [[datetime(2022, 1, 1), datetime(2022, 1, 5)]] ... ) >>> s = nw.from_native(s_native, series_only=True) >>> s.dt.day().to_native() # doctest: +ELLIPSIS [ [ 1, 5 ] ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.day() ) def hour(self) -> SeriesT: """Extracts the hour in a datetime series. Examples: >>> from datetime import datetime >>> import pyarrow as pa >>> import narwhals as nw >>> s_native = pa.chunked_array( ... [[datetime(2022, 1, 1, 5, 3), datetime(2022, 1, 5, 9, 12)]] ... ) >>> s = nw.from_native(s_native, series_only=True) >>> s.dt.hour().to_native() # doctest: +ELLIPSIS [ [ 5, 9 ] ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.hour() ) def minute(self) -> SeriesT: """Extracts the minute in a datetime series. Examples: >>> from datetime import datetime >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series( ... [datetime(2022, 1, 1, 5, 3), datetime(2022, 1, 5, 9, 12)] ... ) >>> s = nw.from_native(s_native, series_only=True) >>> s.dt.minute().to_native() 0 3 1 12 dtype: int32 """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.minute() ) def second(self) -> SeriesT: """Extracts the seconds in a datetime series. Examples: >>> from datetime import datetime >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series( ... [datetime(2022, 1, 1, 5, 3, 10), datetime(2022, 1, 5, 9, 12, 4)] ... ) >>> s = nw.from_native(s_native, series_only=True) >>> s.dt.second().to_native() 0 10 1 4 dtype: int32 """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.second() ) def millisecond(self) -> SeriesT: """Extracts the milliseconds in a datetime series. Examples: >>> from datetime import datetime >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series( ... [ ... datetime(2022, 1, 1, 5, 3, 7, 400000), ... datetime(2022, 1, 1, 5, 3, 7, 0), ... ] ... ) >>> s = nw.from_native(s_native, series_only=True) >>> s.dt.millisecond().alias("datetime").to_native() 0 400 1 0 Name: datetime, dtype: int32 """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.millisecond() ) def microsecond(self) -> SeriesT: """Extracts the microseconds in a datetime series. Examples: >>> from datetime import datetime >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series( ... [ ... datetime(2022, 1, 1, 5, 3, 7, 400000), ... datetime(2022, 1, 1, 5, 3, 7, 0), ... ] ... ) >>> s = nw.from_native(s_native, series_only=True) >>> s.dt.microsecond().alias("datetime").to_native() 0 400000 1 0 Name: datetime, dtype: int32 """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.microsecond() ) def nanosecond(self) -> SeriesT: """Extract the nanoseconds in a date series. Examples: >>> from datetime import datetime >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series( ... [ ... datetime(2022, 1, 1, 5, 3, 7, 400000), ... datetime(2022, 1, 1, 5, 3, 7, 0), ... ] ... ) >>> s = nw.from_native(s_native, series_only=True) >>> s.dt.nanosecond().alias("datetime").to_native() 0 400000000 1 0 Name: datetime, dtype: int32 """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.nanosecond() ) def ordinal_day(self) -> SeriesT: """Get ordinal day. Examples: >>> from datetime import datetime >>> import pyarrow as pa >>> import narwhals as nw >>> s_native = pa.chunked_array( ... [[datetime(2020, 1, 1), datetime(2020, 8, 3)]] ... ) >>> s = nw.from_native(s_native, series_only=True) >>> s.dt.ordinal_day().to_native() # doctest: +ELLIPSIS [ [ 1, 216 ] ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.ordinal_day() ) def weekday(self) -> SeriesT: """Extract the week day in a datetime series. Note that Monday = 1 and Sunday = 7. Examples: >>> from datetime import datetime >>> import pyarrow as pa >>> import narwhals as nw >>> s_native = pa.chunked_array( ... [[datetime(2020, 1, 1), datetime(2020, 8, 3)]] ... ) >>> s = nw.from_native(s_native, series_only=True) >>> s.dt.weekday().to_native() # doctest: +ELLIPSIS [ [ 3, 1 ] ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.weekday() ) def total_minutes(self) -> SeriesT: """Get total minutes. Notes: The function outputs the total minutes in the int dtype by default, however, pandas may change the dtype to float when there are missing values, consider using `fill_null()` in this case. Examples: >>> from datetime import timedelta >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series( ... [timedelta(minutes=10), timedelta(minutes=20, seconds=40)] ... ) >>> s = nw.from_native(s_native, series_only=True) >>> s.dt.total_minutes().to_native() # doctest: +NORMALIZE_WHITESPACE shape: (2,) Series: '' [i64] [ 10 20 ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.total_minutes() ) def total_seconds(self) -> SeriesT: """Get total seconds. Notes: The function outputs the total seconds in the int dtype by default, however, pandas may change the dtype to float when there are missing values, consider using `fill_null()` in this case. Examples: >>> from datetime import timedelta >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series( ... [timedelta(minutes=10), timedelta(minutes=20, seconds=40)] ... ) >>> s = nw.from_native(s_native, series_only=True) >>> s.dt.total_seconds().to_native() # doctest: +NORMALIZE_WHITESPACE shape: (2,) Series: '' [i64] [ 600 1240 ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.total_seconds() ) def total_milliseconds(self) -> SeriesT: """Get total milliseconds. Notes: The function outputs the total milliseconds in the int dtype by default, however, pandas may change the dtype to float when there are missing values, consider using `fill_null()` in this case. Examples: >>> from datetime import timedelta >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series( ... [ ... timedelta(milliseconds=10), ... timedelta(milliseconds=20, microseconds=40), ... ] ... ) >>> s = nw.from_native(s_native, series_only=True) >>> s.dt.total_milliseconds().to_native() # doctest: +NORMALIZE_WHITESPACE shape: (2,) Series: '' [i64] [ 10 20 ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.total_milliseconds() ) def total_microseconds(self) -> SeriesT: """Get total microseconds. Notes: The function outputs the total microseconds in the int dtype by default, however, pandas may change the dtype to float when there are missing values, consider using `fill_null()` in this case. Examples: >>> from datetime import timedelta >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series( ... [ ... timedelta(microseconds=10), ... timedelta(milliseconds=1, microseconds=200), ... ] ... ) >>> s = nw.from_native(s_native, series_only=True) >>> s.dt.total_microseconds().to_native() # doctest: +NORMALIZE_WHITESPACE shape: (2,) Series: '' [i64] [ 10 1200 ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.total_microseconds() ) def total_nanoseconds(self) -> SeriesT: """Get total nanoseconds. Notes: The function outputs the total nanoseconds in the int dtype by default, however, pandas may change the dtype to float when there are missing values, consider using `fill_null()` in this case. Examples: >>> from datetime import datetime >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series( ... ["2024-01-01 00:00:00.000000001", "2024-01-01 00:00:00.000000002"] ... ).str.to_datetime(time_unit="ns") >>> s = nw.from_native(s_native, series_only=True) >>> s.diff().dt.total_nanoseconds().to_native() # doctest: +NORMALIZE_WHITESPACE shape: (2,) Series: '' [i64] [ null 1 ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.total_nanoseconds() ) def to_string(self, format: str) -> SeriesT: """Convert a Date/Time/Datetime series into a String series with the given format. Arguments: format: Format string for converting the datetime to string. Notes: Unfortunately, different libraries interpret format directives a bit differently. - Chrono, the library used by Polars, uses `"%.f"` for fractional seconds, whereas pandas and Python stdlib use `".%f"`. - PyArrow interprets `"%S"` as "seconds, including fractional seconds" whereas most other tools interpret it as "just seconds, as 2 digits". --- Therefore, we make the following adjustments. - for pandas-like libraries, we replace `"%S.%f"` with `"%S%.f"`. - for PyArrow, we replace `"%S.%f"` with `"%S"`. --- Workarounds like these don't make us happy, and we try to avoid them as much as possible, but here we feel like it's the best compromise. If you just want to format a date/datetime Series as a local datetime string, and have it work as consistently as possible across libraries, we suggest using: - `"%Y-%m-%dT%H:%M:%S%.f"` for datetimes - `"%Y-%m-%d"` for dates --- Though note that, even then, different tools may return a different number of trailing zeros. Nonetheless, this is probably consistent enough for most applications. If you have an application where this is not enough, please open an issue and let us know. Examples: >>> from datetime import datetime >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series([datetime(2020, 3, 1), datetime(2020, 4, 1)]) >>> s = nw.from_native(s_native, series_only=True) >>> s.dt.to_string("%Y/%m/%d") ┌───────────────┐ |Narwhals Series| |---------------| |0 2020/03/01| |1 2020/04/01| |dtype: object | └───────────────┘ """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.to_string(format) ) def replace_time_zone(self, time_zone: str | None) -> SeriesT: """Replace time zone. Arguments: time_zone: Target time zone. Examples: >>> from datetime import datetime, timezone >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series( ... [ ... datetime(2024, 1, 1, tzinfo=timezone.utc), ... datetime(2024, 1, 2, tzinfo=timezone.utc), ... ] ... ) >>> s = nw.from_native(s_native, series_only=True) >>> s.dt.replace_time_zone( ... "Asia/Kathmandu" ... ).to_native() # doctest: +NORMALIZE_WHITESPACE shape: (2,) Series: '' [datetime[Ξs, Asia/Kathmandu]] [ 2024-01-01 00:00:00 +0545 2024-01-02 00:00:00 +0545 ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.replace_time_zone(time_zone) ) def convert_time_zone(self, time_zone: str) -> SeriesT: """Convert time zone. If converting from a time-zone-naive column, then conversion happens as if converting from UTC. Arguments: time_zone: Target time zone. Examples: >>> from datetime import datetime, timezone >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series( ... [ ... datetime(2024, 1, 1, tzinfo=timezone.utc), ... datetime(2024, 1, 2, tzinfo=timezone.utc), ... ] ... ) >>> s = nw.from_native(s_native, series_only=True) >>> s.dt.convert_time_zone("Asia/Kathmandu").to_native() 0 2024-01-01 05:45:00+05:45 1 2024-01-02 05:45:00+05:45 dtype: datetime64[ns, Asia/Kathmandu] """ if time_zone is None: msg = "Target `time_zone` cannot be `None` in `convert_time_zone`. Please use `replace_time_zone(None)` if you want to remove the time zone." raise TypeError(msg) return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.convert_time_zone(time_zone) ) def timestamp(self, time_unit: TimeUnit) -> SeriesT: """Return a timestamp in the given time unit. Arguments: time_unit: One of - 'ns': nanosecond. - 'us': microsecond. - 'ms': millisecond. Examples: >>> from datetime import date >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series( ... [date(2001, 1, 1), None, date(2001, 1, 3)], dtype="datetime64[ns]" ... ) >>> s = nw.from_native(s_native, series_only=True) >>> s.dt.timestamp("ms").to_native() 0 9.783072e+11 1 NaN 2 9.784800e+11 dtype: float64 """ if time_unit not in {"ns", "us", "ms"}: msg = ( "invalid `time_unit`" f"\n\nExpected one of {{'ns', 'us', 'ms'}}, got {time_unit!r}." ) raise ValueError(msg) return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.timestamp(time_unit) ) def truncate(self, every: str) -> SeriesT: """Divide the date/datetime range into buckets. Arguments: every: Length of bucket. Must be of form ``, where `multiple` is a positive integer and `unit` is one of - 'ns': nanosecond. - 'us': microsecond. - 'ms': millisecond. - 's': second. - 'm': minute. - 'h': hour. - 'd': day. - 'mo': month. - 'q': quarter. - 'y': year. Examples: >>> from datetime import datetime >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series([datetime(2021, 3, 1, 12, 34)]) >>> s = nw.from_native(s_native, series_only=True) >>> s.dt.truncate("1h").to_native() 0 2021-03-01 12:00:00 dtype: datetime64[ns] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.truncate(every) ) def offset_by(self, by: str) -> SeriesT: """Offset this date by a relative time offset. Arguments: by: The offset. Must be of form ``, where `multiple` is a positive integer and `unit` is one of - 'ns': nanosecond. - 'us': microsecond. - 'ms': millisecond. - 's': second. - 'm': minute. - 'h': hour. - 'd': day. - 'mo': month. - 'q': quarter. - 'y': year. Examples: >>> from datetime import datetime >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series([datetime(2021, 3, 1, 12, 34)]) >>> s = nw.from_native(s_native, series_only=True) >>> s.dt.offset_by("1h").to_native() 0 2021-03-01 13:34:00 dtype: datetime64[ns] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.dt.offset_by(by) ) narwhals-2.11.0/narwhals/series_list.py0000644000000000000000000000725613615410400015062 0ustar00from __future__ import annotations from typing import TYPE_CHECKING, Generic from narwhals.typing import SeriesT if TYPE_CHECKING: from narwhals.typing import NonNestedLiteral class SeriesListNamespace(Generic[SeriesT]): def __init__(self, series: SeriesT) -> None: self._narwhals_series = series def len(self) -> SeriesT: """Return the number of elements in each list. Null values count towards the total. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> s_native = pa.chunked_array([[[1, 2], [3, 4, None], None, []]]) >>> s = nw.from_native(s_native, series_only=True) >>> s.list.len().to_native() # doctest: +ELLIPSIS [ [ 2, 3, null, 0 ] ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.list.len() ) def unique(self) -> SeriesT: """Get the unique/distinct values in the list. Null values are included in the result. The order of unique values is not guaranteed. Examples: >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series([[1, 1, 2], [3, 3, None], None, []]) >>> s = nw.from_native(s_native, series_only=True) >>> s.list.unique().to_native() # doctest: +NORMALIZE_WHITESPACE shape: (4,) Series: '' [list[i64]] [ [1, 2] [null, 3] null [] ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.list.unique() ) def contains(self, item: NonNestedLiteral) -> SeriesT: """Check if sublists contain the given item. Arguments: item: Item that will be checked for membership. Examples: >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series([[1, 2], None, []]) >>> s = nw.from_native(s_native, series_only=True) >>> s.list.contains(1).to_native() # doctest: +NORMALIZE_WHITESPACE shape: (3,) Series: '' [bool] [ true null false ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.list.contains(item) ) def get(self, index: int) -> SeriesT: """Return the value by index in each list. Negative indices are not accepted. Examples: >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series([[1, 2], [3, 4, None], [None, 5]]) >>> s = nw.from_native(s_native, series_only=True) >>> s.list.get(1).to_native() # doctest: +NORMALIZE_WHITESPACE shape: (3,) Series: '' [i64] [ 2 4 5 ] """ if not isinstance(index, int): msg = ( f"Index must be of type 'int'. Got type '{type(index).__name__}' instead." ) raise TypeError(msg) if index < 0: msg = f"Index {index} is out of bounds: should be greater than or equal to 0." raise ValueError(msg) return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.list.get(index) ) narwhals-2.11.0/narwhals/series_str.py0000644000000000000000000004060513615410400014712 0ustar00from __future__ import annotations from typing import Any, Generic from narwhals.dependencies import is_narwhals_series from narwhals.typing import SeriesT class SeriesStringNamespace(Generic[SeriesT]): def __init__(self, series: SeriesT) -> None: self._narwhals_series = series def len_chars(self) -> SeriesT: r"""Return the length of each string as the number of characters. Examples: >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series(["foo", "345", None]) >>> s = nw.from_native(s_native, series_only=True) >>> s.str.len_chars().to_native() # doctest: +NORMALIZE_WHITESPACE shape: (3,) Series: '' [u32] [ 3 3 null ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.str.len_chars() ) def _extract_compliant(self, arg: Any) -> Any: return arg._compliant_series if is_narwhals_series(arg) else arg def replace( self, pattern: str, value: str | SeriesT, *, literal: bool = False, n: int = 1 ) -> SeriesT: r"""Replace first matching regex/literal substring with a new string value. Arguments: pattern: A valid regular expression pattern. value: String that will replace the matched substring. literal: Treat `pattern` as a literal string. n: Number of matches to replace. Examples: >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series(["123abc", "abc abc123"]) >>> s = nw.from_native(s_native, series_only=True) >>> s.str.replace("abc", "").to_native() 0 123 1 abc123 dtype: object """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.str.replace( self._extract_compliant(value), pattern=pattern, literal=literal, n=n ) ) def replace_all( self, pattern: str, value: str | SeriesT, *, literal: bool = False ) -> SeriesT: r"""Replace all matching regex/literal substring with a new string value. Arguments: pattern: A valid regular expression pattern. value: String that will replace the matched substring. literal: Treat `pattern` as a literal string. Examples: >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series(["123abc", "abc abc123"]) >>> s = nw.from_native(s_native, series_only=True) >>> s.str.replace_all("abc", "").to_native() 0 123 1 123 dtype: object """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.str.replace_all( self._extract_compliant(value), pattern, literal=literal ) ) def strip_chars(self, characters: str | None = None) -> SeriesT: r"""Remove leading and trailing characters. Arguments: characters: The set of characters to be removed. All combinations of this set of characters will be stripped from the start and end of the string. If set to None (default), all leading and trailing whitespace is removed instead. Examples: >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series(["apple", "\nmango"]) >>> s = nw.from_native(s_native, series_only=True) >>> s.str.strip_chars().to_native() # doctest: +NORMALIZE_WHITESPACE shape: (2,) Series: '' [str] [ "apple" "mango" ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.str.strip_chars(characters) ) def starts_with(self, prefix: str) -> SeriesT: r"""Check if string values start with a substring. Arguments: prefix: prefix substring Examples: >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series(["apple", "mango", None]) >>> s = nw.from_native(s_native, series_only=True) >>> s.str.starts_with("app").to_native() 0 True 1 False 2 None dtype: object """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.str.starts_with(prefix) ) def ends_with(self, suffix: str) -> SeriesT: r"""Check if string values end with a substring. Arguments: suffix: suffix substring Examples: >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series(["apple", "mango", None]) >>> s = nw.from_native(s_native, series_only=True) >>> s.str.ends_with("ngo").to_native() 0 False 1 True 2 None dtype: object """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.str.ends_with(suffix) ) def contains(self, pattern: str, *, literal: bool = False) -> SeriesT: r"""Check if string contains a substring that matches a pattern. Arguments: pattern: A Character sequence or valid regular expression pattern. literal: If True, treats the pattern as a literal string. If False, assumes the pattern is a regular expression. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> s_native = pa.chunked_array([["cat", "dog", "rabbit and parrot"]]) >>> s = nw.from_native(s_native, series_only=True) >>> s.str.contains("cat|parrot").to_native() # doctest: +ELLIPSIS [ [ true, false, true ] ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.str.contains(pattern, literal=literal) ) def slice(self, offset: int, length: int | None = None) -> SeriesT: r"""Create subslices of the string values of a Series. Arguments: offset: Start index. Negative indexing is supported. length: Length of the slice. If set to `None` (default), the slice is taken to the end of the string. Examples: >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series(["pear", None, "papaya"]) >>> s = nw.from_native(s_native, series_only=True) >>> s.str.slice(4, 3).to_native() # doctest: +NORMALIZE_WHITESPACE 0 1 None 2 ya dtype: object """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.str.slice( offset=offset, length=length ) ) def split(self, by: str) -> SeriesT: r"""Split the string values of a Series by a substring. Arguments: by: Substring to split by. Examples: >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series(["foo bar", "foo_bar"]) >>> s = nw.from_native(s_native, series_only=True) >>> s.str.split("_").to_native() # doctest: +NORMALIZE_WHITESPACE shape: (2,) Series: '' [list[str]] [ ["foo bar"] ["foo", "bar"] ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.str.split(by=by) ) def head(self, n: int = 5) -> SeriesT: r"""Take the first n elements of each string. Arguments: n: Number of elements to take. Negative indexing is supported (see note (1.)) Notes: 1. When the `n` input is negative, `head` returns characters up to the n-th from the end of the string. For example, if `n = -3`, then all characters except the last three are returned. 2. If the length of the string has fewer than `n` characters, the full string is returned. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> s_native = pa.chunked_array([["taata", "taatatata", "zukkyun"]]) >>> s = nw.from_native(s_native, series_only=True) >>> s.str.head().to_native() # doctest: +ELLIPSIS [ [ "taata", "taata", "zukky" ] ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.str.slice(offset=0, length=n) ) def tail(self, n: int = 5) -> SeriesT: r"""Take the last n elements of each string. Arguments: n: Number of elements to take. Negative indexing is supported (see note (1.)) Notes: 1. When the `n` input is negative, `tail` returns characters starting from the n-th from the beginning of the string. For example, if `n = -3`, then all characters except the first three are returned. 2. If the length of the string has fewer than `n` characters, the full string is returned. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> s_native = pa.chunked_array([["taata", "taatatata", "zukkyun"]]) >>> s = nw.from_native(s_native, series_only=True) >>> s.str.tail().to_native() # doctest: +ELLIPSIS [ [ "taata", "atata", "kkyun" ] ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.str.slice(offset=-n, length=None) ) def to_uppercase(self) -> SeriesT: r"""Transform string to uppercase variant. Notes: The PyArrow backend will convert 'ß' to '᚞' instead of 'SS'. For more info see: https://github.com/apache/arrow/issues/34599 There may be other unicode-edge-case-related variations across implementations. Examples: >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series(["apple", None]) >>> s = nw.from_native(s_native, series_only=True) >>> s.str.to_uppercase().to_native() 0 APPLE 1 None dtype: object """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.str.to_uppercase() ) def to_lowercase(self) -> SeriesT: r"""Transform string to lowercase variant. Examples: >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series(["APPLE", None]) >>> s = nw.from_native(s_native, series_only=True) >>> s.str.to_lowercase().to_native() 0 apple 1 None dtype: object """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.str.to_lowercase() ) def to_datetime(self, format: str | None = None) -> SeriesT: """Parse Series with strings to a Series with Datetime dtype. Notes: - pandas defaults to nanosecond time unit, Polars to microsecond. Prior to pandas 2.0, nanoseconds were the only time unit supported in pandas, with no ability to set any other one. The ability to set the time unit in pandas, if the version permits, will arrive. - timezone-aware strings are all converted to and parsed as UTC. Warning: As different backends auto-infer format in different ways, if `format=None` there is no guarantee that the result will be equal. Arguments: format: Format to use for conversion. If set to None (default), the format is inferred from the data. Examples: >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series(["2020-01-01", "2020-01-02"]) >>> s = nw.from_native(s_native, series_only=True) >>> s.str.to_datetime( ... format="%Y-%m-%d" ... ).to_native() # doctest: +NORMALIZE_WHITESPACE shape: (2,) Series: '' [datetime[Ξs]] [ 2020-01-01 00:00:00 2020-01-02 00:00:00 ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.str.to_datetime(format=format) ) def to_date(self, format: str | None = None) -> SeriesT: """Convert to date dtype. Warning: As different backends auto-infer format in different ways, if `format=None` there is no guarantee that the result will be equal. Arguments: format: Format to use for conversion. If set to None (default), the format is inferred from the data. Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> s_native = pa.chunked_array([["2020-01-01", "2020-01-02"]]) >>> s = nw.from_native(s_native, series_only=True) >>> s.str.to_date(format="%Y-%m-%d").to_native() # doctest: +ELLIPSIS [ [ 2020-01-01, 2020-01-02 ] ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.str.to_date(format=format) ) def to_titlecase(self) -> SeriesT: """Modify strings to their titlecase equivalent. Notes: This is a form of case transform where the first letter of each word is capitalized, with the rest of the word in lowercase. Word boundaries are defined by non-**alphabetic** characters, matching the behavior of [`str.title`](https://docs.python.org/3/library/stdtypes.html#str.title) Examples: >>> import pyarrow as pa >>> import narwhals as nw >>> s_native = pa.chunked_array( ... [ ... [ ... "'e.t. phone home'", ... "you talkin' to me?", ... "to infinity,and BEYOND!", ... ] ... ] ... ) >>> s = nw.from_native(s_native, series_only=True) >>> s.str.to_titlecase().to_native() # doctest: +ELLIPSIS [ [ "'E.T. Phone Home'", "You Talkin' To Me?", "To Infinity,And Beyond!" ] ] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.str.to_titlecase() ) def zfill(self, width: int) -> SeriesT: r"""Pad strings with zeros on the left. Arguments: width: The target width of the string. If the string is shorter than this width, it will be padded with zeros on the left. Examples: >>> import pandas as pd >>> import narwhals as nw >>> s_native = pd.Series(["+1", "-23", "456", "123456"]) >>> s = nw.from_native(s_native, series_only=True) >>> s.str.zfill(5).to_native() 0 +0001 1 -0023 2 00456 3 123456 dtype: object """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.str.zfill(width) ) narwhals-2.11.0/narwhals/series_struct.py0000644000000000000000000000164213615410400015424 0ustar00from __future__ import annotations from typing import Generic from narwhals.typing import SeriesT class SeriesStructNamespace(Generic[SeriesT]): def __init__(self, series: SeriesT) -> None: self._narwhals_series = series def field(self, name: str) -> SeriesT: r"""Retrieve a Struct field as a new expression. Arguments: name: Name of the struct field to retrieve. Examples: >>> import polars as pl >>> import narwhals as nw >>> s_native = pl.Series( ... [{"id": "0", "name": "john"}, {"id": "1", "name": "jane"}] ... ) >>> s = nw.from_native(s_native, series_only=True) >>> s.struct.field("name").to_list() ['john', 'jane'] """ return self._narwhals_series._with_compliant( self._narwhals_series._compliant_series.struct.field(name) ) narwhals-2.11.0/narwhals/this.py0000644000000000000000000000306013615410400013471 0ustar00# ruff: noqa ZEN = """\ âĢŋâĢŋâĢŋâĢŋâĢŋ⠘âĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋ THE ZEN OF NARWHALS âĢŋâĢŋâĢŋâĢŋâĢŋâ  âĒđâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋ Keep it simple âĢŋâĢŋâĢŋâĢŋâĢŋ⡀⡄âĒŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋ Move slowly and keep things working âĢŋâĢŋâĢŋâĢŋâĢŋ⡇⡾⡘⠛â ŋâ ŋâ ŋâĒŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋ A good API is an honest one âĢŋâĢŋâĢŋâĄŋâĢŦ⡄â ūâĢĢâ đâĢŋâĢŋâĢŋâĢķâĢŪâĢ™â ŧâĒŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋ Yes, that needs documenting âĢŋâĢŋâĒ‹âĢīâĢŋâĢ·âĢŽâĢ­âĢūâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢĶ⡙âĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋ People learn better from examples âĢŋâĒƒâĢūâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢ·âĄŒâĒŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋ than from explanations⠀ ⡏⠀âĒ°â „âĒŧâĢŋâĢŋâĢŋâĢŋâĄŋ⠋âĒ‰â ŧâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋ⡜âĢŋâĢŋâĄŋâĒâĒŧâĢŋâĢŋâĢŋâĢŋâĢŋ If in doubt, better to say 'no' ⡇âĢŒâĢ€âĢ âĢūâĢŋâĢŋâĢŋâĢŋâĢ‡â ķ⠉âĒâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢ§âĄđâĢŋ⡇âĢŋâĢ§â ŧâ ŋâ ŋâ ŋâ ŋ than to risk causing a commotion⠀ ⥧âĒđâĢŋâĢŋâĢŋâĢœâĢŸâĢļâĢŋâĢŋâĢ·âĢķâĢŋâĄŋâĢŋâĢŋâĢâĒŋâĢŋâĢŋâĢ·âĢŽâĢĨâĢŋâĢŋâĢŋâĢŋâĢŋ⡟âĢ° Yes, we need a test for that âĒĄâĢ†âĒŧâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢ§âĄ™âĢŋâĢŋ⡇âĢŋâĢŋâĢŋâĢŋ⠟âĢ‹âĢ­âĢ›â ŧâĢ‹âĢīâĢŋ If you want users âĢķâĢĪâĢĪâĢ™â ŧâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢ·âĢĶâĢâĢĄâĢŋâĄŋâĒ‹âĢīâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋ you need good docs⠀ âĢŋâĢŋâĢŋâĢŋâĢŋâĢķâĢŽâĢ™âĢ›â ŧâ ŋâ ŋâ ŋâ ŋâ ŋ⠟âĢ›âĢĐâĢĨâĢķâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋâĢŋ Our code is not irreplaceable""" print(ZEN) narwhals-2.11.0/narwhals/translate.py0000644000000000000000000006257013615410400014532 0ustar00from __future__ import annotations import datetime as dt from decimal import Decimal from functools import wraps from typing import TYPE_CHECKING, Any, Callable, Literal, TypeVar, overload from narwhals import plugins from narwhals._constants import EPOCH, MS_PER_SECOND from narwhals._native import ( is_native_arrow, is_native_pandas_like, is_native_polars, is_native_spark_like, ) from narwhals._utils import ( Implementation, Version, has_native_namespace, is_compliant_dataframe, is_compliant_lazyframe, is_compliant_series, ) from narwhals.dependencies import ( get_dask_expr, get_numpy, get_pandas, is_cupy_scalar, is_dask_dataframe, is_duckdb_relation, is_ibis_table, is_numpy_scalar, is_pandas_like_dataframe, is_polars_lazyframe, is_polars_series, is_pyarrow_scalar, is_pyarrow_table, ) if TYPE_CHECKING: from narwhals.dataframe import DataFrame, LazyFrame from narwhals.series import Series from narwhals.typing import ( DataFrameT, Frame, IntoDataFrameT, IntoFrame, IntoLazyFrameT, IntoSeries, IntoSeriesT, LazyFrameT, SeriesT, ) T = TypeVar("T") NON_TEMPORAL_SCALAR_TYPES = (bool, bytes, str, int, float, complex, Decimal) TEMPORAL_SCALAR_TYPES = (dt.date, dt.timedelta, dt.time) @overload def to_native( narwhals_object: DataFrame[IntoDataFrameT], *, pass_through: Literal[False] = ... ) -> IntoDataFrameT: ... @overload def to_native( narwhals_object: LazyFrame[IntoLazyFrameT], *, pass_through: Literal[False] = ... ) -> IntoLazyFrameT: ... @overload def to_native( narwhals_object: Series[IntoSeriesT], *, pass_through: Literal[False] = ... ) -> IntoSeriesT: ... @overload def to_native(narwhals_object: Any, *, pass_through: bool) -> Any: ... def to_native( narwhals_object: DataFrame[IntoDataFrameT] | LazyFrame[IntoLazyFrameT] | Series[IntoSeriesT], *, pass_through: bool = False, ) -> IntoDataFrameT | IntoLazyFrameT | IntoSeriesT | Any: """Convert Narwhals object to native one. Arguments: narwhals_object: Narwhals object. pass_through: Determine what happens if `narwhals_object` isn't a Narwhals class - `False` (default): raise an error - `True`: pass object through as-is Returns: Object of class that user started with. """ from narwhals.dataframe import BaseFrame from narwhals.series import Series if isinstance(narwhals_object, BaseFrame): return narwhals_object._compliant_frame._native_frame if isinstance(narwhals_object, Series): return narwhals_object._compliant_series.native if not pass_through: msg = f"Expected Narwhals object, got {type(narwhals_object)}." raise TypeError(msg) return narwhals_object @overload def from_native(native_object: SeriesT, **kwds: Any) -> SeriesT: ... @overload def from_native(native_object: DataFrameT, **kwds: Any) -> DataFrameT: ... @overload def from_native(native_object: LazyFrameT, **kwds: Any) -> LazyFrameT: ... @overload def from_native( native_object: IntoDataFrameT | IntoSeriesT, *, pass_through: Literal[True], eager_only: Literal[True], series_only: Literal[False] = ..., allow_series: Literal[True], ) -> DataFrame[IntoDataFrameT] | Series[IntoSeriesT]: ... @overload def from_native( native_object: IntoDataFrameT, *, pass_through: Literal[True], eager_only: Literal[False] = ..., series_only: Literal[False] = ..., allow_series: None = ..., ) -> DataFrame[IntoDataFrameT]: ... @overload def from_native( native_object: T, *, pass_through: Literal[True], eager_only: Literal[False] = ..., series_only: Literal[False] = ..., allow_series: None = ..., ) -> T: ... @overload def from_native( native_object: IntoDataFrameT, *, pass_through: Literal[True], eager_only: Literal[True], series_only: Literal[False] = ..., allow_series: None = ..., ) -> DataFrame[IntoDataFrameT]: ... @overload def from_native( native_object: T, *, pass_through: Literal[True], eager_only: Literal[True], series_only: Literal[False] = ..., allow_series: None = ..., ) -> T: ... @overload def from_native( native_object: IntoDataFrameT | IntoLazyFrameT | IntoSeriesT, *, pass_through: Literal[True], eager_only: Literal[False] = ..., series_only: Literal[False] = ..., allow_series: Literal[True], ) -> DataFrame[IntoDataFrameT] | LazyFrame[IntoLazyFrameT] | Series[IntoSeriesT]: ... @overload def from_native( native_object: IntoSeriesT, *, pass_through: Literal[True], eager_only: Literal[False] = ..., series_only: Literal[True], allow_series: None = ..., ) -> Series[IntoSeriesT]: ... @overload def from_native( native_object: IntoLazyFrameT, *, pass_through: Literal[False] = ..., eager_only: Literal[False] = ..., series_only: Literal[False] = ..., allow_series: None = ..., ) -> LazyFrame[IntoLazyFrameT]: ... @overload def from_native( native_object: IntoDataFrameT, *, pass_through: Literal[False] = ..., eager_only: Literal[False] = ..., series_only: Literal[False] = ..., allow_series: None = ..., ) -> DataFrame[IntoDataFrameT]: ... @overload def from_native( native_object: IntoDataFrameT, *, pass_through: Literal[False] = ..., eager_only: Literal[True], series_only: Literal[False] = ..., allow_series: None = ..., ) -> DataFrame[IntoDataFrameT]: ... @overload def from_native( native_object: IntoFrame | IntoSeries, *, pass_through: Literal[False] = ..., eager_only: Literal[False] = ..., series_only: Literal[False] = ..., allow_series: Literal[True], ) -> DataFrame[Any] | LazyFrame[Any] | Series[Any]: ... @overload def from_native( native_object: IntoSeriesT, *, pass_through: Literal[False] = ..., eager_only: Literal[False] = ..., series_only: Literal[True], allow_series: None = ..., ) -> Series[IntoSeriesT]: ... # All params passed in as variables @overload def from_native( native_object: Any, *, pass_through: bool, eager_only: bool, series_only: bool, allow_series: bool | None, ) -> Any: ... def from_native( # noqa: D417 native_object: IntoLazyFrameT | IntoDataFrameT | IntoSeriesT | IntoFrame | IntoSeries | T, *, pass_through: bool = False, eager_only: bool = False, series_only: bool = False, allow_series: bool | None = None, **kwds: Any, ) -> LazyFrame[IntoLazyFrameT] | DataFrame[IntoDataFrameT] | Series[IntoSeriesT] | T: """Convert `native_object` to Narwhals Dataframe, Lazyframe, or Series. Arguments: native_object: Raw object from user. Depending on the other arguments, input object can be - a Dataframe / Lazyframe / Series supported by Narwhals (pandas, Polars, PyArrow, ...) - an object which implements `__narwhals_dataframe__`, `__narwhals_lazyframe__`, or `__narwhals_series__` pass_through: Determine what happens if the object can't be converted to Narwhals - `False` (default): raise an error - `True`: pass object through as-is eager_only: Whether to only allow eager objects - `False` (default): don't require `native_object` to be eager - `True`: only convert to Narwhals if `native_object` is eager series_only: Whether to only allow Series - `False` (default): don't require `native_object` to be a Series - `True`: only convert to Narwhals if `native_object` is a Series allow_series: Whether to allow Series (default is only Dataframe / Lazyframe) - `False` or `None` (default): don't convert to Narwhals if `native_object` is a Series - `True`: allow `native_object` to be a Series Returns: DataFrame, LazyFrame, Series, or original object, depending on which combination of parameters was passed. """ if kwds: msg = f"from_native() got an unexpected keyword argument {next(iter(kwds))!r}" raise TypeError(msg) return _from_native_impl( # type: ignore[no-any-return] native_object, pass_through=pass_through, eager_only=eager_only, eager_or_interchange_only=False, series_only=series_only, allow_series=allow_series, version=Version.MAIN, ) def _translate_if_compliant( # noqa: C901,PLR0911 compliant_object: Any, *, pass_through: bool = False, eager_only: bool = False, # Interchange-level was removed after v1 eager_or_interchange_only: bool, series_only: bool, allow_series: bool | None, version: Version, ) -> Any: if is_compliant_dataframe(compliant_object): if series_only: if not pass_through: msg = "Cannot only use `series_only` with dataframe" raise TypeError(msg) return compliant_object return version.dataframe( compliant_object.__narwhals_dataframe__()._with_version(version), level="full" ) if is_compliant_lazyframe(compliant_object): if series_only: if not pass_through: msg = "Cannot only use `series_only` with lazyframe" raise TypeError(msg) return compliant_object if eager_only or eager_or_interchange_only: if not pass_through: msg = "Cannot only use `eager_only` or `eager_or_interchange_only` with lazyframe" raise TypeError(msg) return compliant_object return version.lazyframe( compliant_object.__narwhals_lazyframe__()._with_version(version), level="full" ) if is_compliant_series(compliant_object): if not allow_series: if not pass_through: msg = "Please set `allow_series=True` or `series_only=True`" raise TypeError(msg) return compliant_object return version.series( compliant_object.__narwhals_series__()._with_version(version), level="full" ) # Object wasn't compliant, can't translate here. return None def _from_native_impl( # noqa: C901, PLR0911, PLR0912, PLR0915 native_object: Any, *, pass_through: bool = False, eager_only: bool = False, # Interchange-level was removed after v1 eager_or_interchange_only: bool, series_only: bool, allow_series: bool | None, version: Version, ) -> Any: from narwhals._interchange.dataframe import supports_dataframe_interchange from narwhals.dataframe import DataFrame, LazyFrame from narwhals.series import Series # Early returns if isinstance(native_object, (DataFrame, LazyFrame)) and not series_only: return native_object if isinstance(native_object, Series) and (series_only or allow_series): return native_object if series_only: if allow_series is False: msg = "Invalid parameter combination: `series_only=True` and `allow_series=False`" raise ValueError(msg) allow_series = True if eager_only and eager_or_interchange_only: msg = "Invalid parameter combination: `eager_only=True` and `eager_or_interchange_only=True`" raise ValueError(msg) # Extensions if ( translated := _translate_if_compliant( native_object, pass_through=pass_through, eager_only=eager_only, eager_or_interchange_only=eager_or_interchange_only, series_only=series_only, allow_series=allow_series, version=version, ) ) is not None: return translated # Polars if is_native_polars(native_object): if series_only and not is_polars_series(native_object): if not pass_through: msg = f"Cannot only use `series_only` with {type(native_object).__qualname__}" raise TypeError(msg) return native_object if (eager_only or eager_or_interchange_only) and is_polars_lazyframe( native_object ): if not pass_through: msg = "Cannot only use `eager_only` or `eager_or_interchange_only` with polars.LazyFrame" raise TypeError(msg) return native_object if (not allow_series) and is_polars_series(native_object): if not pass_through: msg = "Please set `allow_series=True` or `series_only=True`" raise TypeError(msg) return native_object return ( version.namespace.from_native_object(native_object) .compliant.from_native(native_object) .to_narwhals() ) # PandasLike if is_native_pandas_like(native_object): if is_pandas_like_dataframe(native_object): if series_only: if not pass_through: msg = f"Cannot only use `series_only` with {type(native_object).__qualname__}" raise TypeError(msg) return native_object elif not allow_series: if not pass_through: msg = "Please set `allow_series=True` or `series_only=True`" raise TypeError(msg) return native_object return ( version.namespace.from_native_object(native_object) .compliant.from_native(native_object) .to_narwhals() ) # PyArrow if is_native_arrow(native_object): if is_pyarrow_table(native_object): if series_only: if not pass_through: msg = f"Cannot only use `series_only` with {type(native_object).__qualname__}" raise TypeError(msg) return native_object elif not allow_series: if not pass_through: msg = "Please set `allow_series=True` or `series_only=True`" raise TypeError(msg) return native_object return ( version.namespace.from_native_object(native_object) .compliant.from_native(native_object) .to_narwhals() ) # Dask if is_dask_dataframe(native_object): if series_only: if not pass_through: msg = "Cannot only use `series_only` with dask DataFrame" raise TypeError(msg) return native_object if eager_only or eager_or_interchange_only: if not pass_through: msg = "Cannot only use `eager_only` or `eager_or_interchange_only` with dask DataFrame" raise TypeError(msg) return native_object if ( Implementation.DASK._backend_version() <= (2024, 12, 1) and get_dask_expr() is None ): # pragma: no cover msg = "Please install dask-expr" raise ImportError(msg) return ( version.namespace.from_backend(Implementation.DASK) .compliant.from_native(native_object) .to_narwhals() ) # DuckDB if is_duckdb_relation(native_object): if eager_only or series_only: # pragma: no cover if not pass_through: msg = "Cannot only use `series_only=True` or `eager_only=False` with DuckDBPyRelation" raise TypeError(msg) return native_object return ( version.namespace.from_native_object(native_object) .compliant.from_native(native_object) .to_narwhals() ) # Ibis if is_ibis_table(native_object): if eager_only or series_only: # pragma: no cover if not pass_through: msg = "Cannot only use `series_only=True` or `eager_only=False` with ibis.Table" raise TypeError(msg) return native_object return ( version.namespace.from_native_object(native_object) .compliant.from_native(native_object) .to_narwhals() ) # PySpark if is_native_spark_like(native_object): # pragma: no cover ns_spark = version.namespace.from_native_object(native_object) if series_only or eager_only or eager_or_interchange_only: if not pass_through: msg = ( "Cannot only use `series_only`, `eager_only` or `eager_or_interchange_only` " f"with {ns_spark.implementation} DataFrame" ) raise TypeError(msg) return native_object return ns_spark.compliant.from_native(native_object).to_narwhals() # Interchange protocol if supports_dataframe_interchange(native_object): from narwhals._interchange.dataframe import InterchangeFrame if eager_only or series_only: if not pass_through: msg = ( "Cannot only use `series_only=True` or `eager_only=False` " "with object which only implements __dataframe__" ) raise TypeError(msg) return native_object if version is not Version.V1: if pass_through: return native_object msg = ( "The Dataframe Interchange Protocol is no longer supported in the main `narwhals` namespace.\n\n" "You may want to:\n" " - Use `narwhals.stable.v1`, where it is still supported.\n" " - See https://narwhals-dev.github.io/narwhals/backcompat\n" " - Use `pass_through=True` to pass the object through without raising." ) raise TypeError(msg) return Version.V1.dataframe(InterchangeFrame(native_object), level="interchange") compliant_object = plugins.from_native(native_object, version) if compliant_object is not None: return _translate_if_compliant( compliant_object, pass_through=pass_through, eager_only=eager_only, eager_or_interchange_only=eager_or_interchange_only, series_only=series_only, allow_series=allow_series, version=version, ) if not pass_through: msg = f"Unsupported dataframe type, got: {type(native_object)}" if hint := plugins._show_suggestions(type(native_object)): # pragma: no cover msg += "\n\n" + hint raise TypeError(msg) return native_object def get_native_namespace(*obj: Frame | Series[Any] | IntoFrame | IntoSeries) -> Any: """Get native namespace from object. Arguments: obj: Dataframe, Lazyframe, or Series. Multiple objects can be passed positionally, in which case they must all have the same native namespace (else an error is raised). Returns: Native module. Examples: >>> import polars as pl >>> import pandas as pd >>> import narwhals as nw >>> df = nw.from_native(pd.DataFrame({"a": [1, 2, 3]})) >>> nw.get_native_namespace(df) >>> df = nw.from_native(pl.DataFrame({"a": [1, 2, 3]})) >>> nw.get_native_namespace(df) """ if not obj: msg = "At least one object must be passed to `get_native_namespace`." raise ValueError(msg) result = {_get_native_namespace_single_obj(x) for x in obj} if len(result) != 1: msg = f"Found objects with different native namespaces: {result}." raise ValueError(msg) return result.pop() def _get_native_namespace_single_obj( obj: Frame | Series[Any] | IntoFrame | IntoSeries, ) -> Any: if has_native_namespace(obj): return obj.__native_namespace__() return Version.MAIN.namespace.from_native_object( obj ).implementation.to_native_namespace() def narwhalify( func: Callable[..., Any] | None = None, *, pass_through: bool = True, eager_only: bool = False, series_only: bool = False, allow_series: bool | None = True, ) -> Callable[..., Any]: """Decorate function so it becomes dataframe-agnostic. This will try to convert any dataframe/series-like object into the Narwhals respective DataFrame/Series, while leaving the other parameters as they are. Similarly, if the output of the function is a Narwhals DataFrame or Series, it will be converted back to the original dataframe/series type, while if the output is another type it will be left as is. By setting `pass_through=False`, then every input and every output will be required to be a dataframe/series-like object. Arguments: func: Function to wrap in a `from_native`-`to_native` block. pass_through: Determine what happens if the object can't be converted to Narwhals - `False`: raise an error - `True` (default): pass object through as-is eager_only: Whether to only allow eager objects - `False` (default): don't require `native_object` to be eager - `True`: only convert to Narwhals if `native_object` is eager series_only: Whether to only allow Series - `False` (default): don't require `native_object` to be a Series - `True`: only convert to Narwhals if `native_object` is a Series allow_series: Whether to allow Series (default is only Dataframe / Lazyframe) - `False` or `None`: don't convert to Narwhals if `native_object` is a Series - `True` (default): allow `native_object` to be a Series Returns: Decorated function. Examples: Instead of writing >>> import narwhals as nw >>> def agnostic_group_by_sum(df): ... df = nw.from_native(df, pass_through=True) ... df = df.group_by("a").agg(nw.col("b").sum()) ... return nw.to_native(df) you can just write >>> @nw.narwhalify ... def agnostic_group_by_sum(df): ... return df.group_by("a").agg(nw.col("b").sum()) """ def decorator(func: Callable[..., Any]) -> Callable[..., Any]: @wraps(func) def wrapper(*args: Any, **kwargs: Any) -> Any: args = [ from_native( arg, pass_through=pass_through, eager_only=eager_only, series_only=series_only, allow_series=allow_series, ) for arg in args ] # type: ignore[assignment] kwargs = { name: from_native( value, pass_through=pass_through, eager_only=eager_only, series_only=series_only, allow_series=allow_series, ) for name, value in kwargs.items() } backends = { b() for v in (*args, *kwargs.values()) if (b := getattr(v, "__native_namespace__", None)) } if len(backends) > 1: msg = "Found multiple backends. Make sure that all dataframe/series inputs come from the same backend." raise ValueError(msg) result = func(*args, **kwargs) return to_native(result, pass_through=pass_through) return wrapper if func is None: return decorator # If func is not None, it means the decorator is used without arguments return decorator(func) def to_py_scalar(scalar_like: Any) -> Any: """If a scalar is not Python native, converts it to Python native. Arguments: scalar_like: Scalar-like value. Raises: ValueError: If the object is not convertible to a scalar. Examples: >>> import narwhals as nw >>> import pandas as pd >>> df = nw.from_native(pd.DataFrame({"a": [1, 2, 3]})) >>> nw.to_py_scalar(df["a"].item(0)) 1 >>> import pyarrow as pa >>> df = nw.from_native(pa.table({"a": [1, 2, 3]})) >>> nw.to_py_scalar(df["a"].item(0)) 1 >>> nw.to_py_scalar(1) 1 """ scalar: Any pd = get_pandas() if scalar_like is None or isinstance(scalar_like, NON_TEMPORAL_SCALAR_TYPES): scalar = scalar_like elif ( (np := get_numpy()) and isinstance(scalar_like, np.datetime64) and scalar_like.dtype == "datetime64[ns]" ): ms = scalar_like.item() // MS_PER_SECOND scalar = EPOCH + dt.timedelta(microseconds=ms) elif is_numpy_scalar(scalar_like) or is_cupy_scalar(scalar_like): scalar = scalar_like.item() elif pd and isinstance(scalar_like, pd.Timestamp): scalar = scalar_like.to_pydatetime() elif pd and isinstance(scalar_like, pd.Timedelta): scalar = scalar_like.to_pytimedelta() # pd.Timestamp and pd.Timedelta subclass datetime and timedelta, # so we need to check this separately elif isinstance(scalar_like, TEMPORAL_SCALAR_TYPES): scalar = scalar_like elif _is_pandas_na(scalar_like): scalar = None elif is_pyarrow_scalar(scalar_like): scalar = scalar_like.as_py() else: msg = ( f"Expected object convertible to a scalar, found {type(scalar_like)}.\n" f"{scalar_like!r}" ) raise ValueError(msg) return scalar def _is_pandas_na(obj: Any) -> bool: return bool((pd := get_pandas()) and pd.api.types.is_scalar(obj) and pd.isna(obj)) __all__ = ["get_native_namespace", "narwhalify", "to_native", "to_py_scalar"] narwhals-2.11.0/narwhals/typing.py0000644000000000000000000003251513615410400014043 0ustar00from __future__ import annotations from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeVar, Union from narwhals._compliant import CompliantDataFrame, CompliantLazyFrame, CompliantSeries from narwhals._native import ( IntoDataFrame, IntoDataFrameT, IntoFrame, IntoFrameT, IntoLazyFrame, IntoLazyFrameT, IntoSeries, IntoSeriesT, ) from narwhals._typing import Backend, EagerAllowed, IntoBackend, LazyAllowed if TYPE_CHECKING: import datetime as dt import os from collections.abc import Sequence from decimal import Decimal from types import ModuleType import numpy as np import pandas as pd import polars as pl import pyarrow as pa from typing_extensions import TypeAlias from narwhals import dtypes from narwhals.dataframe import DataFrame, LazyFrame from narwhals.expr import Expr from narwhals.schema import Schema from narwhals.series import Series class SupportsNativeNamespace(Protocol): def __native_namespace__(self) -> ModuleType: ... # ruff: noqa: N802 class DTypes(Protocol): @property def Decimal(self) -> type[dtypes.Decimal]: ... @property def Int128(self) -> type[dtypes.Int128]: ... @property def Int64(self) -> type[dtypes.Int64]: ... @property def Int32(self) -> type[dtypes.Int32]: ... @property def Int16(self) -> type[dtypes.Int16]: ... @property def Int8(self) -> type[dtypes.Int8]: ... @property def UInt128(self) -> type[dtypes.UInt128]: ... @property def UInt64(self) -> type[dtypes.UInt64]: ... @property def UInt32(self) -> type[dtypes.UInt32]: ... @property def UInt16(self) -> type[dtypes.UInt16]: ... @property def UInt8(self) -> type[dtypes.UInt8]: ... @property def Float64(self) -> type[dtypes.Float64]: ... @property def Float32(self) -> type[dtypes.Float32]: ... @property def String(self) -> type[dtypes.String]: ... @property def Boolean(self) -> type[dtypes.Boolean]: ... @property def Object(self) -> type[dtypes.Object]: ... @property def Categorical(self) -> type[dtypes.Categorical]: ... @property def Enum(self) -> type[dtypes.Enum]: ... @property def Datetime(self) -> type[dtypes.Datetime]: ... @property def Duration(self) -> type[dtypes.Duration]: ... @property def Date(self) -> type[dtypes.Date]: ... @property def Field(self) -> type[dtypes.Field]: ... @property def Struct(self) -> type[dtypes.Struct]: ... @property def List(self) -> type[dtypes.List]: ... @property def Array(self) -> type[dtypes.Array]: ... @property def Unknown(self) -> type[dtypes.Unknown]: ... @property def Time(self) -> type[dtypes.Time]: ... @property def Binary(self) -> type[dtypes.Binary]: ... _ShapeT = TypeVar("_ShapeT", bound="tuple[int, ...]") _NDArray: TypeAlias = "np.ndarray[_ShapeT, Any]" _1DArray: TypeAlias = "_NDArray[tuple[int]]" _1DArrayInt: TypeAlias = "np.ndarray[tuple[int], np.dtype[np.integer[Any]]]" _2DArray: TypeAlias = "_NDArray[tuple[int, int]]" # noqa: PYI047 _AnyDArray: TypeAlias = "_NDArray[tuple[int, ...]]" # noqa: PYI047 _NumpyScalar: TypeAlias = "np.generic[Any]" Into1DArray: TypeAlias = "_1DArray | _NumpyScalar" """A 1-dimensional `numpy.ndarray` or scalar that can be converted into one.""" IntoExpr: TypeAlias = Union["Expr", str, "Series[Any]", _1DArray] """Anything which can be converted to an expression. Use this to mean "either a Narwhals expression, or something which can be converted into one". For example, `exprs` in `DataFrame.select` is typed to accept `IntoExpr`, as it can either accept a `nw.Expr` (e.g. `df.select(nw.col('a'))`) or a string which will be interpreted as a `nw.Expr`, e.g. `df.select('a')`. """ Frame: TypeAlias = Union["DataFrame[Any]", "LazyFrame[Any]"] """Narwhals DataFrame or Narwhals LazyFrame. Use this if your function can work with either and your function doesn't care about its backend. Examples: >>> import narwhals as nw >>> from narwhals.typing import Frame >>> @nw.narwhalify ... def agnostic_columns(df: Frame) -> list[str]: ... return df.columns """ FrameT = TypeVar("FrameT", "DataFrame[Any]", "LazyFrame[Any]") """TypeVar bound to Narwhals DataFrame or Narwhals LazyFrame. Use this if your function accepts either `nw.DataFrame` or `nw.LazyFrame` and returns an object of the same kind. Examples: >>> import narwhals as nw >>> from narwhals.typing import FrameT >>> @nw.narwhalify ... def agnostic_func(df: FrameT) -> FrameT: ... return df.with_columns(c=nw.col("a") + 1) """ DataFrameT = TypeVar("DataFrameT", bound="DataFrame[Any]") """TypeVar bound to Narwhals DataFrame. Use this if your function can accept a Narwhals DataFrame and returns a Narwhals DataFrame backed by the same backend. Examples: >>> import narwhals as nw >>> from narwhals.typing import DataFrameT >>> @nw.narwhalify >>> def func(df: DataFrameT) -> DataFrameT: ... return df.with_columns(c=df["a"] + 1) """ LazyFrameT = TypeVar("LazyFrameT", bound="LazyFrame[Any]") SeriesT = TypeVar("SeriesT", bound="Series[Any]") DTypeBackend: TypeAlias = 'Literal["pyarrow", "numpy_nullable"] | None' SizeUnit: TypeAlias = Literal[ "b", "kb", "mb", "gb", "tb", "bytes", "kilobytes", "megabytes", "gigabytes", "terabytes", ] TimeUnit: TypeAlias = Literal["ns", "us", "ms", "s"] AsofJoinStrategy: TypeAlias = Literal["backward", "forward", "nearest"] """Join strategy. - *"backward"*: Selects the last row in the right DataFrame whose `on` key is less than or equal to the left's key. - *"forward"*: Selects the first row in the right DataFrame whose `on` key is greater than or equal to the left's key. - *"nearest"*: Search selects the last row in the right DataFrame whose value is nearest to the left's key. """ ClosedInterval: TypeAlias = Literal["left", "right", "none", "both"] """Define which sides of the interval are closed (inclusive).""" ConcatMethod: TypeAlias = Literal["horizontal", "vertical", "diagonal"] """Concatenating strategy. - *"vertical"*: Concatenate vertically. Column names must match. - *"horizontal"*: Concatenate horizontally. If lengths don't match, then missing rows are filled with null values. - *"diagonal"*: Finds a union between the column schemas and fills missing column values with null. """ FillNullStrategy: TypeAlias = Literal["forward", "backward"] """Strategy used to fill null values.""" JoinStrategy: TypeAlias = Literal["inner", "left", "full", "cross", "semi", "anti"] """Join strategy. - *"inner"*: Returns rows that have matching values in both tables. - *"left"*: Returns all rows from the left table, and the matched rows from the right table. - *"full"*: Returns all rows in both dataframes, with the `suffix` appended to the right join keys. - *"cross"*: Returns the Cartesian product of rows from both tables. - *"semi"*: Filter rows that have a match in the right table. - *"anti"*: Filter rows that do not have a match in the right table. """ PivotAgg: TypeAlias = Literal[ "min", "max", "first", "last", "sum", "mean", "median", "len" ] """A predefined aggregate function string.""" RankMethod: TypeAlias = Literal["average", "min", "max", "dense", "ordinal"] """The method used to assign ranks to tied elements. - *"average"*: The average of the ranks that would have been assigned to all the tied values is assigned to each value. - *"min"*: The minimum of the ranks that would have been assigned to all the tied values is assigned to each value. (This is also referred to as "competition" ranking.) - *"max"*: The maximum of the ranks that would have been assigned to all the tied values is assigned to each value. - *"dense"*: Like "min", but the rank of the next highest element is assigned the rank immediately after those assigned to the tied elements. - *"ordinal"*: All values are given a distinct rank, corresponding to the order that the values occur in the Series. """ RollingInterpolationMethod: TypeAlias = Literal[ "nearest", "higher", "lower", "midpoint", "linear" ] """Interpolation method.""" UniqueKeepStrategy: TypeAlias = Literal["any", "first", "last", "none"] """Which of the duplicate rows to keep. - *"any"*: Does not give any guarantee of which row is kept. This allows more optimizations. - *"none"*: Don't keep duplicate rows. - *"first"*: Keep first unique row. - *"last"*: Keep last unique row. """ ModeKeepStrategy: TypeAlias = Literal["any", "all"] """Which of the mode's to keep. - *"any"*: Does not give any guarantee of which mode is kept. - *"all"*: Keeps all the mode's. """ PandasLikeDType: TypeAlias = "pd.api.extensions.ExtensionDtype | np.dtype[Any]" NumericLiteral: TypeAlias = "int | float | Decimal" TemporalLiteral: TypeAlias = "dt.date | dt.datetime | dt.time | dt.timedelta" NonNestedLiteral: TypeAlias = ( "NumericLiteral | TemporalLiteral | str | bool | bytes | None" ) PythonLiteral: TypeAlias = "NonNestedLiteral | list[Any] | tuple[Any, ...]" NonNestedDType: TypeAlias = "dtypes.NumericType | dtypes.TemporalType | dtypes.String | dtypes.Boolean | dtypes.Binary | dtypes.Categorical | dtypes.Unknown | dtypes.Object" """Any Narwhals DType that does not have required arguments.""" IntoDType: TypeAlias = "dtypes.DType | type[NonNestedDType]" """Anything that can be converted into a Narwhals DType. Examples: >>> import polars as pl >>> import narwhals as nw >>> df_native = pl.DataFrame({"a": [1, 2, 3], "b": [4.0, 5.0, 6.0]}) >>> df = nw.from_native(df_native) >>> df.select( ... nw.col("a").cast(nw.Int32), ... nw.col("b").cast(nw.String()).str.split(".").cast(nw.List(nw.Int8)), ... ) ┌──────────────────┐ |Narwhals DataFrame| |------------------| |shape: (3, 2) | |┌─────┮──────────┐| |│ a ┆ b │| |│ --- ┆ --- │| |│ i32 ┆ list[i8] │| |╞═════╩══════════╡| |│ 1 ┆ [4, 0] │| |│ 2 ┆ [5, 0] │| |│ 3 ┆ [6, 0] │| |└─────â”ī──────────┘| └──────────────────┘ """ # TODO @dangotbanned: fix this? # Constructor allows tuples, but we don't support that *everywhere* yet IntoSchema: TypeAlias = "Mapping[str, dtypes.DType] | Schema" """Anything that can be converted into a Narwhals Schema. Defined by column names and their associated *instantiated* Narwhals DType. Examples: >>> import narwhals as nw >>> import pyarrow as pa >>> data = {"a": [1, 2, 3], "b": [None, "hi", "howdy"], "c": [2.1, 2.0, None]} >>> nw.DataFrame.from_dict( ... data, ... schema={"a": nw.UInt8(), "b": nw.String(), "c": nw.Float32()}, ... backend="pyarrow", ... ) ┌────────────────────────┐ | Narwhals DataFrame | |------------------------| |pyarrow.Table | |a: uint8 | |b: string | |c: float | |---- | |a: [[1,2,3]] | |b: [[null,"hi","howdy"]]| |c: [[2.1,2,null]] | └────────────────────────┘ """ IntoArrowSchema: TypeAlias = "pa.Schema | Mapping[str, pa.DataType]" IntoPolarsSchema: TypeAlias = "pl.Schema | Mapping[str, pl.DataType]" IntoPandasSchema: TypeAlias = Mapping[str, PandasLikeDType] FileSource: TypeAlias = "str | os.PathLike[str]" """Path to a file. Either a string or an object that implements [`__fspath__`], such as [`pathlib.Path`]. [`__fspath__`]: https://docs.python.org/3/library/os.html#os.PathLike [`pathlib.Path`]: https://docs.python.org/3/library/pathlib.html#pathlib.Path """ # Annotations for `__getitem__` methods _T = TypeVar("_T") _Slice: TypeAlias = "slice[_T, Any, Any] | slice[Any, _T, Any] | slice[None, None, _T]" _SliceNone: TypeAlias = "slice[None, None, None]" # Index/column positions SingleIndexSelector: TypeAlias = int _SliceIndex: TypeAlias = "_Slice[int] | _SliceNone" """E.g. `[1:]` or `[:3]` or `[::2]`.""" SizedMultiIndexSelector: TypeAlias = "Sequence[int] | _T | _1DArrayInt" MultiIndexSelector: TypeAlias = "_SliceIndex | SizedMultiIndexSelector[_T]" # Labels/column names SingleNameSelector: TypeAlias = str _SliceName: TypeAlias = "_Slice[str] | _SliceNone" SizedMultiNameSelector: TypeAlias = "Sequence[str] | _T | _1DArray" MultiNameSelector: TypeAlias = "_SliceName | SizedMultiNameSelector[_T]" # Mixed selectors SingleColSelector: TypeAlias = "SingleIndexSelector | SingleNameSelector" MultiColSelector: TypeAlias = "MultiIndexSelector[_T] | MultiNameSelector[_T]" __all__ = [ "Backend", "CompliantDataFrame", "CompliantLazyFrame", "CompliantSeries", "DataFrameT", "EagerAllowed", "Frame", "FrameT", "IntoBackend", "IntoDataFrame", "IntoDataFrameT", "IntoExpr", "IntoFrame", "IntoFrameT", "IntoLazyFrame", "IntoLazyFrameT", "IntoSeries", "IntoSeriesT", "LazyAllowed", ] narwhals-2.11.0/narwhals/utils.py0000644000000000000000000000033713615410400013666 0ustar00# Re-export some functions from `_utils` to make them public. from __future__ import annotations from narwhals._utils import Implementation, Version, parse_version __all__ = ["Implementation", "Version", "parse_version"] narwhals-2.11.0/narwhals/_arrow/__init__.py0000644000000000000000000000000013615410400015541 0ustar00narwhals-2.11.0/narwhals/_arrow/dataframe.py0000644000000000000000000007457513615410400015762 0ustar00from __future__ import annotations from collections.abc import Collection, Iterator, Mapping, Sequence from typing import TYPE_CHECKING, Any, Literal, cast, overload import pyarrow as pa import pyarrow.compute as pc from narwhals._arrow.series import ArrowSeries from narwhals._arrow.utils import ( arange, concat_tables, narwhals_to_native_dtype, native_to_narwhals_dtype, repeat, ) from narwhals._compliant import EagerDataFrame from narwhals._utils import ( Implementation, Version, check_column_names_are_unique, convert_str_slice_to_int_slice, generate_temporary_column_name, not_implemented, parse_columns_to_drop, scale_bytes, supports_arrow_c_stream, zip_strict, ) from narwhals.dependencies import is_numpy_array_1d from narwhals.exceptions import ShapeError if TYPE_CHECKING: from collections.abc import Iterable from io import BytesIO from pathlib import Path from types import ModuleType import pandas as pd import polars as pl from typing_extensions import Self, TypeAlias, TypeIs from narwhals._arrow.expr import ArrowExpr from narwhals._arrow.group_by import ArrowGroupBy from narwhals._arrow.namespace import ArrowNamespace from narwhals._arrow.typing import ( # type: ignore[attr-defined] ChunkedArrayAny, Order, ) from narwhals._compliant.typing import CompliantDataFrameAny, CompliantLazyFrameAny from narwhals._spark_like.utils import SparkSession from narwhals._translate import IntoArrowTable from narwhals._typing import _EagerAllowedImpl, _LazyAllowedImpl from narwhals._utils import Version, _LimitedContext from narwhals.dtypes import DType from narwhals.typing import ( IntoSchema, JoinStrategy, SizedMultiIndexSelector, SizedMultiNameSelector, SizeUnit, UniqueKeepStrategy, _1DArray, _2DArray, _SliceIndex, _SliceName, ) JoinType: TypeAlias = Literal[ "left semi", "right semi", "left anti", "right anti", "inner", "left outer", "right outer", "full outer", ] class ArrowDataFrame( EagerDataFrame["ArrowSeries", "ArrowExpr", "pa.Table", "ChunkedArrayAny"] ): _implementation = Implementation.PYARROW def __init__( self, native_dataframe: pa.Table, *, version: Version, validate_column_names: bool, validate_backend_version: bool = False, ) -> None: if validate_column_names: check_column_names_are_unique(native_dataframe.column_names) if validate_backend_version: self._validate_backend_version() self._native_frame = native_dataframe self._version = version @classmethod def from_arrow(cls, data: IntoArrowTable, /, *, context: _LimitedContext) -> Self: backend_version = context._implementation._backend_version() if cls._is_native(data): native = data elif backend_version >= (14,) or isinstance(data, Collection): native = pa.table(data) elif supports_arrow_c_stream(data): # pragma: no cover msg = f"'pyarrow>=14.0.0' is required for `from_arrow` for object of type {type(data).__name__!r}." raise ModuleNotFoundError(msg) else: # pragma: no cover msg = f"`from_arrow` is not supported for object of type {type(data).__name__!r}." raise TypeError(msg) return cls.from_native(native, context=context) @classmethod def from_dict( cls, data: Mapping[str, Any], /, *, context: _LimitedContext, schema: IntoSchema | Mapping[str, DType | None] | None, ) -> Self: if not schema and not data: return cls.from_native(pa.table({}), context=context) if not schema: return cls.from_native(pa.table(data), context=context) # type: ignore[arg-type] if not any(dtype is None for dtype in schema.values()): from narwhals.schema import Schema pa_schema = Schema(cast("IntoSchema", schema)).to_arrow() if pa_schema and not data: native = pa_schema.empty_table() else: native = pa.Table.from_pydict(data, schema=pa_schema) return cls.from_native(native, context=context) if context._implementation._backend_version() < (14,): msg = "Passing `None` dtype in `from_dict` requires PyArrow>=14" raise NotImplementedError(msg) res = pa.table( { name: pa.chunked_array( # type: ignore[misc] [data[name] if data else []], type=narwhals_to_native_dtype(nw_dtype, version=context._version) if nw_dtype is not None else None, ) for name, nw_dtype in schema.items() } ) return cls.from_native(pa.table(res), context=context) @classmethod def from_dicts( cls, data: Sequence[Mapping[str, Any]], /, *, context: _LimitedContext, schema: IntoSchema | Mapping[str, DType | None] | None, ) -> Self: from narwhals.schema import Schema if schema and any(dtype is None for dtype in schema.values()): msg = "`from_dicts` with `schema` where any dtype is `None` is not supported for PyArrow." raise NotImplementedError(msg) pa_schema = ( Schema(cast("IntoSchema", schema)).to_arrow() if schema is not None else schema ) if pa_schema and not data: native = pa_schema.empty_table() else: native = pa.Table.from_pylist(data, schema=pa_schema) return cls.from_native(native, context=context) @staticmethod def _is_native(obj: pa.Table | Any) -> TypeIs[pa.Table]: return isinstance(obj, pa.Table) @classmethod def from_native(cls, data: pa.Table, /, *, context: _LimitedContext) -> Self: return cls(data, version=context._version, validate_column_names=True) @classmethod def from_numpy( cls, data: _2DArray, /, *, context: _LimitedContext, schema: IntoSchema | Sequence[str] | None, ) -> Self: from narwhals.schema import Schema arrays = [pa.array(val) for val in data.T] if isinstance(schema, (Mapping, Schema)): native = pa.Table.from_arrays(arrays, schema=Schema(schema).to_arrow()) else: native = pa.Table.from_arrays(arrays, cls._numpy_column_names(data, schema)) return cls.from_native(native, context=context) def __narwhals_namespace__(self) -> ArrowNamespace: from narwhals._arrow.namespace import ArrowNamespace return ArrowNamespace(version=self._version) def __native_namespace__(self) -> ModuleType: if self._implementation is Implementation.PYARROW: return self._implementation.to_native_namespace() msg = f"Expected pyarrow, got: {type(self._implementation)}" # pragma: no cover raise AssertionError(msg) def __narwhals_dataframe__(self) -> Self: return self def __narwhals_lazyframe__(self) -> Self: return self def _with_version(self, version: Version) -> Self: return self.__class__(self.native, version=version, validate_column_names=False) def _with_native(self, df: pa.Table, *, validate_column_names: bool = True) -> Self: return self.__class__( df, version=self._version, validate_column_names=validate_column_names ) @property def shape(self) -> tuple[int, int]: return self.native.shape def __len__(self) -> int: return len(self.native) def row(self, index: int) -> tuple[Any, ...]: return tuple(col[index] for col in self.native.itercolumns()) @overload def rows(self, *, named: Literal[True]) -> list[dict[str, Any]]: ... @overload def rows(self, *, named: Literal[False]) -> list[tuple[Any, ...]]: ... @overload def rows(self, *, named: bool) -> list[tuple[Any, ...]] | list[dict[str, Any]]: ... def rows(self, *, named: bool) -> list[tuple[Any, ...]] | list[dict[str, Any]]: if not named: return list(self.iter_rows(named=False, buffer_size=512)) # type: ignore[return-value] return self.native.to_pylist() def iter_columns(self) -> Iterator[ArrowSeries]: for name, series in zip_strict(self.columns, self.native.itercolumns()): yield ArrowSeries.from_native(series, context=self, name=name) _iter_columns = iter_columns def iter_rows( self, *, named: bool, buffer_size: int ) -> Iterator[tuple[Any, ...]] | Iterator[dict[str, Any]]: df = self.native num_rows = df.num_rows if not named: for i in range(0, num_rows, buffer_size): rows = df[i : i + buffer_size].to_pydict().values() yield from zip_strict(*rows) else: for i in range(0, num_rows, buffer_size): yield from df[i : i + buffer_size].to_pylist() def get_column(self, name: str) -> ArrowSeries: if not isinstance(name, str): msg = f"Expected str, got: {type(name)}" raise TypeError(msg) return ArrowSeries.from_native(self.native[name], context=self, name=name) def __array__(self, dtype: Any, *, copy: bool | None) -> _2DArray: return self.native.__array__(dtype, copy=copy) def _gather(self, rows: SizedMultiIndexSelector[ChunkedArrayAny]) -> Self: if len(rows) == 0: return self._with_native(self.native.slice(0, 0)) if self._backend_version < (18,) and isinstance(rows, tuple): rows = list(rows) return self._with_native(self.native.take(rows)) def _gather_slice(self, rows: _SliceIndex | range) -> Self: start = rows.start or 0 stop = rows.stop if rows.stop is not None else len(self.native) if start < 0: start = len(self.native) + start if stop < 0: stop = len(self.native) + stop if rows.step is not None and rows.step != 1: msg = "Slicing with step is not supported on PyArrow tables" raise NotImplementedError(msg) return self._with_native(self.native.slice(start, stop - start)) def _select_slice_name(self, columns: _SliceName) -> Self: start, stop, step = convert_str_slice_to_int_slice(columns, self.columns) return self._with_native(self.native.select(self.columns[start:stop:step])) def _select_slice_index(self, columns: _SliceIndex | range) -> Self: return self._with_native( self.native.select(self.columns[columns.start : columns.stop : columns.step]) ) def _select_multi_index( self, columns: SizedMultiIndexSelector[ChunkedArrayAny] ) -> Self: selector: Sequence[int] if isinstance(columns, pa.ChunkedArray): # TODO @dangotbanned: Fix upstream with `pa.ChunkedArray.to_pylist(self) -> list[Any]:` selector = cast("Sequence[int]", columns.to_pylist()) # TODO @dangotbanned: Fix upstream, it is actually much narrower # **Doesn't accept `ndarray`** elif is_numpy_array_1d(columns): selector = columns.tolist() else: selector = columns return self._with_native(self.native.select(selector)) def _select_multi_name( self, columns: SizedMultiNameSelector[ChunkedArrayAny] ) -> Self: selector: Sequence[str] | _1DArray if isinstance(columns, pa.ChunkedArray): # TODO @dangotbanned: Fix upstream with `pa.ChunkedArray.to_pylist(self) -> list[Any]:` selector = cast("Sequence[str]", columns.to_pylist()) else: selector = columns # NOTE: Fixed in https://github.com/zen-xu/pyarrow-stubs/pull/221 return self._with_native(self.native.select(selector)) # pyright: ignore[reportArgumentType] @property def schema(self) -> dict[str, DType]: return { field.name: native_to_narwhals_dtype(field.type, self._version) for field in self.native.schema } def collect_schema(self) -> dict[str, DType]: return self.schema def estimated_size(self, unit: SizeUnit) -> int | float: sz = self.native.nbytes return scale_bytes(sz, unit) explode = not_implemented() @property def columns(self) -> list[str]: return self.native.column_names def simple_select(self, *column_names: str) -> Self: return self._with_native( self.native.select(list(column_names)), validate_column_names=False ) def select(self, *exprs: ArrowExpr) -> Self: new_series = self._evaluate_exprs(*exprs) if not new_series: # return empty dataframe, like Polars does return self._with_native( self.native.__class__.from_arrays([]), validate_column_names=False ) names = [s.name for s in new_series] align = new_series[0]._align_full_broadcast reshaped = align(*new_series) df = pa.Table.from_arrays([s.native for s in reshaped], names=names) return self._with_native(df, validate_column_names=True) def _extract_comparand(self, other: ArrowSeries) -> ChunkedArrayAny: length = len(self) if not other._broadcast: if (len_other := len(other)) != length: msg = f"Expected object of length {length}, got: {len_other}." raise ShapeError(msg) return other.native value = other.native[0] return pa.chunked_array([pa.repeat(value, length)]) def with_columns(self, *exprs: ArrowExpr) -> Self: # NOTE: We use a faux-mutable variable and repeatedly "overwrite" (native_frame) # All `pyarrow` data is immutable, so this is fine native_frame = self.native new_columns = self._evaluate_exprs(*exprs) columns = self.columns for col_value in new_columns: col_name = col_value.name column = self._extract_comparand(col_value) native_frame = ( native_frame.set_column(columns.index(col_name), col_name, column=column) if col_name in columns else native_frame.append_column(col_name, column=column) ) return self._with_native(native_frame, validate_column_names=False) def group_by( self, keys: Sequence[str] | Sequence[ArrowExpr], *, drop_null_keys: bool ) -> ArrowGroupBy: from narwhals._arrow.group_by import ArrowGroupBy return ArrowGroupBy(self, keys, drop_null_keys=drop_null_keys) def join( self, other: Self, *, how: JoinStrategy, left_on: Sequence[str] | None, right_on: Sequence[str] | None, suffix: str, ) -> Self: how_to_join_map: dict[str, JoinType] = { "anti": "left anti", "semi": "left semi", "inner": "inner", "left": "left outer", "full": "full outer", } if how == "cross": plx = self.__narwhals_namespace__() key_token = generate_temporary_column_name( n_bytes=8, columns=[*self.columns, *other.columns] ) return self._with_native( self.with_columns(plx.lit(0, None).alias(key_token).broadcast()) .native.join( other.with_columns( plx.lit(0, None).alias(key_token).broadcast() ).native, keys=key_token, right_keys=key_token, join_type="inner", right_suffix=suffix, ) .drop([key_token]) ) coalesce_keys = how != "full" # polars full join does not coalesce keys return self._with_native( self.native.join( other.native, keys=left_on or [], # type: ignore[arg-type] right_keys=right_on, # type: ignore[arg-type] join_type=how_to_join_map[how], right_suffix=suffix, coalesce_keys=coalesce_keys, ) ) join_asof = not_implemented() 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), validate_column_names=False) def drop_nulls(self, subset: Sequence[str] | None) -> Self: if subset is None: return self._with_native(self.native.drop_null(), validate_column_names=False) plx = self.__narwhals_namespace__() mask = ~plx.any_horizontal(plx.col(*subset).is_null(), ignore_nulls=True) return self.filter(mask) def sort(self, *by: str, descending: bool | Sequence[bool], nulls_last: bool) -> Self: if isinstance(descending, bool): order: Order = "descending" if descending else "ascending" sorting: list[tuple[str, Order]] = [(key, order) for key in by] else: sorting = [ (key, "descending" if is_descending else "ascending") for key, is_descending in zip_strict(by, descending) ] null_placement = "at_end" if nulls_last else "at_start" return self._with_native( self.native.sort_by(sorting, null_placement=null_placement), validate_column_names=False, ) def top_k(self, k: int, *, by: Iterable[str], reverse: bool | Sequence[bool]) -> Self: if isinstance(reverse, bool): order: Order = "ascending" if reverse else "descending" sorting: list[tuple[str, Order]] = [(key, order) for key in by] else: sorting = [ (key, "ascending" if is_ascending else "descending") for key, is_ascending in zip_strict(by, reverse) ] return self._with_native( self.native.take(pc.select_k_unstable(self.native, k, sorting)), # type: ignore[call-overload] validate_column_names=False, ) def to_pandas(self) -> pd.DataFrame: return self.native.to_pandas() def to_polars(self) -> pl.DataFrame: import polars as pl # ignore-banned-import return pl.from_arrow(self.native) # type: ignore[return-value] def to_numpy(self, dtype: Any = None, *, copy: bool | None = None) -> _2DArray: import numpy as np # ignore-banned-import arr: Any = np.column_stack([col.to_numpy() for col in self.native.columns]) return arr @overload def to_dict(self, *, as_series: Literal[True]) -> dict[str, ArrowSeries]: ... @overload def to_dict(self, *, as_series: Literal[False]) -> dict[str, list[Any]]: ... def to_dict( self, *, as_series: bool ) -> dict[str, ArrowSeries] | dict[str, list[Any]]: it = self.iter_columns() if as_series: return {ser.name: ser for ser in it} return {ser.name: ser.to_list() for ser in it} def with_row_index(self, name: str, order_by: Sequence[str] | None) -> Self: plx = self.__narwhals_namespace__() data = arange(0, len(self), 1) if order_by is None: row_index = plx._expr._from_series( plx._series.from_iterable(data, context=self, name=name) ) return self.select(row_index, plx.all()) indices = pc.sort_indices(self.native, [(by, "ascending") for by in order_by]) if self._backend_version < (20,): new_col = data.take(pc.sort_indices(indices)) else: new_col = pc.scatter(data, indices.cast(pa.int64())) # type: ignore[attr-defined] return self._with_native(self.native.add_column(0, name, new_col)) def filter(self, predicate: ArrowExpr) -> Self: mask_native = self._evaluate_single_output_expr(predicate).native return self._with_native( self.native.filter(mask_native), validate_column_names=False ) def head(self, n: int) -> Self: df = self.native if n >= 0: return self._with_native(df.slice(0, n), validate_column_names=False) num_rows = df.num_rows return self._with_native( df.slice(0, max(0, num_rows + n)), validate_column_names=False ) def tail(self, n: int) -> Self: df = self.native if n >= 0: num_rows = df.num_rows return self._with_native( df.slice(max(0, num_rows - n)), validate_column_names=False ) return self._with_native(df.slice(abs(n)), validate_column_names=False) def lazy( self, backend: _LazyAllowedImpl | None = None, *, session: SparkSession | None = None, ) -> CompliantLazyFrameAny: if backend is None: return self if backend is Implementation.DUCKDB: import duckdb # ignore-banned-import from narwhals._duckdb.dataframe import DuckDBLazyFrame _df = self.native return DuckDBLazyFrame( duckdb.table("_df"), validate_backend_version=True, version=self._version ) if backend is Implementation.POLARS: import polars as pl # ignore-banned-import from narwhals._polars.dataframe import PolarsLazyFrame return PolarsLazyFrame( cast("pl.DataFrame", pl.from_arrow(self.native)).lazy(), validate_backend_version=True, version=self._version, ) if backend is Implementation.DASK: import dask.dataframe as dd # ignore-banned-import from narwhals._dask.dataframe import DaskLazyFrame return DaskLazyFrame( dd.from_pandas(self.native.to_pandas()), validate_backend_version=True, version=self._version, ) if backend is Implementation.IBIS: import ibis # ignore-banned-import from narwhals._ibis.dataframe import IbisLazyFrame return IbisLazyFrame( ibis.memtable(self.native, columns=self.columns), validate_backend_version=True, version=self._version, ) if backend.is_spark_like(): from narwhals._spark_like.dataframe import SparkLikeLazyFrame if session is None: msg = "Spark like backends require `session` to be not None." raise ValueError(msg) return SparkLikeLazyFrame._from_compliant_dataframe( self, session=session, implementation=backend, version=self._version ) raise AssertionError # pragma: no cover def collect( self, backend: _EagerAllowedImpl | None, **kwargs: Any ) -> CompliantDataFrameAny: if backend is Implementation.PYARROW or backend is None: from narwhals._arrow.dataframe import ArrowDataFrame return ArrowDataFrame( self.native, version=self._version, validate_column_names=False ) if backend is Implementation.PANDAS: from narwhals._pandas_like.dataframe import PandasLikeDataFrame return PandasLikeDataFrame( self.native.to_pandas(), implementation=Implementation.PANDAS, validate_backend_version=True, version=self._version, validate_column_names=False, ) if backend is Implementation.POLARS: import polars as pl # ignore-banned-import from narwhals._polars.dataframe import PolarsDataFrame return PolarsDataFrame( cast("pl.DataFrame", pl.from_arrow(self.native)), validate_backend_version=True, version=self._version, ) msg = f"Unsupported `backend` value: {backend}" # pragma: no cover raise AssertionError(msg) # pragma: no cover def clone(self) -> Self: return self._with_native(self.native, validate_column_names=False) def item(self, row: int | None, column: int | str | None) -> Any: from narwhals._arrow.series import maybe_extract_py_scalar if row is None and column is None: if self.shape != (1, 1): msg = ( "can only call `.item()` if the dataframe is of shape (1, 1)," " or if explicit row/col values are provided;" f" frame has shape {self.shape!r}" ) raise ValueError(msg) return maybe_extract_py_scalar(self.native[0][0], return_py_scalar=True) if row is None or column is None: msg = "cannot call `.item()` with only one of `row` or `column`" raise ValueError(msg) _col = self.columns.index(column) if isinstance(column, str) else column return maybe_extract_py_scalar(self.native[_col][row], return_py_scalar=True) def rename(self, mapping: Mapping[str, str]) -> Self: names: dict[str, str] | list[str] if self._backend_version >= (17,): names = cast("dict[str, str]", mapping) else: # pragma: no cover names = [mapping.get(c, c) for c in self.columns] return self._with_native(self.native.rename_columns(names)) def write_parquet(self, file: str | Path | BytesIO) -> None: import pyarrow.parquet as pp pp.write_table(self.native, file) @overload def write_csv(self, file: None) -> str: ... @overload def write_csv(self, file: str | Path | BytesIO) -> None: ... def write_csv(self, file: str | Path | BytesIO | None) -> str | None: import pyarrow.csv as pa_csv if file is None: csv_buffer = pa.BufferOutputStream() pa_csv.write_csv(self.native, csv_buffer) return csv_buffer.getvalue().to_pybytes().decode() pa_csv.write_csv(self.native, file) return None def is_unique(self) -> ArrowSeries: import numpy as np # ignore-banned-import col_token = generate_temporary_column_name(n_bytes=8, columns=self.columns) row_index = pa.array(np.arange(len(self))) keep_idx = ( self.native.append_column(col_token, row_index) .group_by(self.columns) .aggregate([(col_token, "min"), (col_token, "max")]) ) native = pa.chunked_array( pc.and_( pc.is_in(row_index, keep_idx[f"{col_token}_min"]), pc.is_in(row_index, keep_idx[f"{col_token}_max"]), ) ) return ArrowSeries.from_native(native, context=self) def unique( self, subset: Sequence[str] | None, *, keep: UniqueKeepStrategy, maintain_order: bool | None = None, order_by: Sequence[str] | None, ) -> Self: # The param `maintain_order` is only here for compatibility with the Polars API # and has no effect on the output. import numpy as np # ignore-banned-import if subset and (error := self._check_columns_exist(subset)): raise error subset = list(subset or self.columns) if keep in {"any", "first", "last"}: from narwhals._arrow.group_by import ArrowGroupBy agg_func = ArrowGroupBy._REMAP_UNIQUE[keep] col_token = generate_temporary_column_name(n_bytes=8, columns=self.columns) if order_by and maintain_order: idx_token = generate_temporary_column_name( n_bytes=8, columns=[*self.columns, col_token] ) df = ( self.with_row_index(idx_token, order_by=None) .sort(*order_by, nulls_last=False, descending=False) .unique(subset=subset, keep=keep, maintain_order=False, order_by=None) ) return df.sort(idx_token, descending=False, nulls_last=False).drop( [idx_token], strict=False ) if order_by: native = self.sort(*order_by, nulls_last=False, descending=False).native else: native = self.native keep_idx_native = ( native.append_column(col_token, pa.array(np.arange(len(self)))) .group_by(subset) .aggregate([(col_token, agg_func)]) .column(f"{col_token}_{agg_func}") ) return self._with_native( native.take(keep_idx_native), validate_column_names=False ) keep_idx = self.simple_select(*subset).is_unique() plx = self.__narwhals_namespace__() return self.filter(plx._expr._from_series(keep_idx)) def gather_every(self, n: int, offset: int) -> Self: return self._with_native(self.native[offset::n], validate_column_names=False) def to_arrow(self) -> pa.Table: return self.native def sample( self, n: int | None, *, fraction: float | None, with_replacement: bool, seed: int | None, ) -> Self: import numpy as np # ignore-banned-import num_rows = len(self) if n is None and fraction is not None: n = int(num_rows * fraction) rng = np.random.default_rng(seed=seed) idx = np.arange(num_rows) mask = rng.choice(idx, size=n, replace=with_replacement) return self._with_native(self.native.take(mask), validate_column_names=False) def unpivot( self, on: Sequence[str] | None, index: Sequence[str] | None, variable_name: str, value_name: str, ) -> Self: # TODO(Unassigned): Even with promote_options="permissive", pyarrow does not # upcast numeric to non-numeric (e.g. string) datatypes n = len(self) index = [] if index is None else list(index) on_ = (c for c in self.columns if c not in index) if on is None else iter(on) index_cols = self.native.select(index) column = self.native.column tables = ( index_cols.append_column(variable_name, repeat(name, n)).append_column( value_name, column(name) ) for name in on_ ) return self._with_native(concat_tables(tables, "permissive")) pivot = not_implemented() narwhals-2.11.0/narwhals/_arrow/expr.py0000644000000000000000000001406013615410400014773 0ustar00from __future__ import annotations from typing import TYPE_CHECKING, Any import pyarrow as pa import pyarrow.compute as pc from narwhals._arrow.series import ArrowSeries from narwhals._compliant import EagerExpr from narwhals._expression_parsing import evaluate_output_names_and_aliases from narwhals._utils import ( Implementation, generate_temporary_column_name, not_implemented, ) if TYPE_CHECKING: from collections.abc import Sequence from typing_extensions import Self from narwhals._arrow.dataframe import ArrowDataFrame from narwhals._arrow.namespace import ArrowNamespace from narwhals._compliant.typing import AliasNames, EvalNames, EvalSeries from narwhals._utils import Version, _LimitedContext class ArrowExpr(EagerExpr["ArrowDataFrame", ArrowSeries]): _implementation: Implementation = Implementation.PYARROW def __init__( self, call: EvalSeries[ArrowDataFrame, ArrowSeries], *, evaluate_output_names: EvalNames[ArrowDataFrame], alias_output_names: AliasNames | None, version: Version, implementation: Implementation = Implementation.PYARROW, ) -> None: self._call = call self._evaluate_output_names = evaluate_output_names self._alias_output_names = alias_output_names self._version = version @classmethod def from_column_names( cls: type[Self], evaluate_column_names: EvalNames[ArrowDataFrame], /, *, context: _LimitedContext, ) -> Self: def func(df: ArrowDataFrame) -> list[ArrowSeries]: try: return [ ArrowSeries( df.native[column_name], name=column_name, version=df._version ) for column_name in evaluate_column_names(df) ] except KeyError as e: if error := df._check_columns_exist(evaluate_column_names(df)): raise error from e raise return cls( func, evaluate_output_names=evaluate_column_names, alias_output_names=None, version=context._version, ) @classmethod def from_column_indices(cls, *column_indices: int, context: _LimitedContext) -> Self: def func(df: ArrowDataFrame) -> list[ArrowSeries]: tbl = df.native cols = df.columns return [ ArrowSeries.from_native(tbl[i], name=cols[i], context=df) for i in column_indices ] return cls( func, evaluate_output_names=cls._eval_names_indices(column_indices), alias_output_names=None, version=context._version, ) def __narwhals_namespace__(self) -> ArrowNamespace: from narwhals._arrow.namespace import ArrowNamespace return ArrowNamespace(version=self._version) def _reuse_series_extra_kwargs( self, *, returns_scalar: bool = False ) -> dict[str, Any]: return {"_return_py_scalar": False} if returns_scalar else {} def over(self, partition_by: Sequence[str], order_by: Sequence[str]) -> Self: meta = self._metadata if partition_by and not meta.is_scalar_like: msg = "Only aggregation or literal operations are supported in grouped `over` context for PyArrow." raise NotImplementedError(msg) if not partition_by: # e.g. `nw.col('a').cum_sum().order_by(key)` # which we can always easily support, as it doesn't require grouping. assert order_by # noqa: S101 def func(df: ArrowDataFrame) -> Sequence[ArrowSeries]: token = generate_temporary_column_name(8, df.columns) df = df.with_row_index(token, order_by=None).sort( *order_by, descending=False, nulls_last=False ) results = self(df.drop([token], strict=True)) if meta is not None and meta.is_scalar_like: # We need to broadcast the results to the original size, since # `over` is a length-preserving operation. size = len(df) return [s._with_native(pa.repeat(s.item(), size)) for s in results] # TODO(marco): is there a way to do this efficiently without # doing 2 sorts? Here we're sorting the dataframe and then # again calling `sort_indices`. `ArrowSeries.scatter` would also sort. sorting_indices = pc.sort_indices(df.get_column(token).native) return [s._with_native(s.native.take(sorting_indices)) for s in results] else: def func(df: ArrowDataFrame) -> Sequence[ArrowSeries]: if order_by: df = df.sort(*order_by, descending=False, nulls_last=False) output_names, aliases = evaluate_output_names_and_aliases(self, df, []) if overlap := set(output_names).intersection(partition_by): # E.g. `df.select(nw.all().sum().over('a'))`. This is well-defined, # we just don't support it yet. msg = ( f"Column names {overlap} appear in both expression output names and in `over` keys.\n" "This is not yet supported." ) raise NotImplementedError(msg) tmp = df.group_by(partition_by, drop_null_keys=False).agg(self) tmp = df.simple_select(*partition_by).join( tmp, how="left", left_on=partition_by, right_on=partition_by, suffix="_right", ) return [tmp.get_column(alias) for alias in aliases] return self.__class__( func, evaluate_output_names=self._evaluate_output_names, alias_output_names=self._alias_output_names, version=self._version, ) ewm_mean = not_implemented() narwhals-2.11.0/narwhals/_arrow/group_by.py0000644000000000000000000002112413615410400015642 0ustar00from __future__ import annotations import collections from typing import TYPE_CHECKING, Any, ClassVar import pyarrow as pa import pyarrow.compute as pc from narwhals._arrow.utils import cast_to_comparable_string_types, extract_py_scalar from narwhals._compliant import EagerGroupBy from narwhals._expression_parsing import evaluate_output_names_and_aliases from narwhals._utils import generate_temporary_column_name, requires if TYPE_CHECKING: from collections.abc import Iterator, Mapping, Sequence from narwhals._arrow.dataframe import ArrowDataFrame from narwhals._arrow.expr import ArrowExpr from narwhals._arrow.typing import ( # type: ignore[attr-defined] AggregateOptions, Aggregation, Incomplete, ) from narwhals._compliant.typing import NarwhalsAggregation from narwhals.typing import UniqueKeepStrategy class ArrowGroupBy(EagerGroupBy["ArrowDataFrame", "ArrowExpr", "Aggregation"]): _REMAP_AGGS: ClassVar[Mapping[NarwhalsAggregation, Aggregation]] = { "sum": "sum", "mean": "mean", "median": "approximate_median", "max": "max", "min": "min", "std": "stddev", "var": "variance", "len": "count", "n_unique": "count_distinct", "count": "count", "all": "all", "any": "any", "first": "first", "last": "last", } _REMAP_UNIQUE: ClassVar[Mapping[UniqueKeepStrategy, Aggregation]] = { "any": "min", "first": "min", "last": "max", } _OPTION_COUNT_ALL: ClassVar[frozenset[NarwhalsAggregation]] = frozenset( ("len", "n_unique") ) _OPTION_COUNT_VALID: ClassVar[frozenset[NarwhalsAggregation]] = frozenset(("count",)) _OPTION_ORDERED: ClassVar[frozenset[NarwhalsAggregation]] = frozenset( ("first", "last") ) _OPTION_VARIANCE: ClassVar[frozenset[NarwhalsAggregation]] = frozenset(("std", "var")) _OPTION_SCALAR: ClassVar[frozenset[NarwhalsAggregation]] = frozenset(("any", "all")) def __init__( self, df: ArrowDataFrame, keys: Sequence[ArrowExpr] | Sequence[str], /, *, drop_null_keys: bool, ) -> None: self._df = df frame, self._keys, self._output_key_names = self._parse_keys(df, keys=keys) self._compliant_frame = frame.drop_nulls(self._keys) if drop_null_keys else frame self._grouped = pa.TableGroupBy(self.compliant.native, self._keys) self._drop_null_keys = drop_null_keys def _configure_agg( self, grouped: pa.TableGroupBy, expr: ArrowExpr, / ) -> tuple[pa.TableGroupBy, Aggregation, AggregateOptions | None]: option: AggregateOptions | None = None function_name = self._leaf_name(expr) kwargs = self._kwargs(expr) if function_name in self._OPTION_VARIANCE: ddof = kwargs["ddof"] option = pc.VarianceOptions(ddof=ddof) elif function_name in self._OPTION_COUNT_ALL: option = pc.CountOptions(mode="all") elif function_name in self._OPTION_COUNT_VALID: option = pc.CountOptions(mode="only_valid") elif function_name in self._OPTION_SCALAR: option = pc.ScalarAggregateOptions(min_count=0) elif function_name in self._OPTION_ORDERED: grouped, option = self._ordered_agg(grouped, function_name) return grouped, self._remap_expr_name(function_name), option def _ordered_agg( self, grouped: pa.TableGroupBy, name: NarwhalsAggregation, / ) -> tuple[pa.TableGroupBy, AggregateOptions]: """The default behavior of `pyarrow` raises when `first` or `last` are used. You'd see an error like: ArrowNotImplementedError: Using ordered aggregator in multiple threaded execution is not supported We need to **disable** multi-threading to use them, but the ability to do so wasn't possible before `14.0.0` ([pyarrow-36709]) [pyarrow-36709]: https://github.com/apache/arrow/issues/36709 """ backend_version = self.compliant._backend_version if backend_version >= (14, 0) and grouped._use_threads: native = self.compliant.native grouped = pa.TableGroupBy(native, grouped.keys, use_threads=False) elif backend_version < (14, 0): # pragma: no cover msg = ( f"Using `{name}()` in a `group_by().agg(...)` context is only available in 'pyarrow>=14.0.0', " f"found version {requires._unparse_version(backend_version)!r}.\n\n" f"See https://github.com/apache/arrow/issues/36709" ) raise NotImplementedError(msg) return grouped, pc.ScalarAggregateOptions(skip_nulls=False) def agg(self, *exprs: ArrowExpr) -> ArrowDataFrame: self._ensure_all_simple(exprs) aggs: list[tuple[str, Aggregation, AggregateOptions | None]] = [] expected_pyarrow_column_names: list[str] = self._keys.copy() new_column_names: list[str] = self._keys.copy() exclude = (*self._keys, *self._output_key_names) grouped = self._grouped for expr in exprs: output_names, aliases = evaluate_output_names_and_aliases( expr, self.compliant, exclude ) md = expr._metadata op_nodes_reversed = list(md.op_nodes_reversed()) if len(op_nodes_reversed) == 1: # e.g. `agg(nw.len())` if op_nodes_reversed[0].name != "len": # pragma: no cover msg = "Safety assertion failed, please report a bug to https://github.com/narwhals-dev/narwhals/issues" raise AssertionError(msg) new_column_names.append(aliases[0]) expected_pyarrow_column_names.append(f"{self._keys[0]}_count") aggs.append((self._keys[0], "count", pc.CountOptions(mode="all"))) continue grouped, function_name, option = self._configure_agg(grouped, expr) new_column_names.extend(aliases) expected_pyarrow_column_names.extend( [f"{output_name}_{function_name}" for output_name in output_names] ) aggs.extend( [(output_name, function_name, option) for output_name in output_names] ) result_simple = grouped.aggregate(aggs) # Rename columns, being very careful expected_old_names_indices: dict[str, list[int]] = collections.defaultdict(list) for idx, item in enumerate(expected_pyarrow_column_names): expected_old_names_indices[item].append(idx) if not ( set(result_simple.column_names) == set(expected_pyarrow_column_names) and len(result_simple.column_names) == len(expected_pyarrow_column_names) ): # pragma: no cover msg = ( f"Safety assertion failed, expected {expected_pyarrow_column_names} " f"got {result_simple.column_names}, " "please report a bug at https://github.com/narwhals-dev/narwhals/issues" ) raise AssertionError(msg) index_map: list[int] = [ expected_old_names_indices[item].pop(0) for item in result_simple.column_names ] new_column_names = [new_column_names[i] for i in index_map] result_simple = result_simple.rename_columns(new_column_names) return self.compliant._with_native(result_simple).rename( dict(zip(self._keys, self._output_key_names)) ) def __iter__(self) -> Iterator[tuple[Any, ArrowDataFrame]]: col_token = generate_temporary_column_name( n_bytes=8, columns=self.compliant.columns ) null_token: str = "__null_token_value__" # noqa: S105 table = self.compliant.native it, separator_scalar = cast_to_comparable_string_types( *(table[key] for key in self._keys), separator="" ) # NOTE: stubs indicate `separator` must also be a `ChunkedArray` # Reality: `str` is fine concat_str: Incomplete = pc.binary_join_element_wise key_values = concat_str( *it, separator_scalar, null_handling="replace", null_replacement=null_token ) table = table.add_column(i=0, field_=col_token, column=key_values) for v in pc.unique(key_values): t = self.compliant._with_native( table.filter(pc.equal(table[col_token], v)).drop([col_token]) ) row = t.simple_select(*self._keys).row(0) yield ( tuple(extract_py_scalar(el) for el in row), t.simple_select(*self._df.columns), ) narwhals-2.11.0/narwhals/_arrow/namespace.py0000644000000000000000000002345713615410400015763 0ustar00from __future__ import annotations import operator from functools import reduce from itertools import chain from typing import TYPE_CHECKING, Literal import pyarrow as pa import pyarrow.compute as pc from narwhals._arrow.dataframe import ArrowDataFrame from narwhals._arrow.expr import ArrowExpr from narwhals._arrow.selectors import ArrowSelectorNamespace from narwhals._arrow.series import ArrowSeries from narwhals._arrow.utils import cast_to_comparable_string_types from narwhals._compliant import EagerNamespace from narwhals._expression_parsing import ( combine_alias_output_names, combine_evaluate_output_names, ) from narwhals._utils import Implementation if TYPE_CHECKING: from collections.abc import Iterator, Sequence from narwhals._arrow.typing import ChunkedArrayAny, Incomplete, ScalarAny from narwhals._utils import Version from narwhals.typing import IntoDType, NonNestedLiteral class ArrowNamespace( EagerNamespace[ArrowDataFrame, ArrowSeries, ArrowExpr, pa.Table, "ChunkedArrayAny"] ): _implementation = Implementation.PYARROW @property def _dataframe(self) -> type[ArrowDataFrame]: return ArrowDataFrame @property def _expr(self) -> type[ArrowExpr]: return ArrowExpr @property def _series(self) -> type[ArrowSeries]: return ArrowSeries def __init__(self, *, version: Version) -> None: self._version = version def extract_native( self, *series: ArrowSeries ) -> Iterator[ChunkedArrayAny | ScalarAny]: return (s.native[0] if s._broadcast else s.native for s in series) def len(self) -> ArrowExpr: # coverage bug? this is definitely hit return self._expr( # pragma: no cover lambda df: [ ArrowSeries.from_iterable([len(df.native)], name="len", context=self) ], evaluate_output_names=lambda _df: ["len"], alias_output_names=None, version=self._version, ) def lit(self, value: NonNestedLiteral, dtype: IntoDType | None) -> ArrowExpr: def _lit_arrow_series(_: ArrowDataFrame) -> ArrowSeries: arrow_series = ArrowSeries.from_iterable( data=[value], name="literal", context=self ) if dtype: return arrow_series.cast(dtype) return arrow_series return self._expr( lambda df: [_lit_arrow_series(df)], evaluate_output_names=lambda _df: ["literal"], alias_output_names=None, version=self._version, ) def all_horizontal(self, *exprs: ArrowExpr, ignore_nulls: bool) -> ArrowExpr: def func(df: ArrowDataFrame) -> list[ArrowSeries]: series: Iterator[ArrowSeries] = chain.from_iterable(e(df) for e in exprs) if ignore_nulls: series = (s.fill_null(True, None, None) for s in series) return [reduce(operator.and_, series)] return self._expr._from_callable( func=func, evaluate_output_names=combine_evaluate_output_names(*exprs), alias_output_names=combine_alias_output_names(*exprs), context=self, ) def any_horizontal(self, *exprs: ArrowExpr, ignore_nulls: bool) -> ArrowExpr: def func(df: ArrowDataFrame) -> list[ArrowSeries]: series: Iterator[ArrowSeries] = chain.from_iterable(e(df) for e in exprs) if ignore_nulls: series = (s.fill_null(False, None, None) for s in series) return [reduce(operator.or_, series)] return self._expr._from_callable( func=func, evaluate_output_names=combine_evaluate_output_names(*exprs), alias_output_names=combine_alias_output_names(*exprs), context=self, ) def sum_horizontal(self, *exprs: ArrowExpr) -> ArrowExpr: def func(df: ArrowDataFrame) -> list[ArrowSeries]: it = chain.from_iterable(expr(df) for expr in exprs) series = (s.fill_null(0, strategy=None, limit=None) for s in it) return [reduce(operator.add, series)] return self._expr._from_callable( func=func, evaluate_output_names=combine_evaluate_output_names(*exprs), alias_output_names=combine_alias_output_names(*exprs), context=self, ) def mean_horizontal(self, *exprs: ArrowExpr) -> ArrowExpr: int_64 = self._version.dtypes.Int64() def func(df: ArrowDataFrame) -> list[ArrowSeries]: expr_results = tuple(chain.from_iterable(expr(df) for expr in exprs)) series = [s.fill_null(0, strategy=None, limit=None) for s in expr_results] non_na = [1 - s.is_null().cast(int_64) for s in expr_results] return [reduce(operator.add, series) / reduce(operator.add, non_na)] return self._expr._from_callable( func=func, evaluate_output_names=combine_evaluate_output_names(*exprs), alias_output_names=combine_alias_output_names(*exprs), context=self, ) def min_horizontal(self, *exprs: ArrowExpr) -> ArrowExpr: def func(df: ArrowDataFrame) -> list[ArrowSeries]: init_series, *series = tuple(chain.from_iterable(expr(df) for expr in exprs)) native_series = reduce( pc.min_element_wise, [s.native for s in series], init_series.native ) return [ ArrowSeries(native_series, name=init_series.name, version=self._version) ] return self._expr._from_callable( func=func, evaluate_output_names=combine_evaluate_output_names(*exprs), alias_output_names=combine_alias_output_names(*exprs), context=self, ) def max_horizontal(self, *exprs: ArrowExpr) -> ArrowExpr: def func(df: ArrowDataFrame) -> list[ArrowSeries]: init_series, *series = tuple(chain.from_iterable(expr(df) for expr in exprs)) native_series = reduce( pc.max_element_wise, [s.native for s in series], init_series.native ) return [ ArrowSeries(native_series, name=init_series.name, version=self._version) ] return self._expr._from_callable( func=func, evaluate_output_names=combine_evaluate_output_names(*exprs), alias_output_names=combine_alias_output_names(*exprs), context=self, ) def _concat_diagonal(self, dfs: Sequence[pa.Table], /) -> pa.Table: if self._backend_version >= (14,): return pa.concat_tables(dfs, promote_options="default") return pa.concat_tables(dfs, promote=True) # pragma: no cover def _concat_horizontal(self, dfs: Sequence[pa.Table], /) -> pa.Table: names = list(chain.from_iterable(df.column_names for df in dfs)) arrays = tuple(chain.from_iterable(df.itercolumns() for df in dfs)) return pa.Table.from_arrays(arrays, names=names) def _concat_vertical(self, dfs: Sequence[pa.Table], /) -> pa.Table: cols_0 = dfs[0].column_names for i, df in enumerate(dfs[1:], start=1): cols_current = df.column_names if cols_current != cols_0: msg = ( "unable to vstack, column names don't match:\n" f" - dataframe 0: {cols_0}\n" f" - dataframe {i}: {cols_current}\n" ) raise TypeError(msg) return pa.concat_tables(dfs) @property def selectors(self) -> ArrowSelectorNamespace: return ArrowSelectorNamespace.from_namespace(self) def concat_str( self, *exprs: ArrowExpr, separator: str, ignore_nulls: bool ) -> ArrowExpr: def func(df: ArrowDataFrame) -> list[ArrowSeries]: series = list(chain.from_iterable(expr(df) for expr in exprs)) name = series[0].name null_handling: Literal["skip", "emit_null"] = ( "skip" if ignore_nulls else "emit_null" ) it, separator_scalar = cast_to_comparable_string_types( *self.extract_native(*series), separator=separator ) # NOTE: stubs indicate `separator` must also be a `ChunkedArray` # Reality: `str` is fine concat_str: Incomplete = pc.binary_join_element_wise compliant = self._series( concat_str(*it, separator_scalar, null_handling=null_handling), name=name, version=self._version, ) return [compliant] return self._expr._from_callable( func=func, evaluate_output_names=combine_evaluate_output_names(*exprs), alias_output_names=combine_alias_output_names(*exprs), context=self, ) def coalesce(self, *exprs: ArrowExpr) -> ArrowExpr: def func(df: ArrowDataFrame) -> list[ArrowSeries]: align = self._series._align_full_broadcast init_series, *series = align(*chain.from_iterable(expr(df) for expr in exprs)) return [ ArrowSeries( pc.coalesce(init_series.native, *(s.native for s in series)), name=init_series.name, version=self._version, ) ] return self._expr._from_callable( func=func, evaluate_output_names=combine_evaluate_output_names(*exprs), alias_output_names=combine_alias_output_names(*exprs), context=self, ) def _if_then_else( self, when: ChunkedArrayAny, then: ChunkedArrayAny, otherwise: ChunkedArrayAny | None = None, ) -> ChunkedArrayAny: otherwise = pa.nulls(len(when), then.type) if otherwise is None else otherwise return pc.if_else(when, then, otherwise) narwhals-2.11.0/narwhals/_arrow/selectors.py0000644000000000000000000000156213615410400016023 0ustar00from __future__ import annotations from typing import TYPE_CHECKING from narwhals._arrow.expr import ArrowExpr from narwhals._compliant import CompliantSelector, EagerSelectorNamespace if TYPE_CHECKING: from narwhals._arrow.dataframe import ArrowDataFrame # noqa: F401 from narwhals._arrow.series import ArrowSeries # noqa: F401 class ArrowSelectorNamespace(EagerSelectorNamespace["ArrowDataFrame", "ArrowSeries"]): @property def _selector(self) -> type[ArrowSelector]: return ArrowSelector class ArrowSelector(CompliantSelector["ArrowDataFrame", "ArrowSeries"], ArrowExpr): # type: ignore[misc] def _to_expr(self) -> ArrowExpr: return ArrowExpr( self._call, evaluate_output_names=self._evaluate_output_names, alias_output_names=self._alias_output_names, version=self._version, ) narwhals-2.11.0/narwhals/_arrow/series.py0000644000000000000000000013216213615410400015313 0ustar00from __future__ import annotations from typing import TYPE_CHECKING, Any, Callable, Literal, cast, overload import pyarrow as pa import pyarrow.compute as pc from narwhals._arrow.series_cat import ArrowSeriesCatNamespace from narwhals._arrow.series_dt import ArrowSeriesDateTimeNamespace from narwhals._arrow.series_list import ArrowSeriesListNamespace from narwhals._arrow.series_str import ArrowSeriesStringNamespace from narwhals._arrow.series_struct import ArrowSeriesStructNamespace from narwhals._arrow.utils import ( cast_for_truediv, chunked_array, extract_native, floordiv_compat, is_array_or_scalar, lit, narwhals_to_native_dtype, native_to_narwhals_dtype, nulls_like, pad_series, zeros, ) from narwhals._compliant import EagerSeries, EagerSeriesHist from narwhals._typing_compat import assert_never from narwhals._utils import ( Implementation, generate_temporary_column_name, is_list_of, no_default, not_implemented, ) from narwhals.dependencies import is_numpy_array_1d from narwhals.exceptions import InvalidOperationError, ShapeError if TYPE_CHECKING: from collections.abc import Iterable, Iterator, Sequence from types import ModuleType import pandas as pd import polars as pl from typing_extensions import Self, TypeAlias, TypeIs from narwhals._arrow.dataframe import ArrowDataFrame from narwhals._arrow.namespace import ArrowNamespace from narwhals._arrow.typing import ( # type: ignore[attr-defined] ArrayAny, ArrayOrChunkedArray, ArrayOrScalar, ChunkedArrayAny, Incomplete, NullPlacement, Order, ScalarAny, TieBreaker, _AsPyType, _BasicDataType, ) from narwhals._compliant.series import HistData from narwhals._typing import NoDefault from narwhals._utils import Version, _LimitedContext from narwhals.dtypes import DType from narwhals.typing import ( ClosedInterval, FillNullStrategy, Into1DArray, IntoDType, ModeKeepStrategy, NonNestedLiteral, PythonLiteral, RankMethod, RollingInterpolationMethod, SizedMultiIndexSelector, _1DArray, _2DArray, _SliceIndex, ) ArrowHistData: TypeAlias = ( "HistData[ChunkedArrayAny, list[ScalarAny] | pa.Int64Array | list[float]]" ) # TODO @dangotbanned: move into `_arrow.utils` # Lots of modules are importing inline @overload def maybe_extract_py_scalar( value: pa.Scalar[_BasicDataType[_AsPyType]], return_py_scalar: bool, # noqa: FBT001 ) -> _AsPyType: ... @overload def maybe_extract_py_scalar( value: pa.Scalar[pa.StructType], return_py_scalar: bool, # noqa: FBT001 ) -> list[dict[str, Any]]: ... @overload def maybe_extract_py_scalar( value: pa.Scalar[pa.ListType[_BasicDataType[_AsPyType]]], return_py_scalar: bool, # noqa: FBT001 ) -> list[_AsPyType]: ... @overload def maybe_extract_py_scalar( value: pa.Scalar[Any] | Any, return_py_scalar: bool, # noqa: FBT001 ) -> Any: ... def maybe_extract_py_scalar(value: Any, return_py_scalar: bool) -> Any: # noqa: FBT001 if TYPE_CHECKING: return value.as_py() if return_py_scalar: return getattr(value, "as_py", lambda: value)() return value class ArrowSeries(EagerSeries["ChunkedArrayAny"]): _implementation = Implementation.PYARROW def __init__( self, native_series: ChunkedArrayAny, *, name: str, version: Version ) -> None: self._name = name self._native_series: ChunkedArrayAny = native_series self._version = version self._broadcast = False @property def native(self) -> ChunkedArrayAny: return self._native_series def _with_version(self, version: Version) -> Self: return self.__class__(self.native, name=self._name, version=version) def _with_native( self, series: ArrayOrScalar, *, preserve_broadcast: bool = False ) -> Self: result = self.from_native(chunked_array(series), name=self.name, context=self) if preserve_broadcast: result._broadcast = self._broadcast return result def _with_binary(self, op: Callable[..., ArrayOrScalar], other: Any) -> Self: ser, other_native = extract_native(self, other) preserve_broadcast = self._broadcast and getattr(other, "_broadcast", True) return self._with_native( op(ser, other_native), preserve_broadcast=preserve_broadcast ).alias(self.name) def _with_binary_right(self, op: Callable[..., ArrayOrScalar], other: Any) -> Self: return self._with_binary(lambda x, y: op(y, x), other).alias(self.name) @classmethod def from_iterable( cls, data: Iterable[Any], *, context: _LimitedContext, name: str = "", dtype: IntoDType | None = None, ) -> Self: version = context._version if dtype is not None: dtype_pa: pa.DataType | None = narwhals_to_native_dtype(dtype, version) if is_array_or_scalar(data): data = data.cast(dtype_pa) dtype_pa = None native = data if cls._is_native(data) else chunked_array([data], dtype_pa) else: native = chunked_array([data]) return cls.from_native(native, context=context, name=name) def _from_scalar(self, value: Any) -> Self: if hasattr(value, "as_py"): value = value.as_py() return super()._from_scalar(value) @staticmethod def _is_native(obj: ChunkedArrayAny | Any) -> TypeIs[ChunkedArrayAny]: return isinstance(obj, pa.ChunkedArray) @classmethod def from_native( cls, data: ChunkedArrayAny, /, *, context: _LimitedContext, name: str = "" ) -> Self: return cls(data, version=context._version, name=name) @classmethod def from_numpy(cls, data: Into1DArray, /, *, context: _LimitedContext) -> Self: return cls.from_iterable( data if is_numpy_array_1d(data) else [data], context=context ) @classmethod def _align_full_broadcast(cls, *series: Self) -> Sequence[Self]: lengths = [len(s) for s in series] target_length = max( length for length, s in zip(lengths, series) if not s._broadcast ) fast_path = all(_len == target_length for _len in lengths) if fast_path: return series reshaped = [] for s in series: if s._broadcast: compliant = s._with_native(pa.repeat(s.native[0], target_length)) elif (actual_len := len(s)) != target_length: msg = f"Expected object of length {target_length}, got {actual_len}." raise ShapeError(msg) else: compliant = s reshaped.append(compliant) return reshaped def __narwhals_namespace__(self) -> ArrowNamespace: from narwhals._arrow.namespace import ArrowNamespace return ArrowNamespace(version=self._version) def __eq__(self, other: object) -> Self: # type: ignore[override] return self._with_binary(pc.equal, other) def __ne__(self, other: object) -> Self: # type: ignore[override] return self._with_binary(pc.not_equal, other) def __ge__(self, other: Any) -> Self: return self._with_binary(pc.greater_equal, other) def __gt__(self, other: Any) -> Self: return self._with_binary(pc.greater, other) def __le__(self, other: Any) -> Self: return self._with_binary(pc.less_equal, other) def __lt__(self, other: Any) -> Self: return self._with_binary(pc.less, other) def __and__(self, other: Any) -> Self: return self._with_binary(pc.and_kleene, other) def __rand__(self, other: Any) -> Self: return self._with_binary_right(pc.and_kleene, other) def __or__(self, other: Any) -> Self: return self._with_binary_right(pc.or_kleene, other) def __ror__(self, other: Any) -> Self: return self._with_binary_right(pc.or_kleene, other) def __add__(self, other: Any) -> Self: return self._with_binary(pc.add, other) def __radd__(self, other: Any) -> Self: return self._with_binary_right(pc.add, other) def __sub__(self, other: Any) -> Self: return self._with_binary(pc.subtract, other) def __rsub__(self, other: Any) -> Self: return self._with_binary_right(pc.subtract, other) def __mul__(self, other: Any) -> Self: return self._with_binary(pc.multiply, other) def __rmul__(self, other: Any) -> Self: return self._with_binary_right(pc.multiply, other) def __pow__(self, other: Any) -> Self: return self._with_binary(pc.power, other) def __rpow__(self, other: Any) -> Self: return self._with_binary_right(pc.power, other) def __floordiv__(self, other: Any) -> Self: return self._with_binary(floordiv_compat, other) def __rfloordiv__(self, other: Any) -> Self: return self._with_binary_right(floordiv_compat, other) def __truediv__(self, other: Any) -> Self: return self._with_binary(lambda x, y: pc.divide(*cast_for_truediv(x, y)), other) def __rtruediv__(self, other: Any) -> Self: return self._with_binary_right( lambda x, y: pc.divide(*cast_for_truediv(x, y)), other ) def __mod__(self, other: Any) -> Self: preserve_broadcast = self._broadcast and getattr(other, "_broadcast", True) floor_div = (self // other).native ser, other = extract_native(self, other) res = pc.subtract(ser, pc.multiply(floor_div, other)) return self._with_native(res, preserve_broadcast=preserve_broadcast) def __rmod__(self, other: Any) -> Self: preserve_broadcast = self._broadcast and getattr(other, "_broadcast", True) floor_div = (other // self).native ser, other = extract_native(self, other) res = pc.subtract(other, pc.multiply(floor_div, ser)) return self._with_native(res, preserve_broadcast=preserve_broadcast) def __invert__(self) -> Self: return self._with_native(pc.invert(self.native), preserve_broadcast=True) @property def _type(self) -> pa.DataType: return self.native.type def len(self, *, _return_py_scalar: bool = True) -> int: return maybe_extract_py_scalar(len(self.native), _return_py_scalar) def filter(self, predicate: ArrowSeries | list[bool | None]) -> Self: other_native: Any if not is_list_of(predicate, bool): _, other_native = extract_native(self, predicate) else: other_native = predicate return self._with_native(self.native.filter(other_native)) def first(self, *, _return_py_scalar: bool = True) -> PythonLiteral: result = self.native[0] if len(self.native) else None return maybe_extract_py_scalar(result, _return_py_scalar) def last(self, *, _return_py_scalar: bool = True) -> PythonLiteral: ca = self.native result = ca[height - 1] if (height := len(ca)) else None return maybe_extract_py_scalar(result, _return_py_scalar) def mean(self, *, _return_py_scalar: bool = True) -> float: return maybe_extract_py_scalar(pc.mean(self.native), _return_py_scalar) def median(self, *, _return_py_scalar: bool = True) -> float: if not self.dtype.is_numeric(): msg = "`median` operation not supported for non-numeric input type." raise InvalidOperationError(msg) return maybe_extract_py_scalar( pc.approximate_median(self.native), _return_py_scalar ) def min(self, *, _return_py_scalar: bool = True) -> Any: return maybe_extract_py_scalar(pc.min(self.native), _return_py_scalar) def max(self, *, _return_py_scalar: bool = True) -> Any: return maybe_extract_py_scalar(pc.max(self.native), _return_py_scalar) def arg_min(self, *, _return_py_scalar: bool = True) -> int: index_min = pc.index(self.native, pc.min(self.native)) return maybe_extract_py_scalar(index_min, _return_py_scalar) def arg_max(self, *, _return_py_scalar: bool = True) -> int: index_max = pc.index(self.native, pc.max(self.native)) return maybe_extract_py_scalar(index_max, _return_py_scalar) def sum(self, *, _return_py_scalar: bool = True) -> float: return maybe_extract_py_scalar( pc.sum(self.native, min_count=0), _return_py_scalar ) def drop_nulls(self) -> Self: return self._with_native(self.native.drop_null()) def shift(self, n: int) -> Self: if n > 0: arrays = [nulls_like(n, self), *self.native[:-n].chunks] elif n < 0: arrays = [*self.native[-n:].chunks, nulls_like(-n, self)] else: return self._with_native(self.native) return self._with_native(pa.concat_arrays(arrays)) def std(self, *, ddof: int, _return_py_scalar: bool = True) -> float: return maybe_extract_py_scalar( pc.stddev(self.native, ddof=ddof), _return_py_scalar ) def var(self, *, ddof: int, _return_py_scalar: bool = True) -> float: return maybe_extract_py_scalar( pc.variance(self.native, ddof=ddof), _return_py_scalar ) def skew(self, *, _return_py_scalar: bool = True) -> float | None: ser_not_null = self.native.drop_null() if len(ser_not_null) == 0: return None if len(ser_not_null) == 1: return float("nan") if len(ser_not_null) == 2: return 0.0 m = pc.subtract(ser_not_null, pc.mean(ser_not_null)) m2 = pc.mean(pc.power(m, lit(2))) m3 = pc.mean(pc.power(m, lit(3))) biased_population_skewness = pc.divide(m3, pc.power(m2, lit(1.5))) return maybe_extract_py_scalar(biased_population_skewness, _return_py_scalar) def kurtosis(self, *, _return_py_scalar: bool = True) -> float | None: ser_not_null = self.native.drop_null() if len(ser_not_null) == 0: return None if len(ser_not_null) == 1: return float("nan") m = pc.subtract(ser_not_null, pc.mean(ser_not_null)) m2 = pc.mean(pc.power(m, lit(2))) m4 = pc.mean(pc.power(m, lit(4))) k = pc.subtract(pc.divide(m4, pc.power(m2, lit(2))), lit(3)) return maybe_extract_py_scalar(k, _return_py_scalar) def count(self, *, _return_py_scalar: bool = True) -> int: return maybe_extract_py_scalar(pc.count(self.native), _return_py_scalar) def n_unique(self, *, _return_py_scalar: bool = True) -> int: return maybe_extract_py_scalar( pc.count(self.native.unique(), mode="all"), _return_py_scalar ) def __native_namespace__(self) -> ModuleType: if self._implementation is Implementation.PYARROW: return self._implementation.to_native_namespace() msg = f"Expected pyarrow, got: {type(self._implementation)}" # pragma: no cover raise AssertionError(msg) @property def name(self) -> str: return self._name def _gather(self, rows: SizedMultiIndexSelector[ChunkedArrayAny]) -> Self: if len(rows) == 0: return self._with_native(self.native.slice(0, 0)) if self._backend_version < (18,) and isinstance(rows, tuple): rows = list(rows) return self._with_native(self.native.take(rows)) def _gather_slice(self, rows: _SliceIndex | range) -> Self: start = rows.start or 0 stop = rows.stop if rows.stop is not None else len(self.native) if start < 0: start = len(self.native) + start if stop < 0: stop = len(self.native) + stop if rows.step is not None and rows.step != 1: msg = "Slicing with step is not supported on PyArrow tables" raise NotImplementedError(msg) return self._with_native(self.native.slice(start, stop - start)) def scatter(self, indices: int | Sequence[int], values: Any) -> Self: import numpy as np # ignore-banned-import values_native: ArrayAny if isinstance(indices, int): indices_native = pa.array([indices]) values_native = pa.array([values]) else: # TODO(unassigned): we may also want to let `indices` be a Series. # https://github.com/narwhals-dev/narwhals/issues/2155 indices_native = pa.array(indices) if isinstance(values, self.__class__): values_native = values.native.combine_chunks() else: # NOTE: Requires fixes in https://github.com/zen-xu/pyarrow-stubs/pull/209 pa_array: Incomplete = pa.array values_native = pa_array(values) sorting_indices = pc.sort_indices(indices_native) indices_native = indices_native.take(sorting_indices) values_native = values_native.take(sorting_indices) mask: _1DArray = np.zeros(self.len(), dtype=bool) mask[indices_native] = True # NOTE: Multiple issues # - Missing `values` type # - `mask` accepts a `np.ndarray`, but not mentioned in stubs # - Missing `replacements` type # - Missing return type pc_replace_with_mask: Incomplete = pc.replace_with_mask return self._with_native(pc_replace_with_mask(self.native, mask, values_native)) def to_list(self) -> list[Any]: return self.native.to_pylist() def __array__(self, dtype: Any = None, *, copy: bool | None = None) -> _1DArray: return self.native.__array__(dtype=dtype, copy=copy) def to_numpy(self, dtype: Any = None, *, copy: bool | None = None) -> _1DArray: return self.native.to_numpy() def alias(self, name: str) -> Self: ret = self.__class__(self.native, name=name, version=self._version) ret._broadcast = self._broadcast return ret @property def dtype(self) -> DType: return native_to_narwhals_dtype(self.native.type, self._version) def abs(self) -> Self: return self._with_native(pc.abs(self.native)) def cum_sum(self, *, reverse: bool) -> Self: cum_sum = pc.cumulative_sum result = ( cum_sum(self.native, skip_nulls=True) if not reverse else cum_sum(self.native[::-1], skip_nulls=True)[::-1] ) return self._with_native(result) def round(self, decimals: int) -> Self: return self._with_native( pc.round(self.native, decimals, round_mode="half_towards_infinity") ) def floor(self) -> Self: return self._with_native(pc.floor(self.native)) def ceil(self) -> Self: return self._with_native(pc.ceil(self.native)) def diff(self) -> Self: return self._with_native(pc.pairwise_diff(self.native.combine_chunks())) def any(self, *, _return_py_scalar: bool = True) -> bool: return maybe_extract_py_scalar( pc.any(self.native, min_count=0), _return_py_scalar ) def all(self, *, _return_py_scalar: bool = True) -> bool: return maybe_extract_py_scalar( pc.all(self.native, min_count=0), _return_py_scalar ) def is_between( self, lower_bound: Any, upper_bound: Any, closed: ClosedInterval ) -> Self: _, lower_bound = extract_native(self, lower_bound) _, upper_bound = extract_native(self, upper_bound) if closed == "left": ge = pc.greater_equal(self.native, lower_bound) lt = pc.less(self.native, upper_bound) res = pc.and_kleene(ge, lt) elif closed == "right": gt = pc.greater(self.native, lower_bound) le = pc.less_equal(self.native, upper_bound) res = pc.and_kleene(gt, le) elif closed == "none": gt = pc.greater(self.native, lower_bound) lt = pc.less(self.native, upper_bound) res = pc.and_kleene(gt, lt) elif closed == "both": ge = pc.greater_equal(self.native, lower_bound) le = pc.less_equal(self.native, upper_bound) res = pc.and_kleene(ge, le) else: assert_never(closed) return self._with_native(res) def is_null(self) -> Self: return self._with_native(self.native.is_null(), preserve_broadcast=True) def is_nan(self) -> Self: return self._with_native(pc.is_nan(self.native), preserve_broadcast=True) def cast(self, dtype: IntoDType) -> Self: data_type = narwhals_to_native_dtype(dtype, self._version) return self._with_native(pc.cast(self.native, data_type), preserve_broadcast=True) def null_count(self, *, _return_py_scalar: bool = True) -> int: return maybe_extract_py_scalar(self.native.null_count, _return_py_scalar) def head(self, n: int) -> Self: if n >= 0: return self._with_native(self.native.slice(0, n)) num_rows = len(self) return self._with_native(self.native.slice(0, max(0, num_rows + n))) def tail(self, n: int) -> Self: if n >= 0: num_rows = len(self) return self._with_native(self.native.slice(max(0, num_rows - n))) return self._with_native(self.native.slice(abs(n))) def is_in(self, other: Any) -> Self: if self._is_native(other): value_set: ArrayOrChunkedArray = other else: value_set = pa.array(other) return self._with_native(pc.is_in(self.native, value_set=value_set)) def arg_true(self) -> Self: import numpy as np # ignore-banned-import res = np.flatnonzero(self.native) return self.from_iterable(res, name=self.name, context=self) def item(self, index: int | None = None) -> Any: if index is None: if len(self) != 1: msg = ( "can only call '.item()' if the Series is of length 1," f" or an explicit index is provided (Series is of length {len(self)})" ) raise ValueError(msg) return maybe_extract_py_scalar(self.native[0], return_py_scalar=True) return maybe_extract_py_scalar(self.native[index], return_py_scalar=True) def value_counts( self, *, sort: bool, parallel: bool, name: str | None, normalize: bool ) -> ArrowDataFrame: """Parallel is unused, exists for compatibility.""" from narwhals._arrow.dataframe import ArrowDataFrame index_name_ = "index" if self._name is None else self._name value_name_ = name or ("proportion" if normalize else "count") val_counts = pc.value_counts(self.native) values = val_counts.field("values") counts = cast("ChunkedArrayAny", val_counts.field("counts")) if normalize: arrays = [values, pc.divide(*cast_for_truediv(counts, pc.sum(counts)))] else: arrays = [values, counts] val_count = pa.Table.from_arrays(arrays, names=[index_name_, value_name_]) if sort: val_count = val_count.sort_by([(value_name_, "descending")]) return ArrowDataFrame( val_count, version=self._version, validate_column_names=True ) def zip_with(self, mask: Self, other: Self) -> Self: cond = mask.native.combine_chunks() return self._with_native(pc.if_else(cond, self.native, other.native)) def sample( self, n: int | None, *, fraction: float | None, with_replacement: bool, seed: int | None, ) -> Self: import numpy as np # ignore-banned-import num_rows = len(self) if n is None and fraction is not None: n = int(num_rows * fraction) rng = np.random.default_rng(seed=seed) idx = np.arange(num_rows) mask = rng.choice(idx, size=n, replace=with_replacement) return self._with_native(self.native.take(mask)) def fill_nan(self, value: float | None) -> Self: result = pc.if_else(pc.is_nan(self.native), value, self.native) return self._with_native(result, preserve_broadcast=True) def fill_null( self, value: Self | NonNestedLiteral, strategy: FillNullStrategy | None, limit: int | None, ) -> Self: import numpy as np # ignore-banned-import def fill_aux( arr: ChunkedArrayAny, limit: int, direction: FillNullStrategy | None ) -> ArrayAny: # this algorithm first finds the indices of the valid values to fill all the null value positions # then it calculates the distance of each new index and the original index # if the distance is equal to or less than the limit and the original value is null, it is replaced valid_mask = pc.is_valid(arr) indices = pa.array(np.arange(len(arr)), type=pa.int64()) if direction == "forward": valid_index = np.maximum.accumulate(np.where(valid_mask, indices, -1)) distance = indices - valid_index else: valid_index = np.minimum.accumulate( np.where(valid_mask[::-1], indices[::-1], len(arr)) )[::-1] distance = valid_index - indices return pc.if_else( pc.and_(pc.is_null(arr), pc.less_equal(distance, lit(limit))), # pyright: ignore[reportArgumentType, reportCallIssue] arr.take(valid_index), arr, ) if value is not None: _, native_value = extract_native(self, value) series: ArrayOrScalar = pc.fill_null(self.native, native_value) elif limit is None: fill_func = ( pc.fill_null_forward if strategy == "forward" else pc.fill_null_backward ) series = fill_func(self.native) else: series = fill_aux(self.native, limit, strategy) return self._with_native(series, preserve_broadcast=True) def to_frame(self) -> ArrowDataFrame: from narwhals._arrow.dataframe import ArrowDataFrame df = pa.Table.from_arrays([self.native], names=[self.name]) return ArrowDataFrame(df, version=self._version, validate_column_names=False) def to_pandas(self) -> pd.Series[Any]: import pandas as pd # ignore-banned-import() return pd.Series(self.native, name=self.name) def to_polars(self) -> pl.Series: import polars as pl # ignore-banned-import return cast("pl.Series", pl.from_arrow(self.native)) def is_unique(self) -> ArrowSeries: return self.to_frame().is_unique().alias(self.name) def is_first_distinct(self) -> Self: import numpy as np # ignore-banned-import row_number = pa.array(np.arange(len(self))) col_token = generate_temporary_column_name(n_bytes=8, columns=[self.name]) first_distinct_index = ( pa.Table.from_arrays([self.native], names=[self.name]) .append_column(col_token, row_number) .group_by(self.name) .aggregate([(col_token, "min")]) .column(f"{col_token}_min") ) return self._with_native(pc.is_in(row_number, first_distinct_index)) def is_last_distinct(self) -> Self: import numpy as np # ignore-banned-import row_number = pa.array(np.arange(len(self))) col_token = generate_temporary_column_name(n_bytes=8, columns=[self.name]) last_distinct_index = ( pa.Table.from_arrays([self.native], names=[self.name]) .append_column(col_token, row_number) .group_by(self.name) .aggregate([(col_token, "max")]) .column(f"{col_token}_max") ) return self._with_native(pc.is_in(row_number, last_distinct_index)) def is_sorted(self, *, descending: bool) -> bool: if not isinstance(descending, bool): msg = f"argument 'descending' should be boolean, found {type(descending)}" raise TypeError(msg) if descending: result = pc.all(pc.greater_equal(self.native[:-1], self.native[1:])) else: result = pc.all(pc.less_equal(self.native[:-1], self.native[1:])) return maybe_extract_py_scalar(result, return_py_scalar=True) def unique(self, *, maintain_order: bool = True) -> Self: # TODO(marco): `pc.unique` seems to always maintain order, is that guaranteed? return self._with_native(self.native.unique()) def replace_strict( self, default: Any | NoDefault, old: Sequence[Any], new: Sequence[Any], *, return_dtype: IntoDType | None, ) -> Self: # https://stackoverflow.com/a/79111029/4451315 idxs = pc.index_in(self.native, pa.array(old)) # Tracks which values were found in the mapping (even if mapped to None) was_matched = pc.is_valid(idxs) result_native = pc.take(pa.array(new), idxs) if return_dtype is not None: result_native = result_native.cast( narwhals_to_native_dtype(return_dtype, self._version) ) result = self._with_native(result_native) if default is no_default: # Check that all non-null input values were matched # (result may have more nulls if mapping contains {value: None}) unmatched_mask = pc.and_(pc.is_valid(self.native), pc.invert(was_matched)) if maybe_extract_py_scalar( pc.any(unmatched_mask, min_count=0), return_py_scalar=True ): unmatched_values = ( self.filter(self._with_native(unmatched_mask)) .unique(maintain_order=False) .to_list() ) msg = ( "replace_strict did not replace all non-null values.\n\n" f"The following did not get replaced: {unmatched_values}" ) raise InvalidOperationError(msg) else: result_native, default = extract_native(result, default) # Only fill with default where the value wasn't matched (not where result is null due to mapping) # If was_matched, keep result.native; otherwise use default result = self._with_native(pc.if_else(was_matched, result.native, default)) return result def sort(self, *, descending: bool, nulls_last: bool) -> Self: order: Order = "descending" if descending else "ascending" null_placement: NullPlacement = "at_end" if nulls_last else "at_start" sorted_indices = pc.array_sort_indices( self.native, order=order, null_placement=null_placement ) return self._with_native(self.native.take(sorted_indices)) def to_dummies(self, *, separator: str, drop_first: bool) -> ArrowDataFrame: import numpy as np # ignore-banned-import from narwhals._arrow.dataframe import ArrowDataFrame name = self._name # NOTE: stub is missing attributes (https://arrow.apache.org/docs/python/generated/pyarrow.DictionaryArray.html) da: Incomplete = self.native.combine_chunks().dictionary_encode("encode") columns: _2DArray = np.zeros((len(da.dictionary), len(da)), np.int8) columns[da.indices, np.arange(len(da))] = 1 null_col_pa, null_col_pl = f"{name}{separator}None", f"{name}{separator}null" cols = [ {null_col_pa: null_col_pl}.get( f"{name}{separator}{v}", f"{name}{separator}{v}" ) for v in da.dictionary ] output_order = ( [ null_col_pl, *sorted([c for c in cols if c != null_col_pl])[int(drop_first) :], ] if null_col_pl in cols else sorted(cols)[int(drop_first) :] ) return ArrowDataFrame( pa.Table.from_arrays(columns, names=cols), version=self._version, validate_column_names=True, ).simple_select(*output_order) def quantile( self, quantile: float, interpolation: RollingInterpolationMethod, *, _return_py_scalar: bool = True, ) -> float: return maybe_extract_py_scalar( pc.quantile(self.native, q=quantile, interpolation=interpolation)[0], _return_py_scalar, ) def gather_every(self, n: int, offset: int = 0) -> Self: return self._with_native(self.native[offset::n]) def clip(self, lower_bound: Self, upper_bound: Self) -> Self: _, lower = extract_native(self, lower_bound) _, upper = extract_native(self, upper_bound) return self._with_native( pc.max_element_wise(pc.min_element_wise(self.native, upper), lower) ) def clip_lower(self, lower_bound: Self) -> Self: _, lower = extract_native(self, lower_bound) return self._with_native(pc.max_element_wise(self.native, lower)) def clip_upper(self, upper_bound: Self) -> Self: _, upper = extract_native(self, upper_bound) return self._with_native(pc.min_element_wise(self.native, upper)) def to_arrow(self) -> ArrayAny: return self.native.combine_chunks() def mode(self, *, keep: ModeKeepStrategy) -> ArrowSeries: plx = self.__narwhals_namespace__() col_token = generate_temporary_column_name(n_bytes=8, columns=[self.name]) counts = self.value_counts( name=col_token, normalize=False, sort=False, parallel=False ) result = counts.filter( plx.col(col_token) == plx.col(col_token).max().broadcast() ).get_column(self.name) return result.head(1) if keep == "any" else result def is_finite(self) -> Self: return self._with_native(pc.is_finite(self.native)) def cum_count(self, *, reverse: bool) -> Self: dtypes = self._version.dtypes return (~self.is_null()).cast(dtypes.UInt32()).cum_sum(reverse=reverse) def cum_min(self, *, reverse: bool) -> Self: result = ( pc.cumulative_min(self.native, skip_nulls=True) if not reverse else pc.cumulative_min(self.native[::-1], skip_nulls=True)[::-1] ) return self._with_native(result) def cum_max(self, *, reverse: bool) -> Self: result = ( pc.cumulative_max(self.native, skip_nulls=True) if not reverse else pc.cumulative_max(self.native[::-1], skip_nulls=True)[::-1] ) return self._with_native(result) def cum_prod(self, *, reverse: bool) -> Self: result = ( pc.cumulative_prod(self.native, skip_nulls=True) if not reverse else pc.cumulative_prod(self.native[::-1], skip_nulls=True)[::-1] ) return self._with_native(result) def rolling_sum(self, window_size: int, *, min_samples: int, center: bool) -> Self: min_samples = min_samples if min_samples is not None else window_size padded_series, offset = pad_series(self, window_size=window_size, center=center) cum_sum = padded_series.cum_sum(reverse=False).fill_null( value=None, strategy="forward", limit=None ) rolling_sum = ( cum_sum - cum_sum.shift(window_size).fill_null(value=0, strategy=None, limit=None) if window_size != 0 else cum_sum ) valid_count = padded_series.cum_count(reverse=False) count_in_window = valid_count - valid_count.shift(window_size).fill_null( value=0, strategy=None, limit=None ) result = self._with_native( pc.if_else((count_in_window >= min_samples).native, rolling_sum.native, None) ) return result._gather_slice(slice(offset, None)) def rolling_mean(self, window_size: int, *, min_samples: int, center: bool) -> Self: min_samples = min_samples if min_samples is not None else window_size padded_series, offset = pad_series(self, window_size=window_size, center=center) cum_sum = padded_series.cum_sum(reverse=False).fill_null( value=None, strategy="forward", limit=None ) rolling_sum = ( cum_sum - cum_sum.shift(window_size).fill_null(value=0, strategy=None, limit=None) if window_size != 0 else cum_sum ) valid_count = padded_series.cum_count(reverse=False) count_in_window = valid_count - valid_count.shift(window_size).fill_null( value=0, strategy=None, limit=None ) result = ( self._with_native( pc.if_else( (count_in_window >= min_samples).native, rolling_sum.native, None ) ) / count_in_window ) return result._gather_slice(slice(offset, None)) def rolling_var( self, window_size: int, *, min_samples: int, center: bool, ddof: int ) -> Self: min_samples = min_samples if min_samples is not None else window_size padded_series, offset = pad_series(self, window_size=window_size, center=center) cum_sum = padded_series.cum_sum(reverse=False).fill_null( value=None, strategy="forward", limit=None ) rolling_sum = ( cum_sum - cum_sum.shift(window_size).fill_null(value=0, strategy=None, limit=None) if window_size != 0 else cum_sum ) cum_sum_sq = ( pow(padded_series, 2) .cum_sum(reverse=False) .fill_null(value=None, strategy="forward", limit=None) ) rolling_sum_sq = ( cum_sum_sq - cum_sum_sq.shift(window_size).fill_null(value=0, strategy=None, limit=None) if window_size != 0 else cum_sum_sq ) valid_count = padded_series.cum_count(reverse=False) count_in_window = valid_count - valid_count.shift(window_size).fill_null( value=0, strategy=None, limit=None ) result = self._with_native( pc.if_else( (count_in_window >= min_samples).native, (rolling_sum_sq - (rolling_sum**2 / count_in_window)).native, None, ) ) / self._with_native(pc.max_element_wise((count_in_window - ddof).native, 0)) return result._gather_slice(slice(offset, None, None)) def rolling_std( self, window_size: int, *, min_samples: int, center: bool, ddof: int ) -> Self: return ( self.rolling_var( window_size=window_size, min_samples=min_samples, center=center, ddof=ddof ) ** 0.5 ) def rank(self, method: RankMethod, *, descending: bool) -> Self: if method == "average": msg = ( "`rank` with `method='average' is not supported for pyarrow backend. " "The available methods are {'min', 'max', 'dense', 'ordinal'}." ) raise ValueError(msg) sort_keys: Order = "descending" if descending else "ascending" tiebreaker: TieBreaker = "first" if method == "ordinal" else method native_series: ArrayOrChunkedArray if self._backend_version < (14, 0, 0): # pragma: no cover native_series = self.native.combine_chunks() else: native_series = self.native null_mask = pc.is_null(native_series) rank = pc.rank(native_series, sort_keys=sort_keys, tiebreaker=tiebreaker) result = pc.if_else(null_mask, lit(None, rank.type), rank) return self._with_native(result) def hist_from_bins( self, bins: list[float], *, include_breakpoint: bool ) -> ArrowDataFrame: return ( _ArrowHist.from_series(self, include_breakpoint=include_breakpoint) .with_bins(bins) .to_frame() ) def hist_from_bin_count( self, bin_count: int, *, include_breakpoint: bool ) -> ArrowDataFrame: return ( _ArrowHist.from_series(self, include_breakpoint=include_breakpoint) .with_bin_count(bin_count) .to_frame() ) def __iter__(self) -> Iterator[Any]: for x in self.native: yield maybe_extract_py_scalar(x, return_py_scalar=True) def __contains__(self, other: Any) -> bool: from pyarrow import ( ArrowInvalid, # ignore-banned-imports ArrowNotImplementedError, # ignore-banned-imports ArrowTypeError, # ignore-banned-imports ) try: other_ = lit(other) if other is not None else lit(None, type=self._type) return maybe_extract_py_scalar( pc.is_in(other_, self.native), return_py_scalar=True ) except (ArrowInvalid, ArrowNotImplementedError, ArrowTypeError) as exc: msg = f"Unable to compare other of type {type(other)} with series of type {self.dtype}." raise InvalidOperationError(msg) from exc def log(self, base: float) -> Self: return self._with_native(pc.logb(self.native, lit(base))) def exp(self) -> Self: return self._with_native(pc.exp(self.native)) def sqrt(self) -> Self: return self._with_native(pc.sqrt(self.native)) @property def dt(self) -> ArrowSeriesDateTimeNamespace: return ArrowSeriesDateTimeNamespace(self) @property def cat(self) -> ArrowSeriesCatNamespace: return ArrowSeriesCatNamespace(self) @property def str(self) -> ArrowSeriesStringNamespace: return ArrowSeriesStringNamespace(self) @property def list(self) -> ArrowSeriesListNamespace: return ArrowSeriesListNamespace(self) @property def struct(self) -> ArrowSeriesStructNamespace: return ArrowSeriesStructNamespace(self) ewm_mean = not_implemented() class _ArrowHist( EagerSeriesHist["ChunkedArrayAny", "list[ScalarAny] | pa.Int64Array | list[float]"] ): _series: ArrowSeries def to_frame(self) -> ArrowDataFrame: # NOTE: Constructor typing is too strict for `TypedDict` table: Incomplete = pa.Table.from_pydict from_native = self._series.__narwhals_namespace__()._dataframe.from_native return from_native(table(self._data), context=self._series) # NOTE: *Could* be handled at narwhals-level def is_empty_series(self) -> bool: # NOTE: `ChunkedArray.combine_chunks` returns the concrete array type # Stubs say `Array[pa.BooleanScalar]`, which is missing properties # https://github.com/zen-xu/pyarrow-stubs/blob/6bedee748bc74feb8513b24bf43d64b24c7fddc8/pyarrow-stubs/__lib_pxi/array.pyi#L2395-L2399 is_null = self.native.is_null(nan_is_null=True) arr = cast("pa.BooleanArray", is_null.combine_chunks()) return arr.false_count == 0 # NOTE: *Could* be handled at narwhals-level, **iff** we add `nw.repeat`, `nw.linear_space` # See https://github.com/narwhals-dev/narwhals/pull/2839#discussion_r2215630696 def series_empty(self, arg: int | list[float], /) -> ArrowHistData: count = self._zeros(arg) if self._breakpoint: return {"breakpoint": self._calculate_breakpoint(arg), "count": count} return {"count": count} def _zeros(self, arg: int | list[float], /) -> pa.Int64Array: return zeros(arg) if isinstance(arg, int) else zeros(len(arg) - 1) def _linear_space( self, start: float, end: float, num_samples: int, *, closed: Literal["both", "none"] = "both", ) -> _1DArray: from numpy import linspace # ignore-banned-import return linspace(start=start, stop=end, num=num_samples, endpoint=closed == "both") def _calculate_bins(self, bin_count: int) -> _1DArray: """Prepare bins for histogram calculation from bin_count.""" d = pc.min_max(self.native) lower, upper = d["min"].as_py(), d["max"].as_py() if lower == upper: lower -= 0.5 upper += 0.5 return self._linear_space(lower, upper, bin_count + 1) def _calculate_hist(self, bins: list[float] | _1DArray) -> ArrowHistData: ser = self.native # NOTE: `mypy` refuses to resolve `ndarray.__getitem__` # Previously annotated as `list[float]`, but # - wasn't accurate to how we implemented it # - `pa.scalar` overloads fail to match on `float | np.float64` (but runtime is fine) bins = cast("list[float]", bins) # Handle single bin case if len(bins) == 2: is_between_bins = pc.and_( pc.greater_equal(ser, lit(bins[0])), pc.less_equal(ser, lit(bins[1])) ) count = pc.sum(is_between_bins.cast(pa.uint8())) if self._breakpoint: return {"breakpoint": [bins[-1]], "count": [count]} return {"count": [count]} # Handle multiple bins import numpy as np # ignore-banned-import bin_indices = np.searchsorted(bins, ser, side="left") # lowest bin is inclusive bin_indices = pc.if_else(pc.equal(ser, lit(bins[0])), 1, bin_indices) # Align unique categories and counts appropriately obs_cats, obs_counts = np.unique(bin_indices, return_counts=True) obj_cats = np.arange(1, len(bins)) counts = np.zeros_like(obj_cats) counts[np.isin(obj_cats, obs_cats)] = obs_counts[np.isin(obs_cats, obj_cats)] if self._breakpoint: return {"breakpoint": bins[1:], "count": counts} return {"count": counts} narwhals-2.11.0/narwhals/_arrow/series_cat.py0000644000000000000000000000125613615410400016141 0ustar00from __future__ import annotations from typing import TYPE_CHECKING import pyarrow as pa from narwhals._arrow.utils import ArrowSeriesNamespace from narwhals._compliant.any_namespace import CatNamespace if TYPE_CHECKING: from narwhals._arrow.series import ArrowSeries from narwhals._arrow.typing import Incomplete class ArrowSeriesCatNamespace(ArrowSeriesNamespace, CatNamespace["ArrowSeries"]): def get_categories(self) -> ArrowSeries: # NOTE: Should be `list[pa.DictionaryArray]`, but `DictionaryArray` has no attributes chunks: Incomplete = self.native.chunks return self.with_native(pa.concat_arrays(x.dictionary for x in chunks).unique()) narwhals-2.11.0/narwhals/_arrow/series_dt.py0000644000000000000000000002150213615410400015775 0ustar00from __future__ import annotations from typing import TYPE_CHECKING, Any, Callable, ClassVar, cast import pyarrow as pa import pyarrow.compute as pc from narwhals._arrow.utils import UNITS_DICT, ArrowSeriesNamespace, floordiv_compat, lit from narwhals._compliant.any_namespace import DateTimeNamespace from narwhals._constants import ( MS_PER_MINUTE, MS_PER_SECOND, NS_PER_MICROSECOND, NS_PER_MILLISECOND, NS_PER_MINUTE, NS_PER_SECOND, SECONDS_PER_DAY, SECONDS_PER_MINUTE, US_PER_MINUTE, US_PER_SECOND, ) from narwhals._duration import Interval if TYPE_CHECKING: from collections.abc import Mapping from typing_extensions import TypeAlias from narwhals._arrow.series import ArrowSeries from narwhals._arrow.typing import ChunkedArrayAny, ScalarAny from narwhals.dtypes import Datetime from narwhals.typing import TimeUnit UnitCurrent: TypeAlias = TimeUnit UnitTarget: TypeAlias = TimeUnit BinOpBroadcast: TypeAlias = Callable[[ChunkedArrayAny, ScalarAny], ChunkedArrayAny] IntoRhs: TypeAlias = int class ArrowSeriesDateTimeNamespace( ArrowSeriesNamespace, DateTimeNamespace["ArrowSeries"] ): _TIMESTAMP_DATE_FACTOR: ClassVar[Mapping[TimeUnit, int]] = { "ns": NS_PER_SECOND, "us": US_PER_SECOND, "ms": MS_PER_SECOND, "s": 1, } _TIMESTAMP_DATETIME_OP_FACTOR: ClassVar[ Mapping[tuple[UnitCurrent, UnitTarget], tuple[BinOpBroadcast, IntoRhs]] ] = { ("ns", "us"): (floordiv_compat, 1_000), ("ns", "ms"): (floordiv_compat, 1_000_000), ("us", "ns"): (pc.multiply, NS_PER_MICROSECOND), ("us", "ms"): (floordiv_compat, 1_000), ("ms", "ns"): (pc.multiply, NS_PER_MILLISECOND), ("ms", "us"): (pc.multiply, 1_000), ("s", "ns"): (pc.multiply, NS_PER_SECOND), ("s", "us"): (pc.multiply, US_PER_SECOND), ("s", "ms"): (pc.multiply, MS_PER_SECOND), } @property def unit(self) -> TimeUnit: # NOTE: Unsafe (native). return cast("pa.TimestampType[TimeUnit, Any]", self.native.type).unit @property def time_zone(self) -> str | None: # NOTE: Unsafe (narwhals). return cast("Datetime", self.compliant.dtype).time_zone def to_string(self, format: str) -> ArrowSeries: # PyArrow differs from other libraries in that %S also prints out # the fractional part of the second...:'( # https://arrow.apache.org/docs/python/generated/pyarrow.compute.strftime.html format = format.replace("%S.%f", "%S").replace("%S%.f", "%S") return self.with_native(pc.strftime(self.native, format)) def replace_time_zone(self, time_zone: str | None) -> ArrowSeries: if time_zone is not None: result = pc.assume_timezone(pc.local_timestamp(self.native), time_zone) else: result = pc.local_timestamp(self.native) return self.with_native(result) def convert_time_zone(self, time_zone: str) -> ArrowSeries: ser = self.replace_time_zone("UTC") if self.time_zone is None else self.compliant return self.with_native(ser.native.cast(pa.timestamp(self.unit, time_zone))) def timestamp(self, time_unit: TimeUnit) -> ArrowSeries: ser = self.compliant dtypes = ser._version.dtypes if isinstance(ser.dtype, dtypes.Datetime): current = ser.dtype.time_unit s_cast = self.native.cast(pa.int64()) if current == time_unit: result = s_cast elif item := self._TIMESTAMP_DATETIME_OP_FACTOR.get((current, time_unit)): fn, factor = item result = fn(s_cast, lit(factor)) else: # pragma: no cover msg = f"unexpected time unit {current}, please report an issue at https://github.com/narwhals-dev/narwhals" raise AssertionError(msg) return self.with_native(result) if isinstance(ser.dtype, dtypes.Date): time_s = pc.multiply(self.native.cast(pa.int32()), lit(SECONDS_PER_DAY)) factor = self._TIMESTAMP_DATE_FACTOR[time_unit] return self.with_native(pc.multiply(time_s, lit(factor))) msg = "Input should be either of Date or Datetime type" raise TypeError(msg) def date(self) -> ArrowSeries: return self.with_native(self.native.cast(pa.date32())) def year(self) -> ArrowSeries: return self.with_native(pc.year(self.native)) def month(self) -> ArrowSeries: return self.with_native(pc.month(self.native)) def day(self) -> ArrowSeries: return self.with_native(pc.day(self.native)) def hour(self) -> ArrowSeries: return self.with_native(pc.hour(self.native)) def minute(self) -> ArrowSeries: return self.with_native(pc.minute(self.native)) def second(self) -> ArrowSeries: return self.with_native(pc.second(self.native)) def millisecond(self) -> ArrowSeries: return self.with_native(pc.millisecond(self.native)) def microsecond(self) -> ArrowSeries: arr = self.native result = pc.add(pc.multiply(pc.millisecond(arr), lit(1000)), pc.microsecond(arr)) return self.with_native(result) def nanosecond(self) -> ArrowSeries: result = pc.add( pc.multiply(self.microsecond().native, lit(1000)), pc.nanosecond(self.native) ) return self.with_native(result) def ordinal_day(self) -> ArrowSeries: return self.with_native(pc.day_of_year(self.native)) def weekday(self) -> ArrowSeries: return self.with_native(pc.day_of_week(self.native, count_from_zero=False)) def total_minutes(self) -> ArrowSeries: unit_to_minutes_factor = { "s": SECONDS_PER_MINUTE, "ms": MS_PER_MINUTE, "us": US_PER_MINUTE, "ns": NS_PER_MINUTE, } factor = lit(unit_to_minutes_factor[self.unit], type=pa.int64()) return self.with_native(pc.divide(self.native, factor).cast(pa.int64())) def total_seconds(self) -> ArrowSeries: unit_to_seconds_factor = { "s": 1, "ms": MS_PER_SECOND, "us": US_PER_SECOND, "ns": NS_PER_SECOND, } factor = lit(unit_to_seconds_factor[self.unit], type=pa.int64()) return self.with_native(pc.divide(self.native, factor).cast(pa.int64())) def total_milliseconds(self) -> ArrowSeries: unit_to_milli_factor = { "s": 1e3, # seconds "ms": 1, # milli "us": 1e3, # micro "ns": 1e6, # nano } factor = lit(unit_to_milli_factor[self.unit], type=pa.int64()) if self.unit == "s": return self.with_native(pc.multiply(self.native, factor).cast(pa.int64())) return self.with_native(pc.divide(self.native, factor).cast(pa.int64())) def total_microseconds(self) -> ArrowSeries: unit_to_micro_factor = { "s": 1e6, # seconds "ms": 1e3, # milli "us": 1, # micro "ns": 1e3, # nano } factor = lit(unit_to_micro_factor[self.unit], type=pa.int64()) if self.unit in {"s", "ms"}: return self.with_native(pc.multiply(self.native, factor).cast(pa.int64())) return self.with_native(pc.divide(self.native, factor).cast(pa.int64())) def total_nanoseconds(self) -> ArrowSeries: unit_to_nano_factor = { "s": NS_PER_SECOND, "ms": NS_PER_MILLISECOND, "us": NS_PER_MICROSECOND, "ns": 1, } factor = lit(unit_to_nano_factor[self.unit], type=pa.int64()) return self.with_native(pc.multiply(self.native, factor).cast(pa.int64())) def truncate(self, every: str) -> ArrowSeries: interval = Interval.parse(every) return self.with_native( pc.floor_temporal(self.native, interval.multiple, UNITS_DICT[interval.unit]) ) def offset_by(self, by: str) -> ArrowSeries: interval = Interval.parse_no_constraints(by) native = self.native if interval.unit in {"y", "q", "mo"}: msg = f"Offsetting by {interval.unit} is not yet supported for pyarrow." raise NotImplementedError(msg) dtype = self.compliant.dtype datetime_dtype = self.version.dtypes.Datetime if interval.unit == "d" and isinstance(dtype, datetime_dtype) and dtype.time_zone: offset: pa.DurationScalar[Any] = lit(interval.to_timedelta()) native_naive = pc.local_timestamp(native) result = pc.assume_timezone(pc.add(native_naive, offset), dtype.time_zone) return self.with_native(result) if interval.unit == "ns": # pragma: no cover offset = lit(interval.multiple, pa.duration("ns")) # type: ignore[assignment] else: offset = lit(interval.to_timedelta()) return self.with_native(pc.add(native, offset)) narwhals-2.11.0/narwhals/_arrow/series_list.py0000644000000000000000000000134013615410400016337 0ustar00from __future__ import annotations from typing import TYPE_CHECKING import pyarrow as pa import pyarrow.compute as pc from narwhals._arrow.utils import ArrowSeriesNamespace from narwhals._compliant.any_namespace import ListNamespace from narwhals._utils import not_implemented if TYPE_CHECKING: from narwhals._arrow.series import ArrowSeries class ArrowSeriesListNamespace(ArrowSeriesNamespace, ListNamespace["ArrowSeries"]): def len(self) -> ArrowSeries: return self.with_native(pc.list_value_length(self.native).cast(pa.uint32())) def get(self, index: int) -> ArrowSeries: return self.with_native(pc.list_element(self.native, index)) unique = not_implemented() contains = not_implemented() narwhals-2.11.0/narwhals/_arrow/series_str.py0000644000000000000000000001071413615410400016201 0ustar00from __future__ import annotations import string from typing import TYPE_CHECKING import pyarrow as pa import pyarrow.compute as pc from narwhals._arrow.utils import ( ArrowSeriesNamespace, extract_native, lit, parse_datetime_format, ) from narwhals._compliant.any_namespace import StringNamespace if TYPE_CHECKING: from narwhals._arrow.series import ArrowSeries from narwhals._arrow.typing import Incomplete class ArrowSeriesStringNamespace(ArrowSeriesNamespace, StringNamespace["ArrowSeries"]): def len_chars(self) -> ArrowSeries: return self.with_native(pc.utf8_length(self.native)) def replace( self, value: ArrowSeries, pattern: str, *, literal: bool, n: int ) -> ArrowSeries: fn = pc.replace_substring if literal else pc.replace_substring_regex _, value_native = extract_native(self.compliant, value) if not isinstance(value_native, pa.StringScalar): msg = "PyArrow backed `.str.replace` only supports str replacement values" raise TypeError(msg) arr = fn( self.native, pattern, replacement=value_native.as_py(), max_replacements=n ) return self.with_native(arr) def replace_all( self, value: ArrowSeries, pattern: str, *, literal: bool ) -> ArrowSeries: return self.replace(value, pattern, literal=literal, n=-1) def strip_chars(self, characters: str | None) -> ArrowSeries: return self.with_native( pc.utf8_trim(self.native, characters or string.whitespace) ) def starts_with(self, prefix: str) -> ArrowSeries: return self.with_native(pc.equal(self.slice(0, len(prefix)).native, lit(prefix))) def ends_with(self, suffix: str) -> ArrowSeries: return self.with_native( pc.equal(self.slice(-len(suffix), None).native, lit(suffix)) ) def contains(self, pattern: str, *, literal: bool) -> ArrowSeries: check_func = pc.match_substring if literal else pc.match_substring_regex return self.with_native(check_func(self.native, pattern)) def slice(self, offset: int, length: int | None) -> ArrowSeries: stop = offset + length if length is not None else None return self.with_native( pc.utf8_slice_codeunits(self.native, start=offset, stop=stop) ) def split(self, by: str) -> ArrowSeries: split_series = pc.split_pattern(self.native, by) # type: ignore[call-overload] return self.with_native(split_series) def to_datetime(self, format: str | None) -> ArrowSeries: format = parse_datetime_format(self.native) if format is None else format timestamp_array = pc.strptime(self.native, format=format, unit="us") return self.with_native(timestamp_array) def to_date(self, format: str | None) -> ArrowSeries: return self.to_datetime(format=format).dt.date() def to_uppercase(self) -> ArrowSeries: return self.with_native(pc.utf8_upper(self.native)) def to_lowercase(self) -> ArrowSeries: return self.with_native(pc.utf8_lower(self.native)) def to_titlecase(self) -> ArrowSeries: return self.with_native(pc.utf8_title(self.native)) def zfill(self, width: int) -> ArrowSeries: binary_join: Incomplete = pc.binary_join_element_wise native = self.native hyphen, plus = lit("-"), lit("+") first_char, remaining_chars = ( self.slice(0, 1).native, self.slice(1, None).native, ) # Conditions less_than_width = pc.less(pc.utf8_length(native), lit(width)) starts_with_hyphen = pc.equal(first_char, hyphen) starts_with_plus = pc.equal(first_char, plus) conditions = pc.make_struct( pc.and_(starts_with_hyphen, less_than_width), pc.and_(starts_with_plus, less_than_width), less_than_width, ) # Cases padded_remaining_chars = pc.utf8_lpad(remaining_chars, width - 1, padding="0") result = pc.case_when( conditions, binary_join( pa.repeat(hyphen, len(native)), padded_remaining_chars, "" ), # starts with hyphen and less than width binary_join( pa.repeat(plus, len(native)), padded_remaining_chars, "" ), # starts with plus and less than width pc.utf8_lpad(native, width=width, padding="0"), # less than width native, ) return self.with_native(result) narwhals-2.11.0/narwhals/_arrow/series_struct.py0000644000000000000000000000077013615410400016716 0ustar00from __future__ import annotations from typing import TYPE_CHECKING import pyarrow.compute as pc from narwhals._arrow.utils import ArrowSeriesNamespace from narwhals._compliant.any_namespace import StructNamespace if TYPE_CHECKING: from narwhals._arrow.series import ArrowSeries class ArrowSeriesStructNamespace(ArrowSeriesNamespace, StructNamespace["ArrowSeries"]): def field(self, name: str) -> ArrowSeries: return self.with_native(pc.struct_field(self.native, name)).alias(name) narwhals-2.11.0/narwhals/_arrow/typing.py0000644000000000000000000000446713615410400015341 0ustar00from __future__ import annotations # pragma: no cover from typing import ( TYPE_CHECKING, # pragma: no cover Any, # pragma: no cover TypeVar, # pragma: no cover ) if TYPE_CHECKING: import sys from typing import Generic, Literal if sys.version_info >= (3, 10): from typing import TypeAlias else: from typing_extensions import TypeAlias import pyarrow as pa from pyarrow.__lib_pxi.table import ( AggregateOptions, # noqa: F401 Aggregation, # noqa: F401 ) from pyarrow._stubs_typing import ( # pyright: ignore[reportMissingModuleSource] # pyright: ignore[reportMissingModuleSource] # pyright: ignore[reportMissingModuleSource] Indices, # noqa: F401 Mask, # noqa: F401 Order, # noqa: F401 ) from narwhals._arrow.expr import ArrowExpr from narwhals._arrow.series import ArrowSeries IntoArrowExpr: TypeAlias = "ArrowExpr | ArrowSeries" TieBreaker: TypeAlias = Literal["min", "max", "first", "dense"] NullPlacement: TypeAlias = Literal["at_start", "at_end"] NativeIntervalUnit: TypeAlias = Literal[ "year", "quarter", "month", "week", "day", "hour", "minute", "second", "millisecond", "microsecond", "nanosecond", ] PromoteOptions: TypeAlias = Literal["none", "default", "permissive"] ChunkedArrayAny: TypeAlias = pa.ChunkedArray[Any] ArrayAny: TypeAlias = pa.Array[Any] ArrayOrChunkedArray: TypeAlias = "ArrayAny | ChunkedArrayAny" ScalarAny: TypeAlias = pa.Scalar[Any] ArrayOrScalar: TypeAlias = "ArrayOrChunkedArray | ScalarAny" ArrayOrScalarT1 = TypeVar("ArrayOrScalarT1", ArrayAny, ChunkedArrayAny, ScalarAny) ArrayOrScalarT2 = TypeVar("ArrayOrScalarT2", ArrayAny, ChunkedArrayAny, ScalarAny) _AsPyType = TypeVar("_AsPyType") class _BasicDataType(pa.DataType, Generic[_AsPyType]): ... Incomplete: TypeAlias = Any # pragma: no cover """ Marker for working code that fails on the stubs. Common issues: - Annotated for `Array`, but not `ChunkedArray` - Relies on typing information that the stubs don't provide statically - Missing attributes - Incorrect return types - Inconsistent use of generic/concrete types - `_clone_signature` used on signatures that are not identical """ narwhals-2.11.0/narwhals/_arrow/utils.py0000644000000000000000000004357213615410400015167 0ustar00from __future__ import annotations from functools import lru_cache from typing import TYPE_CHECKING, Any, cast import pyarrow as pa import pyarrow.compute as pc from narwhals._compliant import EagerSeriesNamespace from narwhals._utils import Implementation, Version, isinstance_or_issubclass if TYPE_CHECKING: from collections.abc import Iterable, Iterator, Mapping from typing_extensions import TypeAlias, TypeIs from narwhals._arrow.series import ArrowSeries from narwhals._arrow.typing import ( ArrayAny, ArrayOrScalar, ArrayOrScalarT1, ArrayOrScalarT2, ChunkedArrayAny, Incomplete, NativeIntervalUnit, PromoteOptions, ScalarAny, ) from narwhals._duration import IntervalUnit from narwhals.dtypes import DType from narwhals.typing import IntoDType, PythonLiteral # NOTE: stubs don't allow for `ChunkedArray[StructArray]` # Intended to represent the `.chunks` property storing `list[pa.StructArray]` ChunkedArrayStructArray: TypeAlias = ChunkedArrayAny def is_timestamp(t: Any) -> TypeIs[pa.TimestampType[Any, Any]]: ... def is_duration(t: Any) -> TypeIs[pa.DurationType[Any]]: ... def is_list(t: Any) -> TypeIs[pa.ListType[Any]]: ... def is_large_list(t: Any) -> TypeIs[pa.LargeListType[Any]]: ... def is_fixed_size_list(t: Any) -> TypeIs[pa.FixedSizeListType[Any, Any]]: ... def is_dictionary(t: Any) -> TypeIs[pa.DictionaryType[Any, Any, Any]]: ... def extract_regex( strings: ChunkedArrayAny, /, pattern: str, *, options: Any = None, memory_pool: Any = None, ) -> ChunkedArrayStructArray: ... else: from pyarrow.compute import extract_regex from pyarrow.types import ( is_dictionary, # noqa: F401 is_duration, is_fixed_size_list, is_large_list, is_list, is_timestamp, ) BACKEND_VERSION = Implementation.PYARROW._backend_version() """Static backend version for `pyarrow`.""" UNITS_DICT: Mapping[IntervalUnit, NativeIntervalUnit] = { "y": "year", "q": "quarter", "mo": "month", "d": "day", "h": "hour", "m": "minute", "s": "second", "ms": "millisecond", "us": "microsecond", "ns": "nanosecond", } lit = pa.scalar """Alias for `pyarrow.scalar`.""" def extract_py_scalar(value: Any, /) -> Any: from narwhals._arrow.series import maybe_extract_py_scalar return maybe_extract_py_scalar(value, return_py_scalar=True) def is_array_or_scalar(obj: Any) -> TypeIs[ArrayOrScalar]: """Return True for any base `pyarrow` container.""" return isinstance(obj, (pa.ChunkedArray, pa.Array, pa.Scalar)) def chunked_array( arr: ArrayOrScalar | list[Iterable[Any]], dtype: pa.DataType | None = None, / ) -> ChunkedArrayAny: if isinstance(arr, pa.ChunkedArray): return arr if isinstance(arr, list): return pa.chunked_array(arr, dtype) return pa.chunked_array([arr], dtype) def nulls_like(n: int, series: ArrowSeries) -> ArrayAny: """Create a strongly-typed Array instance with all elements null. Uses the type of `series`, without upseting `mypy`. """ return pa.nulls(n, series.native.type) def repeat( value: PythonLiteral | ScalarAny, n: int, /, dtype: pa.DataType | None = None ) -> ArrayAny: """Create an Array instance whose slots are the given scalar. *Optionally*, casting to `dtype` **before** repeating `n` times. """ lit_: Incomplete = lit return pa.repeat(lit_(value, type=dtype), n) def zeros(n: int, /) -> pa.Int64Array: return pa.repeat(0, n) def native_to_narwhals_dtype(dtype: pa.DataType, version: Version) -> DType: if isinstance(dtype, pa.ExtensionType): return version.dtypes.Unknown() return native_non_extension_to_narwhals_dtype(dtype, version) @lru_cache(maxsize=16) def native_non_extension_to_narwhals_dtype(dtype: pa.DataType, version: Version) -> DType: # noqa: C901, PLR0912 dtypes = version.dtypes if pa.types.is_int64(dtype): return dtypes.Int64() if pa.types.is_int32(dtype): return dtypes.Int32() if pa.types.is_int16(dtype): return dtypes.Int16() if pa.types.is_int8(dtype): return dtypes.Int8() if pa.types.is_uint64(dtype): return dtypes.UInt64() if pa.types.is_uint32(dtype): return dtypes.UInt32() if pa.types.is_uint16(dtype): return dtypes.UInt16() if pa.types.is_uint8(dtype): return dtypes.UInt8() if pa.types.is_boolean(dtype): return dtypes.Boolean() if pa.types.is_float64(dtype): return dtypes.Float64() if pa.types.is_float32(dtype): return dtypes.Float32() # bug in coverage? it shows `31->exit` (where `31` is currently the line number of # the next line), even though both when the if condition is true and false are covered if ( # pragma: no cover pa.types.is_string(dtype) or pa.types.is_large_string(dtype) or getattr(pa.types, "is_string_view", lambda _: False)(dtype) ): return dtypes.String() if pa.types.is_date32(dtype): return dtypes.Date() if is_timestamp(dtype): return dtypes.Datetime(time_unit=dtype.unit, time_zone=dtype.tz) if is_duration(dtype): return dtypes.Duration(time_unit=dtype.unit) if pa.types.is_dictionary(dtype): return dtypes.Categorical() if pa.types.is_struct(dtype): return dtypes.Struct( [ dtypes.Field( dtype.field(i).name, native_to_narwhals_dtype(dtype.field(i).type, version), ) for i in range(dtype.num_fields) ] ) if is_list(dtype) or is_large_list(dtype): return dtypes.List(native_to_narwhals_dtype(dtype.value_type, version)) if is_fixed_size_list(dtype): return dtypes.Array( native_to_narwhals_dtype(dtype.value_type, version), dtype.list_size ) if pa.types.is_decimal(dtype): return dtypes.Decimal() if pa.types.is_time32(dtype) or pa.types.is_time64(dtype): return dtypes.Time() if pa.types.is_binary(dtype): return dtypes.Binary() return dtypes.Unknown() # pragma: no cover dtypes = Version.MAIN.dtypes NW_TO_PA_DTYPES: Mapping[type[DType], pa.DataType] = { dtypes.Float64: pa.float64(), dtypes.Float32: pa.float32(), dtypes.Binary: pa.binary(), dtypes.String: pa.string(), dtypes.Boolean: pa.bool_(), dtypes.Categorical: pa.dictionary(pa.uint32(), pa.string()), dtypes.Date: pa.date32(), dtypes.Time: pa.time64("ns"), dtypes.Int8: pa.int8(), dtypes.Int16: pa.int16(), dtypes.Int32: pa.int32(), dtypes.Int64: pa.int64(), dtypes.UInt8: pa.uint8(), dtypes.UInt16: pa.uint16(), dtypes.UInt32: pa.uint32(), dtypes.UInt64: pa.uint64(), } UNSUPPORTED_DTYPES = (dtypes.Decimal, dtypes.Object) def narwhals_to_native_dtype(dtype: IntoDType, version: Version) -> pa.DataType: dtypes = version.dtypes base_type = dtype.base_type() if pa_type := NW_TO_PA_DTYPES.get(base_type): return pa_type if isinstance_or_issubclass(dtype, dtypes.Datetime): unit = dtype.time_unit return pa.timestamp(unit, tz) if (tz := dtype.time_zone) else pa.timestamp(unit) if isinstance_or_issubclass(dtype, dtypes.Duration): return pa.duration(dtype.time_unit) if isinstance_or_issubclass(dtype, dtypes.List): return pa.list_(value_type=narwhals_to_native_dtype(dtype.inner, version=version)) if isinstance_or_issubclass(dtype, dtypes.Struct): return pa.struct( [ (field.name, narwhals_to_native_dtype(field.dtype, version=version)) for field in dtype.fields ] ) if isinstance_or_issubclass(dtype, dtypes.Array): # pragma: no cover inner = narwhals_to_native_dtype(dtype.inner, version=version) list_size = dtype.size return pa.list_(inner, list_size=list_size) if issubclass(base_type, UNSUPPORTED_DTYPES): msg = f"Converting to {base_type.__name__} dtype is not supported for PyArrow." raise NotImplementedError(msg) msg = f"Unknown dtype: {dtype}" # pragma: no cover raise AssertionError(msg) def extract_native( lhs: ArrowSeries, rhs: ArrowSeries | PythonLiteral | ScalarAny ) -> tuple[ChunkedArrayAny, ChunkedArrayAny | ScalarAny]: """Extract native objects in binary operation. If the comparison isn't supported, return `NotImplemented` so that the "right-hand-side" operation (e.g. `__radd__`) can be tried. If one of the two sides has a `_broadcast` flag, then extract the scalar underneath it so that PyArrow can do its own broadcasting. """ from narwhals._arrow.series import ArrowSeries if rhs is None: # pragma: no cover return lhs.native, lit(None, type=lhs._type) if isinstance(rhs, ArrowSeries): if lhs._broadcast and not rhs._broadcast: return lhs.native[0], rhs.native if rhs._broadcast: return lhs.native, rhs.native[0] return lhs.native, rhs.native if isinstance(rhs, list): msg = "Expected Series or scalar, got list." raise TypeError(msg) return lhs.native, rhs if isinstance(rhs, pa.Scalar) else lit(rhs) def floordiv_compat(left: ArrayOrScalar, right: ArrayOrScalar, /) -> Any: # The following lines are adapted from pandas' pyarrow implementation. # Ref: https://github.com/pandas-dev/pandas/blob/262fcfbffcee5c3116e86a951d8b693f90411e68/pandas/core/arrays/arrow/array.py#L124-L154 # We modify it to add `safe_mask` so as to align with Polars' behaviour when dividing by zero. safe_mask = pc.not_equal(right, lit(0, type=right.type)) safe_right = pc.if_else(safe_mask, right, lit(1, type=right.type)) # Dummy value non_safe_result = lit(None, type=left.type) if pa.types.is_integer(left.type) and pa.types.is_integer(right.type): divided = pc.if_else( safe_mask, pc.divide_checked(left, safe_right), non_safe_result ) # TODO @dangotbanned: Use a `TypeVar` in guards # Narrowing to a `Union` isn't interacting well with the rest of the stubs # https://github.com/zen-xu/pyarrow-stubs/pull/215 if pa.types.is_signed_integer(divided.type): div_type = cast("pa._lib.Int64Type", divided.type) has_remainder = pc.not_equal(pc.multiply(divided, right), left) has_one_negative_operand = pc.less( pc.bit_wise_xor(left, right), lit(0, div_type) ) result = pc.if_else( pc.and_(has_remainder, has_one_negative_operand), pc.subtract(divided, lit(1, div_type)), divided, ) else: result = divided # pragma: no cover result = result.cast(left.type) else: divided = pc.if_else(safe_mask, pc.divide(left, safe_right), non_safe_result) result = pc.floor(divided) return result def cast_for_truediv( arrow_array: ArrayOrScalarT1, pa_object: ArrayOrScalarT2 ) -> tuple[ArrayOrScalarT1, ArrayOrScalarT2]: # Lifted from: # https://github.com/pandas-dev/pandas/blob/262fcfbffcee5c3116e86a951d8b693f90411e68/pandas/core/arrays/arrow/array.py#L108-L122 # Ensure int / int -> float mirroring Python/Numpy behavior # as pc.divide_checked(int, int) -> int if pa.types.is_integer(arrow_array.type) and pa.types.is_integer(pa_object.type): # GH: 56645. # noqa: ERA001 # https://github.com/apache/arrow/issues/35563 return arrow_array.cast(pa.float64(), safe=False), pa_object.cast( pa.float64(), safe=False ) return arrow_array, pa_object # Regex for date, time, separator and timezone components DATE_RE = r"(?P\d{1,4}[-/.]\d{1,2}[-/.]\d{1,4}|\d{8})" SEP_RE = r"(?P\s|T)" TIME_RE = r"(?P