Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 15 additions & 75 deletions narwhals/dtypes/_supertyping.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,13 @@
"""Rules for safe type promotion.

Follows a subset of `polars`' [`get_supertype_with_options`].

See [Data type promotion rules] for an in-depth explanation.

[`get_supertype_with_options`]: https://github.com/pola-rs/polars/blob/529f7ec642912a2f15656897d06f1532c2f5d4c4/crates/polars-core/src/utils/supertype.rs#L142-L543
[Data type promotion rules]: https://narwhals-dev.github.io/narwhals/concepts/promotion-rules/
"""

from __future__ import annotations

from collections import deque
Expand Down Expand Up @@ -149,21 +159,7 @@ def dtype_eq(left: DType, right: DType, /) -> bool:

@cache
def _integer_supertyping() -> Mapping[FrozenDTypes, type[Int | Float64]]:
"""Generate the supertype conversion table for all integer data type pairs.

The rules:

# pick the highest bit-width for matching signs
(Int*, Int*) -> Int*
(UInt*, UInt*) -> UInt*
# pick a *strictly higher* signed bit-width
(UInt<lower>, Int<higher>) -> Int<higher>
# promote unsigned to the next highest bit-width,
# but do not exceed `Int64`
(UInt{8,16,32}, Int*) -> Int{16,32,64}
# all others
(UInt{64,128}, Int*) -> Float64
"""
"""Generate the supertype conversion table for all integer data type pairs."""
tps_int = SignedIntegerType.__subclasses__()
tps_uint = UnsignedIntegerType.__subclasses__()
get_bits: attrgetter[_Bits] = attrgetter("_bits")
Expand All @@ -190,14 +186,7 @@ def _integer_supertyping() -> Mapping[FrozenDTypes, type[Int | Float64]]:

@cache
def _primitive_numeric_supertyping() -> Mapping[FrozenDTypes, type[Float]]:
"""Generate the supertype conversion table for all (integer, float) data type pairs.

The rules:

(Integer{8,16}, Float32) -> Float32
(Integer{32,64,128}, Float32) -> Float64
(Integer, Float64) -> Float64
"""
"""Generate the supertype conversion table for all (integer, float) data type pairs."""
F32, F64 = Float32, Float64 # noqa: N806
small_int = (Int8, Int16, UInt8, UInt16)
small_int_f32 = ((frozen_dtypes(tp, F32), F32) for tp in small_int)
Expand Down Expand Up @@ -225,10 +214,7 @@ def has_nested(base_types: FrozenDTypes, /) -> bool:
def _struct_fields_union(
left: Collection[Field], right: Collection[Field], /
) -> Struct | None:
"""Adapted from [`union_struct_fields`].

[`union_struct_fields`]: https://github.com/pola-rs/polars/blob/c2412600210a21143835c9dfcb0a9182f462b619/crates/polars-core/src/utils/supertype.rs#L559-L586
"""
"""Adapted from [`union_struct_fields`](https://github.com/pola-rs/polars/blob/c2412600210a21143835c9dfcb0a9182f462b619/crates/polars-core/src/utils/supertype.rs#L559-L586)."""
longest, shortest = (left, right) if len(left) >= len(right) else (right, left)
longest_map = {f.name: f.dtype() for f in longest}
for f in shortest:
Expand All @@ -245,16 +231,7 @@ def _struct_fields_union(
def _struct_supertype(left: Struct, right: Struct, /) -> Struct | None:
"""Get the supertype of two struct data types.

Adapted from [`super_type_structs`]

Unlike all other rules, the order of *operands* is meaningful.
`left` defines the field order of the output `Struct` *unless* `right` has more fields.

We can derive a supertype with each `Struct`'s fields in arbitrary order,
and even with disjoint field names.
*But*, **all fields that intersect** (*by name*) must satisfy all other supertyping rules (*by dtype*).

[`super_type_structs`]: https://github.com/pola-rs/polars/blob/c2412600210a21143835c9dfcb0a9182f462b619/crates/polars-core/src/utils/supertype.rs#L588-L603
Adapted from [`super_type_structs`](https://github.com/pola-rs/polars/blob/c2412600210a21143835c9dfcb0a9182f462b619/crates/polars-core/src/utils/supertype.rs#L588-L603)
"""
left_fields, right_fields = left.fields, right.fields
if len(left_fields) != len(right_fields):
Expand Down Expand Up @@ -343,18 +320,6 @@ def _numeric_supertype(base_types: FrozenDTypes) -> DType | None:
`_{primitive_numeric,integer}_supertyping` define most valid numeric supertypes.

We generate these on first use, with all subsequent calls returning the same mapping.

The rules defined here are:

(Float32, Float64) -> Float64
(Decimal, Float*) -> Float64
(Decimal, Integer) -> Decimal
(Boolean, Numeric) -> Numeric

Important:
`Decimal` behavior is expected to change following [#3377]

[#3377]: https://github.com/narwhals-dev/narwhals/pull/3377
"""
if NUMERIC.issuperset(base_types):
if INTEGER.issuperset(base_types):
Expand All @@ -370,28 +335,7 @@ def _numeric_supertype(base_types: FrozenDTypes) -> DType | None:


def _mixed_supertype(st: _SupertypeCase[DType, DType]) -> DType | None:
"""Get the supertype of two data types that do not share the same class.

We support only one combination that requires preservation of instance attributes:

(Date, Datetime) -> Datetime

All others can match using *only* the class pairs themselves.

The following are supported in `polars`, but are not planned to be implemented here (see [#121]):

(Date, {UInt,Int,Float}{32,64}) -> {Int,Float}{32,64}
(Time, {Int,Float}{32,64}) -> {Int,Float}64
(Datetime, {UInt,Int,Float}{32,64}) -> {Int,Float}64
(Duration, {UInt,Int,Float}{32,64}) -> {Int,Float}64

We also reject all nested data types here, *whereas* [`polars` supports mixed `Struct`]:

(Struct, DType) -> Struct

[#121]: https://github.com/narwhals-dev/narwhals/issues/121
[`polars` supports mixed `Struct`]: https://github.com/pola-rs/polars/blob/d6d9d8a2c7d3e416488388a0114c6ff3eafcb66c/crates/polars-core/src/utils/supertype.rs#L499-L507
"""
"""Get the supertype of two data types that do not share the same class."""
if Date in st.base_types and _has_intersection(st.base_types, DATETIME):
return st.left if isinstance(st.left, Datetime) else st.right
if NUMERIC.isdisjoint(st.base_types):
Expand Down Expand Up @@ -429,16 +373,12 @@ def get_supertype(
) -> DTypeT1 | DTypeT2 | DType | None:
"""Given two data types, determine the data type that both types can reasonably safely be cast to.

Aims to follow the rules defined by [`polars_core::utils::supertype::get_supertype_with_options`].

Arguments:
left: First data type.
right: Second data type.

Returns:
The common supertype that both types can be safely cast to, or None if no such type exists.

[`polars_core::utils::supertype::get_supertype_with_options`]: https://github.com/pola-rs/polars/blob/529f7ec642912a2f15656897d06f1532c2f5d4c4/crates/polars-core/src/utils/supertype.rs#L142-L543
"""
st_case = _SupertypeCase(left, right)
if Unknown in st_case.base_types:
Expand Down
8 changes: 7 additions & 1 deletion utils/promotion-rules.md.jinja
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,11 @@ st(nw.Float64(), nw.Float32())
st(nw.Decimal(), nw.Float64())
```

<!--
TODO(dangotbanned): Update once Decimal is parametric
https://github.com/narwhals-dev/narwhals/pull/3377
-->

## Temporal Types

### Duration
Expand Down Expand Up @@ -349,7 +354,7 @@ The following combinations have **no valid supertype** and will result in `None`
* `Datetime` values with different time zones
* `Enum` values with different categories
* `Array` values with different shapes
* Temporal types combined with numeric types (unlike Polars, Narwhals does not support these conversions)
* Temporal types combined with numeric types (unlike Polars, see [narwhals-dev/narwhals#121])
* Any other unlisted combination of different base types

## Everything in a single table
Expand All @@ -358,6 +363,7 @@ The following combinations have **no valid supertype** and will result in `None`

[get_supertype_with_options]: https://github.com/pola-rs/polars/blob/529f7ec642912a2f15656897d06f1532c2f5d4c4/crates/polars-core/src/utils/supertype.rs#L142-L543
[base-type]: ../api-reference/dtypes.md#narwhals.dtypes.DType.base_type
[narwhals-dev/narwhals#121]: https://github.com/narwhals-dev/narwhals/issues/121
[^1]: Given two data types `A` and `B`, their supertype `S` is the smallest (most specific)
type such that both `A` and `B` can be losslessly cast to without risk of overflow,
truncation, or loss of precision. If no such type exists, the supertype is `None`.
Expand Down
Loading