polars-runtime-compat 1.34.0b2__cp39-abi3-manylinux_2_24_aarch64.whl

polars/interchange/from_dataframe.py

@@ -0,0 +1,328 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import polars._reexport as pl
+import polars.functions as F
+from polars._utils.various import qualified_type_name
+from polars.datatypes import Boolean, Enum, Int64, String, UInt8, UInt32
+from polars.exceptions import InvalidOperationError
+from polars.interchange.dataframe import PolarsDataFrame
+from polars.interchange.protocol import ColumnNullType, CopyNotAllowedError, DtypeKind
+from polars.interchange.utils import (
+    dtype_to_polars_dtype,
+    get_buffer_length_in_elements,
+    polars_dtype_to_data_buffer_dtype,
+)
+
+if TYPE_CHECKING:
+    from polars import DataFrame, Series
+    from polars._typing import PolarsDataType
+    from polars.interchange.protocol import Buffer, Column, Dtype, SupportsInterchange
+    from polars.interchange.protocol import DataFrame as InterchangeDataFrame
+
+
+def from_dataframe(df: SupportsInterchange, *, allow_copy: bool = True) -> DataFrame:
+    """
+    Build a Polars DataFrame from any dataframe supporting the interchange protocol.
+
+    Parameters
+    ----------
+    df
+        Object supporting the dataframe interchange protocol, i.e. must have implemented
+        the `__dataframe__` method.
+    allow_copy
+        Allow memory to be copied to perform the conversion. If set to False, causes
+        conversions that are not zero-copy to fail.
+    """
+    if isinstance(df, pl.DataFrame):
+        return df
+    elif isinstance(df, PolarsDataFrame):
+        return df._df
+
+    if not hasattr(df, "__dataframe__"):
+        msg = f"`df` of type {qualified_type_name(df)!r} does not support the dataframe interchange protocol"
+        raise TypeError(msg)
+
+    return _from_dataframe(
+        df.__dataframe__(allow_copy=allow_copy),  # type: ignore[arg-type]
+        allow_copy=allow_copy,
+    )
+
+
+def _from_dataframe(df: InterchangeDataFrame, *, allow_copy: bool) -> DataFrame:
+    chunks = []
+    for chunk in df.get_chunks():
+        polars_chunk = _protocol_df_chunk_to_polars(chunk, allow_copy=allow_copy)
+        chunks.append(polars_chunk)
+
+    # Handle implementations that incorrectly yield no chunks for an empty dataframe
+    if not chunks:
+        polars_chunk = _protocol_df_chunk_to_polars(df, allow_copy=allow_copy)
+        chunks.append(polars_chunk)
+
+    return F.concat(chunks, rechunk=False)
+
+
+def _protocol_df_chunk_to_polars(
+    df: InterchangeDataFrame, *, allow_copy: bool
+) -> DataFrame:
+    columns = []
+    for column, name in zip(df.get_columns(), df.column_names()):
+        dtype = dtype_to_polars_dtype(column.dtype)
+        if dtype == String:
+            s = _string_column_to_series(column, allow_copy=allow_copy)
+        elif dtype == Enum:
+            s = _categorical_column_to_series(column, allow_copy=allow_copy)
+        else:
+            s = _column_to_series(column, dtype, allow_copy=allow_copy)
+        columns.append(s.alias(name))
+
+    return pl.DataFrame(columns)
+
+
+def _column_to_series(
+    column: Column, dtype: PolarsDataType, *, allow_copy: bool
+) -> Series:
+    buffers = column.get_buffers()
+    offset = column.offset
+
+    data_buffer = _construct_data_buffer(
+        *buffers["data"], column.size(), offset, allow_copy=allow_copy
+    )
+    validity_buffer = _construct_validity_buffer(
+        buffers["validity"], column, dtype, data_buffer, offset, allow_copy=allow_copy
+    )
+    return pl.Series._from_buffers(dtype, data=data_buffer, validity=validity_buffer)
+
+
+def _string_column_to_series(column: Column, *, allow_copy: bool) -> Series:
+    if column.size() == 0:
+        return pl.Series(dtype=String)
+    elif not allow_copy:
+        msg = "string buffers must be converted"
+        raise CopyNotAllowedError(msg)
+
+    buffers = column.get_buffers()
+    offset = column.offset
+
+    offsets_buffer_info = buffers["offsets"]
+    if offsets_buffer_info is None:
+        msg = "cannot create String column without an offsets buffer"
+        raise RuntimeError(msg)
+    offsets_buffer = _construct_offsets_buffer(
+        *offsets_buffer_info, offset, allow_copy=allow_copy
+    )
+
+    buffer, dtype = buffers["data"]
+    data_buffer = _construct_data_buffer(
+        buffer, dtype, buffer.bufsize, offset=0, allow_copy=allow_copy
+    )
+
+    # First construct a Series without a validity buffer
+    # to allow constructing the validity buffer from a sentinel value
+    data_buffers = [data_buffer, offsets_buffer]
+    data = pl.Series._from_buffers(String, data=data_buffers, validity=None)
+
+    # Add the validity buffer if present
+    validity_buffer = _construct_validity_buffer(
+        buffers["validity"], column, String, data, offset, allow_copy=allow_copy
+    )
+    if validity_buffer is not None:
+        data = pl.Series._from_buffers(
+            String, data=data_buffers, validity=validity_buffer
+        )
+
+    return data
+
+
+def _categorical_column_to_series(column: Column, *, allow_copy: bool) -> Series:
+    categorical = column.describe_categorical
+    if not categorical["is_dictionary"]:
+        msg = "non-dictionary categoricals are not yet supported"
+        raise NotImplementedError(msg)
+
+    categories_col = categorical["categories"]
+    if categories_col.size() == 0:
+        dtype = Enum([])
+    elif categories_col.dtype[0] != DtypeKind.STRING:
+        msg = "non-string categories are not supported"
+        raise NotImplementedError(msg)
+    else:
+        categories = _string_column_to_series(categories_col, allow_copy=allow_copy)
+        dtype = Enum(categories)
+
+    buffers = column.get_buffers()
+    offset = column.offset
+
+    data_buffer = _construct_data_buffer(
+        *buffers["data"], column.size(), offset, allow_copy=allow_copy
+    )
+    validity_buffer = _construct_validity_buffer(
+        buffers["validity"], column, dtype, data_buffer, offset, allow_copy=allow_copy
+    )
+
+    # First construct a physical Series without categories
+    # to allow for sentinel values that do not fit in UInt32
+    data_dtype = data_buffer.dtype
+    out = pl.Series._from_buffers(
+        data_dtype, data=data_buffer, validity=validity_buffer
+    )
+
+    # Polars only supports UInt32 categoricals
+    if data_dtype != UInt32:
+        if not allow_copy and column.size() > 0:
+            msg = f"data buffer must be cast from {data_dtype} to UInt32"
+            raise CopyNotAllowedError(msg)
+
+        # TODO: Cast directly to Enum
+        # https://github.com/pola-rs/polars/issues/13409
+        out = out.cast(UInt32)
+
+    return out.cast(dtype)
+
+
+def _construct_data_buffer(
+    buffer: Buffer,
+    dtype: Dtype,
+    length: int,
+    offset: int = 0,
+    *,
+    allow_copy: bool,
+) -> Series:
+    polars_dtype = dtype_to_polars_dtype(dtype)
+
+    # Handle implementations that incorrectly set the data buffer dtype
+    # to the column dtype
+    # https://github.com/pola-rs/polars/pull/10787
+    polars_dtype = polars_dtype_to_data_buffer_dtype(polars_dtype)
+
+    buffer_info = (buffer.ptr, offset, length)
+
+    # Handle byte-packed boolean buffer
+    if polars_dtype == Boolean and dtype[1] == 8:
+        if length == 0:
+            return pl.Series(dtype=Boolean)
+        elif not allow_copy:
+            msg = "byte-packed boolean buffer must be converted to bit-packed boolean"
+            raise CopyNotAllowedError(msg)
+        return pl.Series._from_buffer(UInt8, buffer_info, owner=buffer).cast(Boolean)
+
+    return pl.Series._from_buffer(polars_dtype, buffer_info, owner=buffer)
+
+
+def _construct_offsets_buffer(
+    buffer: Buffer,
+    dtype: Dtype,
+    offset: int,
+    *,
+    allow_copy: bool,
+) -> Series:
+    polars_dtype = dtype_to_polars_dtype(dtype)
+    length = get_buffer_length_in_elements(buffer.bufsize, dtype) - offset
+
+    buffer_info = (buffer.ptr, offset, length)
+    s = pl.Series._from_buffer(polars_dtype, buffer_info, owner=buffer)
+
+    # Polars only supports Int64 offsets
+    if polars_dtype != Int64:
+        if not allow_copy:
+            msg = f"offsets buffer must be cast from {polars_dtype} to Int64"
+            raise CopyNotAllowedError(msg)
+        s = s.cast(Int64)
+
+    return s
+
+
+def _construct_validity_buffer(
+    validity_buffer_info: tuple[Buffer, Dtype] | None,
+    column: Column,
+    column_dtype: PolarsDataType,
+    data: Series,
+    offset: int = 0,
+    *,
+    allow_copy: bool,
+) -> Series | None:
+    null_type, null_value = column.describe_null
+    if null_type == ColumnNullType.NON_NULLABLE or column.null_count == 0:
+        return None
+
+    elif null_type == ColumnNullType.USE_BITMASK:
+        if validity_buffer_info is None:
+            return None
+        buffer = validity_buffer_info[0]
+        return _construct_validity_buffer_from_bitmask(
+            buffer, null_value, column.size(), offset, allow_copy=allow_copy
+        )
+
+    elif null_type == ColumnNullType.USE_BYTEMASK:
+        if validity_buffer_info is None:
+            return None
+        buffer = validity_buffer_info[0]
+        return _construct_validity_buffer_from_bytemask(
+            buffer, null_value, allow_copy=allow_copy
+        )
+
+    elif null_type == ColumnNullType.USE_NAN:
+        if not allow_copy:
+            msg = "bitmask must be constructed"
+            raise CopyNotAllowedError(msg)
+        return data.is_not_nan()
+
+    elif null_type == ColumnNullType.USE_SENTINEL:
+        if not allow_copy:
+            msg = "bitmask must be constructed"
+            raise CopyNotAllowedError(msg)
+
+        sentinel = pl.Series([null_value])
+        try:
+            if column_dtype.is_temporal():
+                sentinel = sentinel.cast(column_dtype)
+            return data != sentinel  # noqa: TRY300
+        except InvalidOperationError as e:
+            msg = f"invalid sentinel value for column of type {column_dtype}: {null_value!r}"
+            raise TypeError(msg) from e
+
+    else:
+        msg = f"unsupported null type: {null_type!r}"
+        raise NotImplementedError(msg)
+
+
+def _construct_validity_buffer_from_bitmask(
+    buffer: Buffer,
+    null_value: int,
+    length: int,
+    offset: int = 0,
+    *,
+    allow_copy: bool,
+) -> Series:
+    buffer_info = (buffer.ptr, offset, length)
+    s = pl.Series._from_buffer(Boolean, buffer_info, buffer)
+
+    if null_value != 0:
+        if not allow_copy:
+            msg = "bitmask must be inverted"
+            raise CopyNotAllowedError(msg)
+        s = ~s
+
+    return s
+
+
+def _construct_validity_buffer_from_bytemask(
+    buffer: Buffer,
+    null_value: int,
+    *,
+    allow_copy: bool,
+) -> Series:
+    if not allow_copy:
+        msg = "bytemask must be converted into a bitmask"
+        raise CopyNotAllowedError(msg)
+
+    buffer_info = (buffer.ptr, 0, buffer.bufsize)
+    s = pl.Series._from_buffer(UInt8, buffer_info, owner=buffer)
+    s = s.cast(Boolean)
+
+    if null_value != 0:
+        s = ~s
+
+    return s

polars/interchange/protocol.py

@@ -0,0 +1,303 @@
+from __future__ import annotations
+
+from enum import IntEnum
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    Literal,
+    Protocol,
+    TypedDict,
+)
+
+from polars._utils.unstable import issue_unstable_warning
+
+if TYPE_CHECKING:
+    import sys
+    from collections.abc import Iterable, Sequence
+
+    from polars.interchange.buffer import PolarsBuffer
+    from polars.interchange.column import PolarsColumn
+
+    if sys.version_info >= (3, 10):
+        from typing import TypeAlias
+    else:
+        from typing_extensions import TypeAlias
+
+
+class DlpackDeviceType(IntEnum):
+    """Integer enum for device type codes matching DLPack."""
+
+    CPU = 1
+    CUDA = 2
+    CPU_PINNED = 3
+    OPENCL = 4
+    VULKAN = 7
+    METAL = 8
+    VPI = 9
+    ROCM = 10
+
+
+class DtypeKind(IntEnum):
+    """
+    Integer enum for data types.
+
+    Attributes
+    ----------
+    INT : int
+        Matches to signed integer data type.
+    UINT : int
+        Matches to unsigned integer data type.
+    FLOAT : int
+        Matches to floating point data type.
+    BOOL : int
+        Matches to boolean data type.
+    STRING : int
+        Matches to string data type (UTF-8 encoded).
+    DATETIME : int
+        Matches to datetime data type.
+    CATEGORICAL : int
+        Matches to categorical data type.
+    """
+
+    INT = 0
+    UINT = 1
+    FLOAT = 2
+    BOOL = 20
+    STRING = 21  # UTF-8
+    DATETIME = 22
+    CATEGORICAL = 23
+
+
+Dtype: TypeAlias = tuple[DtypeKind, int, str, str]  # see Column.dtype
+
+
+class ColumnNullType(IntEnum):
+    """
+    Integer enum for null type representation.
+
+    Attributes
+    ----------
+    NON_NULLABLE : int
+        Non-nullable column.
+    USE_NAN : int
+        Use explicit float NaN value.
+    USE_SENTINEL : int
+        Sentinel value besides NaN.
+    USE_BITMASK : int
+        The bit is set/unset representing a null on a certain position.
+    USE_BYTEMASK : int
+        The byte is set/unset representing a null on a certain position.
+    """
+
+    NON_NULLABLE = 0
+    USE_NAN = 1
+    USE_SENTINEL = 2
+    USE_BITMASK = 3
+    USE_BYTEMASK = 4
+
+
+class ColumnBuffers(TypedDict):
+    """Buffers backing a column."""
+
+    # first element is a buffer containing the column data;
+    # second element is the data buffer's associated dtype
+    data: tuple[PolarsBuffer, Dtype]
+
+    # first element is a buffer containing mask values indicating missing data;
+    # second element is the mask value buffer's associated dtype.
+    # None if the null representation is not a bit or byte mask
+    validity: tuple[PolarsBuffer, Dtype] | None
+
+    # first element is a buffer containing the offset values for
+    # variable-size binary data (e.g., variable-length strings);
+    # second element is the offsets buffer's associated dtype.
+    # None if the data buffer does not have an associated offsets buffer
+    offsets: tuple[PolarsBuffer, Dtype] | None
+
+
+class CategoricalDescription(TypedDict):
+    """Description of a categorical column."""
+
+    # whether the ordering of dictionary indices is semantically meaningful
+    is_ordered: bool
+    # whether a dictionary-style mapping of categorical values to other objects exists
+    is_dictionary: Literal[True]
+    # Python-level only (e.g. `{int: str}`).
+    # None if not a dictionary-style categorical.
+    categories: PolarsColumn
+
+
+class Buffer(Protocol):
+    """Interchange buffer object."""
+
+    @property
+    def bufsize(self) -> int:
+        """Buffer size in bytes."""
+
+    @property
+    def ptr(self) -> int:
+        """Pointer to start of the buffer as an integer."""
+
+    def __dlpack__(self) -> Any:
+        """Represent this structure as DLPack interface."""
+
+    def __dlpack_device__(self) -> tuple[DlpackDeviceType, int | None]:
+        """Device type and device ID for where the data in the buffer resides."""
+
+
+class Column(Protocol):
+    """Interchange column object."""
+
+    def size(self) -> int:
+        """Size of the column in elements."""
+
+    @property
+    def offset(self) -> int:
+        """Offset of the first element with respect to the start of the underlying buffer."""  # noqa: W505
+
+    @property
+    def dtype(self) -> Dtype:
+        """Data type of the column."""
+
+    @property
+    def describe_categorical(self) -> CategoricalDescription:
+        """Description of the categorical data type of the column."""
+
+    @property
+    def describe_null(self) -> tuple[ColumnNullType, Any]:
+        """Description of the null representation the column uses."""
+
+    @property
+    def null_count(self) -> int | None:
+        """Number of null elements, if known."""
+
+    @property
+    def metadata(self) -> dict[str, Any]:
+        """The metadata for the column."""
+
+    def num_chunks(self) -> int:
+        """Return the number of chunks the column consists of."""
+
+    def get_chunks(self, n_chunks: int | None = None) -> Iterable[Column]:
+        """Return an iterator yielding the column chunks."""
+
+    def get_buffers(self) -> ColumnBuffers:
+        """Return a dictionary containing the underlying buffers."""
+
+
+class DataFrame(Protocol):
+    """Interchange dataframe object."""
+
+    version: ClassVar[int]  # Version of the protocol
+
+    def __dataframe__(
+        self,
+        nan_as_null: bool = False,  # noqa: FBT001
+        allow_copy: bool = True,  # noqa: FBT001
+    ) -> DataFrame:
+        """Convert to a dataframe object implementing the dataframe interchange protocol."""  # noqa: W505
+
+    @property
+    def metadata(self) -> dict[str, Any]:
+        """The metadata for the dataframe."""
+
+    def num_columns(self) -> int:
+        """Return the number of columns in the dataframe."""
+
+    def num_rows(self) -> int | None:
+        """Return the number of rows in the dataframe, if available."""
+
+    def num_chunks(self) -> int:
+        """Return the number of chunks the dataframe consists of.."""
+
+    def column_names(self) -> Iterable[str]:
+        """Return the column names."""
+
+    def get_column(self, i: int) -> Column:
+        """Return the column at the indicated position."""
+
+    def get_column_by_name(self, name: str) -> Column:
+        """Return the column with the given name."""
+
+    def get_columns(self) -> Iterable[Column]:
+        """Return an iterator yielding the columns."""
+
+    def select_columns(self, indices: Sequence[int]) -> DataFrame:
+        """Create a new dataframe by selecting a subset of columns by index."""
+
+    def select_columns_by_name(self, names: Sequence[str]) -> DataFrame:
+        """Create a new dataframe by selecting a subset of columns by name."""
+
+    def get_chunks(self, n_chunks: int | None = None) -> Iterable[DataFrame]:
+        """Return an iterator yielding the chunks of the dataframe."""
+
+
+class SupportsInterchange(Protocol):
+    """Dataframe that supports conversion into an interchange dataframe object."""
+
+    def __dataframe__(
+        self,
+        nan_as_null: bool = False,  # noqa: FBT001
+        allow_copy: bool = True,  # noqa: FBT001
+    ) -> SupportsInterchange:
+        """Convert to a dataframe object implementing the dataframe interchange protocol."""  # noqa: W505
+
+
+class Endianness:
+    """Enum indicating the byte-order of a data type."""
+
+    LITTLE = "<"
+    BIG = ">"
+    NATIVE = "="
+    NA = "|"
+
+
+class CopyNotAllowedError(RuntimeError):
+    """Exception raised when a copy is required, but `allow_copy` is set to `False`."""
+
+
+class CompatLevel:
+    """Data structure compatibility level."""
+
+    _version: int
+
+    def __init__(self) -> None:
+        msg = "it is not allowed to create a CompatLevel object"
+        raise TypeError(msg)
+
+    @staticmethod
+    def _with_version(version: int) -> CompatLevel:
+        compat_level = CompatLevel.__new__(CompatLevel)
+        compat_level._version = version
+        return compat_level
+
+    @staticmethod
+    def _newest() -> CompatLevel:
+        return CompatLevel._future1  # type: ignore[attr-defined]
+
+    @staticmethod
+    def newest() -> CompatLevel:
+        """
+        Get the highest supported compatibility level.
+
+        .. warning::
+            Highest compatibility level is considered **unstable**. It may be changed
+            at any point without it being considered a breaking change.
+        """
+        issue_unstable_warning(
+            "using the highest compatibility level is considered unstable."
+        )
+        return CompatLevel._newest()
+
+    @staticmethod
+    def oldest() -> CompatLevel:
+        """Get the most compatible level."""
+        return CompatLevel._compatible  # type: ignore[attr-defined]
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__module__}.{self.__class__.__qualname__}: {self._version}>"
+
+
+CompatLevel._compatible = CompatLevel._with_version(0)  # type: ignore[attr-defined]
+CompatLevel._future1 = CompatLevel._with_version(1)  # type: ignore[attr-defined]