366 lines
12 KiB
Python
366 lines
12 KiB
Python
|
"""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()
|
||
|
|
)
|