352 lines
12 KiB
Python
352 lines
12 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import abc
|
||
|
|
from functools import lru_cache
|
||
|
|
from typing import TYPE_CHECKING, Any, ClassVar, Final, Protocol, TypeVar, overload
|
||
|
|
|
||
|
|
import polars as pl
|
||
|
|
|
||
|
|
from narwhals._duration import Interval
|
||
|
|
from narwhals._utils import (
|
||
|
|
Implementation,
|
||
|
|
Version,
|
||
|
|
_DeferredIterable,
|
||
|
|
_StoresCompliant,
|
||
|
|
_StoresNative,
|
||
|
|
deep_getattr,
|
||
|
|
isinstance_or_issubclass,
|
||
|
|
)
|
||
|
|
from narwhals.exceptions import (
|
||
|
|
ColumnNotFoundError,
|
||
|
|
ComputeError,
|
||
|
|
DuplicateError,
|
||
|
|
InvalidOperationError,
|
||
|
|
NarwhalsError,
|
||
|
|
ShapeError,
|
||
|
|
)
|
||
|
|
|
||
|
|
if TYPE_CHECKING:
|
||
|
|
from collections.abc import Callable, Iterable, Iterator, Mapping
|
||
|
|
|
||
|
|
from typing_extensions import TypeIs
|
||
|
|
|
||
|
|
from narwhals._polars.dataframe import Method
|
||
|
|
from narwhals._polars.expr import PolarsExpr
|
||
|
|
from narwhals._polars.series import PolarsSeries
|
||
|
|
from narwhals._polars.typing import NativeAccessor
|
||
|
|
from narwhals.dtypes import DType
|
||
|
|
from narwhals.typing import IntoDType
|
||
|
|
|
||
|
|
T = TypeVar("T")
|
||
|
|
NativeT = TypeVar(
|
||
|
|
"NativeT", bound="pl.DataFrame | pl.LazyFrame | pl.Series | pl.Expr"
|
||
|
|
)
|
||
|
|
|
||
|
|
NativeT_co = TypeVar("NativeT_co", "pl.Series", "pl.Expr", covariant=True)
|
||
|
|
CompliantT_co = TypeVar("CompliantT_co", "PolarsSeries", "PolarsExpr", covariant=True)
|
||
|
|
CompliantT = TypeVar("CompliantT", "PolarsSeries", "PolarsExpr")
|
||
|
|
|
||
|
|
BACKEND_VERSION = Implementation.POLARS._backend_version()
|
||
|
|
"""Static backend version for `polars`."""
|
||
|
|
|
||
|
|
SERIES_RESPECTS_DTYPE: Final[bool] = BACKEND_VERSION >= (0, 20, 26)
|
||
|
|
"""`pl.Series(dtype=...)` fixed in https://github.com/pola-rs/polars/pull/15962
|
||
|
|
|
||
|
|
Includes `SERIES_ACCEPTS_PD_INDEX`.
|
||
|
|
"""
|
||
|
|
|
||
|
|
SERIES_ACCEPTS_PD_INDEX: Final[bool] = BACKEND_VERSION >= (0, 20, 7)
|
||
|
|
"""`pl.Series(values: pd.Index)` fixed in https://github.com/pola-rs/polars/pull/14087"""
|
||
|
|
|
||
|
|
|
||
|
|
@overload
|
||
|
|
def extract_native(obj: _StoresNative[NativeT]) -> NativeT: ...
|
||
|
|
@overload
|
||
|
|
def extract_native(obj: T) -> T: ...
|
||
|
|
def extract_native(obj: _StoresNative[NativeT] | T) -> NativeT | T:
|
||
|
|
return obj.native if _is_compliant_polars(obj) else obj
|
||
|
|
|
||
|
|
|
||
|
|
def _is_compliant_polars(
|
||
|
|
obj: _StoresNative[NativeT] | Any,
|
||
|
|
) -> TypeIs[_StoresNative[NativeT]]:
|
||
|
|
from narwhals._polars.dataframe import PolarsDataFrame, PolarsLazyFrame
|
||
|
|
from narwhals._polars.expr import PolarsExpr
|
||
|
|
from narwhals._polars.series import PolarsSeries
|
||
|
|
|
||
|
|
return isinstance(obj, (PolarsDataFrame, PolarsLazyFrame, PolarsSeries, PolarsExpr))
|
||
|
|
|
||
|
|
|
||
|
|
def extract_args_kwargs(
|
||
|
|
args: Iterable[Any], kwds: Mapping[str, Any], /
|
||
|
|
) -> tuple[Iterator[Any], dict[str, Any]]:
|
||
|
|
it_args = (extract_native(arg) for arg in args)
|
||
|
|
return it_args, {k: extract_native(v) for k, v in kwds.items()}
|
||
|
|
|
||
|
|
|
||
|
|
@lru_cache(maxsize=16)
|
||
|
|
def native_to_narwhals_dtype( # noqa: C901, PLR0912
|
||
|
|
dtype: pl.DataType, version: Version
|
||
|
|
) -> DType:
|
||
|
|
dtypes = version.dtypes
|
||
|
|
if dtype == pl.Float64:
|
||
|
|
return dtypes.Float64()
|
||
|
|
if dtype == pl.Float32:
|
||
|
|
return dtypes.Float32()
|
||
|
|
if hasattr(pl, "Int128") and dtype == pl.Int128: # pragma: no cover
|
||
|
|
# Not available for Polars pre 1.8.0
|
||
|
|
return dtypes.Int128()
|
||
|
|
if dtype == pl.Int64:
|
||
|
|
return dtypes.Int64()
|
||
|
|
if dtype == pl.Int32:
|
||
|
|
return dtypes.Int32()
|
||
|
|
if dtype == pl.Int16:
|
||
|
|
return dtypes.Int16()
|
||
|
|
if dtype == pl.Int8:
|
||
|
|
return dtypes.Int8()
|
||
|
|
if hasattr(pl, "UInt128") and dtype == pl.UInt128: # pragma: no cover
|
||
|
|
# Not available for Polars pre 1.8.0
|
||
|
|
return dtypes.UInt128()
|
||
|
|
if dtype == pl.UInt64:
|
||
|
|
return dtypes.UInt64()
|
||
|
|
if dtype == pl.UInt32:
|
||
|
|
return dtypes.UInt32()
|
||
|
|
if dtype == pl.UInt16:
|
||
|
|
return dtypes.UInt16()
|
||
|
|
if dtype == pl.UInt8:
|
||
|
|
return dtypes.UInt8()
|
||
|
|
if dtype == pl.String:
|
||
|
|
return dtypes.String()
|
||
|
|
if dtype == pl.Boolean:
|
||
|
|
return dtypes.Boolean()
|
||
|
|
if dtype == pl.Object:
|
||
|
|
return dtypes.Object()
|
||
|
|
if dtype == pl.Categorical:
|
||
|
|
return dtypes.Categorical()
|
||
|
|
if isinstance_or_issubclass(dtype, pl.Enum):
|
||
|
|
if version is Version.V1:
|
||
|
|
return dtypes.Enum() # type: ignore[call-arg]
|
||
|
|
categories = _DeferredIterable(dtype.categories.to_list)
|
||
|
|
return dtypes.Enum(categories)
|
||
|
|
if dtype == pl.Date:
|
||
|
|
return dtypes.Date()
|
||
|
|
if isinstance_or_issubclass(dtype, pl.Datetime):
|
||
|
|
return (
|
||
|
|
dtypes.Datetime()
|
||
|
|
if dtype is pl.Datetime
|
||
|
|
else dtypes.Datetime(dtype.time_unit, dtype.time_zone)
|
||
|
|
)
|
||
|
|
if isinstance_or_issubclass(dtype, pl.Duration):
|
||
|
|
return (
|
||
|
|
dtypes.Duration()
|
||
|
|
if dtype is pl.Duration
|
||
|
|
else dtypes.Duration(dtype.time_unit)
|
||
|
|
)
|
||
|
|
if isinstance_or_issubclass(dtype, pl.Struct):
|
||
|
|
fields = [
|
||
|
|
dtypes.Field(name, native_to_narwhals_dtype(tp, version))
|
||
|
|
for name, tp in dtype
|
||
|
|
]
|
||
|
|
return dtypes.Struct(fields)
|
||
|
|
if isinstance_or_issubclass(dtype, pl.List):
|
||
|
|
return dtypes.List(native_to_narwhals_dtype(dtype.inner, version))
|
||
|
|
if isinstance_or_issubclass(dtype, pl.Array):
|
||
|
|
outer_shape = dtype.width if BACKEND_VERSION < (0, 20, 30) else dtype.size
|
||
|
|
return dtypes.Array(native_to_narwhals_dtype(dtype.inner, version), outer_shape)
|
||
|
|
if dtype == pl.Decimal:
|
||
|
|
return dtypes.Decimal()
|
||
|
|
if dtype == pl.Time:
|
||
|
|
return dtypes.Time()
|
||
|
|
if dtype == pl.Binary:
|
||
|
|
return dtypes.Binary()
|
||
|
|
return dtypes.Unknown()
|
||
|
|
|
||
|
|
|
||
|
|
dtypes = Version.MAIN.dtypes
|
||
|
|
NW_TO_PL_DTYPES: Mapping[type[DType], pl.DataType] = {
|
||
|
|
dtypes.Float64: pl.Float64(),
|
||
|
|
dtypes.Float32: pl.Float32(),
|
||
|
|
dtypes.Binary: pl.Binary(),
|
||
|
|
dtypes.String: pl.String(),
|
||
|
|
dtypes.Boolean: pl.Boolean(),
|
||
|
|
dtypes.Categorical: pl.Categorical(),
|
||
|
|
dtypes.Date: pl.Date(),
|
||
|
|
dtypes.Time: pl.Time(),
|
||
|
|
dtypes.Int8: pl.Int8(),
|
||
|
|
dtypes.Int16: pl.Int16(),
|
||
|
|
dtypes.Int32: pl.Int32(),
|
||
|
|
dtypes.Int64: pl.Int64(),
|
||
|
|
dtypes.UInt8: pl.UInt8(),
|
||
|
|
dtypes.UInt16: pl.UInt16(),
|
||
|
|
dtypes.UInt32: pl.UInt32(),
|
||
|
|
dtypes.UInt64: pl.UInt64(),
|
||
|
|
dtypes.Object: pl.Object(),
|
||
|
|
dtypes.Unknown: pl.Unknown(),
|
||
|
|
}
|
||
|
|
UNSUPPORTED_DTYPES = (dtypes.Decimal,)
|
||
|
|
|
||
|
|
|
||
|
|
def narwhals_to_native_dtype( # noqa: C901
|
||
|
|
dtype: IntoDType, version: Version
|
||
|
|
) -> pl.DataType:
|
||
|
|
dtypes = version.dtypes
|
||
|
|
base_type = dtype.base_type()
|
||
|
|
if pl_type := NW_TO_PL_DTYPES.get(base_type):
|
||
|
|
return pl_type
|
||
|
|
if dtype == dtypes.Int128 and hasattr(pl, "Int128"):
|
||
|
|
# Not available for Polars pre 1.8.0
|
||
|
|
return pl.Int128()
|
||
|
|
if isinstance_or_issubclass(dtype, dtypes.Enum):
|
||
|
|
if version is Version.V1:
|
||
|
|
msg = "Converting to Enum is not supported in narwhals.stable.v1"
|
||
|
|
raise NotImplementedError(msg)
|
||
|
|
if isinstance(dtype, dtypes.Enum):
|
||
|
|
return pl.Enum(dtype.categories)
|
||
|
|
msg = "Can not cast / initialize Enum without categories present"
|
||
|
|
raise ValueError(msg)
|
||
|
|
if isinstance_or_issubclass(dtype, dtypes.Datetime):
|
||
|
|
return pl.Datetime(dtype.time_unit, dtype.time_zone) # type: ignore[arg-type]
|
||
|
|
if isinstance_or_issubclass(dtype, dtypes.Duration):
|
||
|
|
return pl.Duration(dtype.time_unit) # type: ignore[arg-type]
|
||
|
|
if isinstance_or_issubclass(dtype, dtypes.List):
|
||
|
|
return pl.List(narwhals_to_native_dtype(dtype.inner, version))
|
||
|
|
if isinstance_or_issubclass(dtype, dtypes.Struct):
|
||
|
|
fields = [
|
||
|
|
pl.Field(field.name, narwhals_to_native_dtype(field.dtype, version))
|
||
|
|
for field in dtype.fields
|
||
|
|
]
|
||
|
|
return pl.Struct(fields)
|
||
|
|
if isinstance_or_issubclass(dtype, dtypes.Array): # pragma: no cover
|
||
|
|
size = dtype.size
|
||
|
|
kwargs = {"width": size} if BACKEND_VERSION < (0, 20, 30) else {"shape": size}
|
||
|
|
return pl.Array(narwhals_to_native_dtype(dtype.inner, version), **kwargs)
|
||
|
|
if issubclass(base_type, UNSUPPORTED_DTYPES):
|
||
|
|
msg = f"Converting to {base_type.__name__} dtype is not supported for Polars."
|
||
|
|
raise NotImplementedError(msg)
|
||
|
|
return pl.Unknown() # pragma: no cover
|
||
|
|
|
||
|
|
|
||
|
|
def _is_polars_exception(exception: Exception) -> bool:
|
||
|
|
if BACKEND_VERSION >= (1,):
|
||
|
|
# Old versions of Polars didn't have PolarsError.
|
||
|
|
return isinstance(exception, pl.exceptions.PolarsError)
|
||
|
|
# Last attempt, for old Polars versions.
|
||
|
|
return "polars.exceptions" in str(type(exception)) # pragma: no cover
|
||
|
|
|
||
|
|
|
||
|
|
def _is_cudf_exception(exception: Exception) -> bool:
|
||
|
|
# These exceptions are raised when running polars on GPUs via cuDF
|
||
|
|
return str(exception).startswith("CUDF failure")
|
||
|
|
|
||
|
|
|
||
|
|
def catch_polars_exception(exception: Exception) -> NarwhalsError | Exception:
|
||
|
|
if isinstance(exception, pl.exceptions.ColumnNotFoundError):
|
||
|
|
return ColumnNotFoundError(str(exception))
|
||
|
|
if isinstance(exception, pl.exceptions.ShapeError):
|
||
|
|
return ShapeError(str(exception))
|
||
|
|
if isinstance(exception, pl.exceptions.InvalidOperationError):
|
||
|
|
return InvalidOperationError(str(exception))
|
||
|
|
if isinstance(exception, pl.exceptions.DuplicateError):
|
||
|
|
return DuplicateError(str(exception))
|
||
|
|
if isinstance(exception, pl.exceptions.ComputeError):
|
||
|
|
return ComputeError(str(exception))
|
||
|
|
if _is_polars_exception(exception) or _is_cudf_exception(exception):
|
||
|
|
return NarwhalsError(str(exception)) # pragma: no cover
|
||
|
|
# Just return exception as-is.
|
||
|
|
return exception
|
||
|
|
|
||
|
|
|
||
|
|
class PolarsAnyNamespace(
|
||
|
|
_StoresCompliant[CompliantT_co],
|
||
|
|
_StoresNative[NativeT_co],
|
||
|
|
Protocol[CompliantT_co, NativeT_co],
|
||
|
|
):
|
||
|
|
_accessor: ClassVar[NativeAccessor]
|
||
|
|
|
||
|
|
def __getattr__(self, attr: str) -> Callable[..., CompliantT_co]:
|
||
|
|
def func(*args: Any, **kwargs: Any) -> CompliantT_co:
|
||
|
|
pos, kwds = extract_args_kwargs(args, kwargs)
|
||
|
|
method = deep_getattr(self.native, self._accessor, attr)
|
||
|
|
return self.compliant._with_native(method(*pos, **kwds))
|
||
|
|
|
||
|
|
return func
|
||
|
|
|
||
|
|
|
||
|
|
class PolarsDateTimeNamespace(PolarsAnyNamespace[CompliantT, NativeT_co]):
|
||
|
|
_accessor: ClassVar[NativeAccessor] = "dt"
|
||
|
|
|
||
|
|
def truncate(self, every: str) -> CompliantT:
|
||
|
|
# Ensure consistent error message is raised.
|
||
|
|
Interval.parse(every)
|
||
|
|
return self.__getattr__("truncate")(every)
|
||
|
|
|
||
|
|
def offset_by(self, by: str) -> CompliantT:
|
||
|
|
# Ensure consistent error message is raised.
|
||
|
|
Interval.parse_no_constraints(by)
|
||
|
|
return self.__getattr__("offset_by")(by)
|
||
|
|
|
||
|
|
to_string: Method[CompliantT]
|
||
|
|
replace_time_zone: Method[CompliantT]
|
||
|
|
convert_time_zone: Method[CompliantT]
|
||
|
|
timestamp: Method[CompliantT]
|
||
|
|
date: Method[CompliantT]
|
||
|
|
year: Method[CompliantT]
|
||
|
|
month: Method[CompliantT]
|
||
|
|
day: Method[CompliantT]
|
||
|
|
hour: Method[CompliantT]
|
||
|
|
minute: Method[CompliantT]
|
||
|
|
second: Method[CompliantT]
|
||
|
|
millisecond: Method[CompliantT]
|
||
|
|
microsecond: Method[CompliantT]
|
||
|
|
nanosecond: Method[CompliantT]
|
||
|
|
ordinal_day: Method[CompliantT]
|
||
|
|
weekday: Method[CompliantT]
|
||
|
|
total_minutes: Method[CompliantT]
|
||
|
|
total_seconds: Method[CompliantT]
|
||
|
|
total_milliseconds: Method[CompliantT]
|
||
|
|
total_microseconds: Method[CompliantT]
|
||
|
|
total_nanoseconds: Method[CompliantT]
|
||
|
|
|
||
|
|
|
||
|
|
class PolarsStringNamespace(PolarsAnyNamespace[CompliantT, NativeT_co]):
|
||
|
|
_accessor: ClassVar[NativeAccessor] = "str"
|
||
|
|
|
||
|
|
# NOTE: Use `abstractmethod` if we have defs to implement, but also `Method` usage
|
||
|
|
@abc.abstractmethod
|
||
|
|
def zfill(self, width: int) -> CompliantT: ...
|
||
|
|
|
||
|
|
len_chars: Method[CompliantT]
|
||
|
|
replace: Method[CompliantT]
|
||
|
|
replace_all: Method[CompliantT]
|
||
|
|
strip_chars: Method[CompliantT]
|
||
|
|
starts_with: Method[CompliantT]
|
||
|
|
ends_with: Method[CompliantT]
|
||
|
|
contains: Method[CompliantT]
|
||
|
|
slice: Method[CompliantT]
|
||
|
|
split: Method[CompliantT]
|
||
|
|
to_date: Method[CompliantT]
|
||
|
|
to_datetime: Method[CompliantT]
|
||
|
|
to_lowercase: Method[CompliantT]
|
||
|
|
to_uppercase: Method[CompliantT]
|
||
|
|
|
||
|
|
|
||
|
|
class PolarsCatNamespace(PolarsAnyNamespace[CompliantT, NativeT_co]):
|
||
|
|
_accessor: ClassVar[NativeAccessor] = "cat"
|
||
|
|
get_categories: Method[CompliantT]
|
||
|
|
|
||
|
|
|
||
|
|
class PolarsListNamespace(PolarsAnyNamespace[CompliantT, NativeT_co]):
|
||
|
|
_accessor: ClassVar[NativeAccessor] = "list"
|
||
|
|
|
||
|
|
@abc.abstractmethod
|
||
|
|
def len(self) -> CompliantT: ...
|
||
|
|
|
||
|
|
get: Method[CompliantT]
|
||
|
|
|
||
|
|
unique: Method[CompliantT]
|
||
|
|
|
||
|
|
|
||
|
|
class PolarsStructNamespace(PolarsAnyNamespace[CompliantT, NativeT_co]):
|
||
|
|
_accessor: ClassVar[NativeAccessor] = "struct"
|
||
|
|
field: Method[CompliantT]
|