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

polars/ml/torch.py

@@ -0,0 +1,213 @@
+# mypy: disable-error-code="unused-ignore"
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from polars._utils.unstable import issue_unstable_warning
+from polars.dataframe import DataFrame
+from polars.expr import Expr
+from polars.selectors import exclude
+
+if TYPE_CHECKING:
+    import sys
+    from collections.abc import Sequence
+
+    from torch import Tensor, memory_format
+
+    if sys.version_info >= (3, 11):
+        from typing import Self
+    else:
+        from typing_extensions import Self
+try:
+    import torch
+    from torch.utils.data import TensorDataset
+except ImportError:
+    msg = (
+        "Required package 'torch' not installed.\n"
+        "Please install it using the command `pip install torch`."
+    )
+    raise ImportError(msg) from None
+
+
+__all__ = ["PolarsDataset"]
+
+
+class PolarsDataset(TensorDataset):  # type: ignore[misc]
+    """
+    TensorDataset class specialized for use with Polars DataFrames.
+
+    .. warning::
+        This functionality is considered **unstable**. It may be changed
+        at any point without it being considered a breaking change.
+
+    Parameters
+    ----------
+    frame
+        Polars DataFrame containing the data that will be retrieved as Tensors.
+    label
+        One or more column names or expressions that label the feature data; results
+        in `(features,label)` tuples, where all non-label columns are considered
+        to be features. If no label is designated then each returned item is a
+        simple `(features,)` tuple containing all row elements.
+    features
+        One or more column names or expressions that represent the feature data.
+        If not provided, all columns not designated as labels are considered to be
+        features.
+
+    Notes
+    -----
+    * Integer, slice, range, integer list/Tensor Dataset indexing is all supported.
+    * Designating multi-element labels is also supported.
+
+    Examples
+    --------
+    >>> from torch.utils.data import DataLoader
+    >>> df = pl.DataFrame(
+    ...     data=[
+    ...         (0, 1, 1.5),
+    ...         (1, 0, -0.5),
+    ...         (2, 0, 0.0),
+    ...         (3, 1, -2.25),
+    ...     ],
+    ...     schema=["lbl", "feat1", "feat2"],
+    ...     orient="row",
+    ... )
+
+    Create a Dataset from a Polars DataFrame, standardising the dtype and
+    separating the label/feature columns.
+
+    >>> ds = df.to_torch("dataset", label="lbl", dtype=pl.Float32)
+    >>> ds  # doctest: +IGNORE_RESULT
+    <PolarsDataset [len:4, features:2, labels:1] at 0x156B033B0>
+    >>> ds.features
+    tensor([[ 1.0000,  1.5000],
+            [ 0.0000, -0.5000],
+            [ 0.0000,  0.0000],
+            [ 1.0000, -2.2500]])
+    >>> ds[0]
+    (tensor([1.0000, 1.5000]), tensor(0.))
+
+    The Dataset can be used standalone, or in conjunction with a DataLoader.
+
+    >>> dl = DataLoader(ds, batch_size=2)
+    >>> list(dl)
+    [[tensor([[ 1.0000,  1.5000],
+              [ 0.0000, -0.5000]]),
+      tensor([0., 1.])],
+     [tensor([[ 0.0000,  0.0000],
+              [ 1.0000, -2.2500]]),
+      tensor([2., 3.])]]
+
+    Note that the label can be given as an expression as well as a column name,
+    allowing for independent transform and dtype adjustment from the feature
+    columns.
+
+    >>> ds = df.to_torch(
+    ...     "dataset",
+    ...     dtype=pl.Float32,
+    ...     label=(pl.col("lbl") * 8).cast(pl.Int16),
+    ... )
+    >>> ds[:2]
+    (tensor([[ 1.0000,  1.5000],
+    [ 0.0000, -0.5000]]), tensor([0, 8], dtype=torch.int16))
+    """
+
+    tensors: tuple[Tensor, ...]
+    labels: Tensor | None
+    features: Tensor
+
+    def __init__(
+        self,
+        frame: DataFrame,
+        *,
+        label: str | Expr | Sequence[str | Expr] | None = None,
+        features: str | Expr | Sequence[str | Expr] | None = None,
+    ) -> None:
+        issue_unstable_warning("`PolarsDataset` is considered unstable.")
+        if isinstance(label, (str, Expr)):
+            label = [label]
+
+        label_frame: DataFrame | None = None
+        if not label:
+            feature_frame = frame.select(features) if features else frame
+            self.features = feature_frame.to_torch()
+            self.tensors = (self.features,)
+            self.labels = None
+        else:
+            label_frame = frame.select(*label)
+            self.labels = (  # type: ignore[attr-defined]
+                label_frame if len(label) > 1 else label_frame.to_series()
+            ).to_torch()
+
+            feature_frame = frame.select(
+                features
+                if (isinstance(features, Expr) or features)
+                else exclude(label_frame.columns)
+            )
+            self.features = feature_frame.to_torch()
+            self.tensors = (self.features, self.labels)  # type: ignore[assignment]
+
+        self._n_labels = 0 if (label_frame is None) else label_frame.width
+        self._n_features = feature_frame.width
+
+    def __copy__(self) -> Self:
+        """Return a shallow copy of this PolarsDataset."""
+        dummy_frame = DataFrame({"blank": [0]})
+        dataset_copy = self.__class__(dummy_frame)
+        for attr in (
+            "tensors",
+            "labels",
+            "features",
+            "_n_labels",
+            "_n_features",
+        ):
+            setattr(dataset_copy, attr, getattr(self, attr))
+        return dataset_copy
+
+    def __repr__(self) -> str:
+        """Return a string representation of the PolarsDataset."""
+        return (
+            f"<{type(self).__name__} "
+            f"[len:{len(self)},"
+            f" features:{self._n_features},"
+            f" labels:{self._n_labels}"
+            f"] at 0x{id(self):X}>"
+        )
+
+    def half(
+        self,
+        *,
+        features: bool = True,
+        labels: bool = True,
+        memory_format: memory_format = torch.preserve_format,
+    ) -> Self:
+        """
+        Return a copy of this PolarsDataset with the numeric data converted to f16.
+
+        Parameters
+        ----------
+        features
+            Convert feature data to half precision (f16).
+        labels
+            Convert label data to half precision (f16).
+        memory_format
+            Desired memory format for the modified tensors.
+        """
+        ds = self.__copy__()
+        if features:
+            ds.features = self.features.to(torch.float16, memory_format=memory_format)
+        if self.labels is not None:
+            if labels:
+                ds.labels = self.labels.to(torch.float16, memory_format=memory_format)
+            ds.tensors = (ds.features, ds.labels)  # type: ignore[assignment]
+        else:
+            ds.tensors = (ds.features,)
+        return ds
+
+    @property
+    def schema(self) -> dict[str, torch.dtype | None]:
+        """Return the features/labels schema."""
+        return {
+            "features": self.features.dtype,
+            "labels": self.labels.dtype if self.labels is not None else None,
+        }

polars/ml/utilities.py

@@ -0,0 +1,30 @@
+from typing import Any
+
+from polars import DataFrame
+from polars._dependencies import numpy as np
+from polars._typing import IndexOrder
+from polars.datatypes import Array, List
+
+
+def frame_to_numpy(
+    df: DataFrame,
+    *,
+    writable: bool,
+    target: str,
+    order: IndexOrder = "fortran",
+) -> np.ndarray[Any, Any]:
+    """Convert a DataFrame to a NumPy array for use with Jax or PyTorch."""
+    for nm, tp in df.schema.items():
+        if tp == List:
+            msg = f"cannot convert List column {nm!r} to {target} (use Array dtype instead)"
+            raise TypeError(msg) from None
+
+    if df.width == 1 and df.schema.dtypes()[0] == Array:
+        arr = df[df.columns[0]].to_numpy(writable=writable)
+    else:
+        arr = df.to_numpy(writable=writable, order=order)
+
+    if arr.dtype == object:
+        msg = f"cannot convert DataFrame to {target} (mixed type columns result in `object` dtype)\n{df.schema!r}"
+        raise TypeError(msg)
+    return arr

polars/plugins.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+import contextlib
+import sys
+from functools import lru_cache
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from polars._utils.parse import parse_into_list_of_expressions
+from polars._utils.wrap import wrap_expr
+
+with contextlib.suppress(ImportError):  # Module not available when building docs
+    import polars._plr as plr
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+    from polars import Expr
+    from polars._typing import IntoExpr
+
+__all__ = ["register_plugin_function"]
+
+
+def register_plugin_function(
+    *,
+    plugin_path: Path | str,
+    function_name: str,
+    args: IntoExpr | Iterable[IntoExpr],
+    kwargs: dict[str, Any] | None = None,
+    is_elementwise: bool = False,
+    changes_length: bool = False,
+    returns_scalar: bool = False,
+    cast_to_supertype: bool = False,
+    input_wildcard_expansion: bool = False,
+    pass_name_to_apply: bool = False,
+    use_abs_path: bool = False,
+) -> Expr:
+    """
+    Register a plugin function.
+
+    See the `user guide <https://docs.pola.rs/user-guide/plugins/expr_plugins>`_
+    for more information about plugins.
+
+    Parameters
+    ----------
+    plugin_path
+        Path to the plugin package. Accepts either the file path to the dynamic library
+        file or the path to the directory containing it.
+    function_name
+        The name of the Rust function to register.
+    args
+        The arguments passed to this function. These get passed to the `input`
+        argument on the Rust side, and have to be expressions (or be convertible
+        to expressions).
+    kwargs
+        Non-expression arguments to the plugin function. These must be
+        JSON serializable.
+    is_elementwise
+        Indicate that the function operates on scalars only. This will potentially
+        trigger fast paths.
+    changes_length
+        Indicate that the function will change the length of the expression.
+        For example, a `unique` or `slice` operation.
+    returns_scalar
+        Automatically explode on unit length if the function ran as final aggregation.
+        This is the case for aggregations like `sum`, `min`, `covariance` etc.
+    cast_to_supertype
+        Cast the input expressions to their supertype.
+    input_wildcard_expansion
+        Expand wildcard expressions before executing the function.
+    pass_name_to_apply
+        If set to `True`, the `Series` passed to the function in a group-by operation
+        will ensure the name is set. This is an extra heap allocation per group.
+    use_abs_path
+        If set to `True`, the path will be resolved to an absolute path.
+        The path to the dynamic library is relative to the virtual environment by
+        default.
+
+    Returns
+    -------
+    Expr
+
+    Warnings
+    --------
+    This is highly unsafe as this will call the C function loaded by
+    `plugin::function_name`.
+
+    The parameters you set dictate how Polars will handle the function.
+    Make sure they are correct!
+    """
+    pyexprs = parse_into_list_of_expressions(args)
+    serialized_kwargs = _serialize_kwargs(kwargs)
+    plugin_path = _resolve_plugin_path(plugin_path, use_abs_path=use_abs_path)
+
+    return wrap_expr(
+        plr.register_plugin_function(
+            plugin_path=str(plugin_path),
+            function_name=function_name,
+            args=pyexprs,
+            kwargs=serialized_kwargs,
+            is_elementwise=is_elementwise,
+            input_wildcard_expansion=input_wildcard_expansion,
+            returns_scalar=returns_scalar,
+            cast_to_supertype=cast_to_supertype,
+            pass_name_to_apply=pass_name_to_apply,
+            changes_length=changes_length,
+        )
+    )
+
+
+def _serialize_kwargs(kwargs: dict[str, Any] | None) -> bytes:
+    """Serialize the function's keyword arguments."""
+    if not kwargs:
+        return b""
+
+    import pickle
+
+    # Use the highest pickle protocol supported the serde-pickle crate:
+    # https://docs.rs/serde-pickle/latest/serde_pickle/
+    return pickle.dumps(kwargs, protocol=5)
+
+
+@lru_cache(maxsize=16)
+def _resolve_plugin_path(path: Path | str, *, use_abs_path: bool = False) -> Path:
+    """Get the file path of the dynamic library file."""
+    if not isinstance(path, Path):
+        path = Path(path)
+
+    if path.is_file():
+        return _resolve_file_path(path, use_abs_path=use_abs_path)
+
+    for p in path.iterdir():
+        if _is_dynamic_lib(p):
+            return _resolve_file_path(p, use_abs_path=use_abs_path)
+
+    msg = f"no dynamic library found at path: {path}"
+    raise FileNotFoundError(msg)
+
+
+def _is_dynamic_lib(path: Path) -> bool:
+    return path.is_file() and path.suffix in (".so", ".dll", ".pyd")
+
+
+def _resolve_file_path(path: Path, *, use_abs_path: bool = False) -> Path:
+    venv_path = Path(sys.prefix)
+
+    if use_abs_path:
+        return path.resolve()
+    else:
+        try:
+            file_path = path.relative_to(venv_path)
+        except ValueError:  # Fallback
+            file_path = path.resolve()
+
+    return file_path

polars/pyproject.toml

@@ -0,0 +1,96 @@
+[project]
+name = "polars"
+description = "Blazingly fast DataFrame library"
+readme = "README.md"
+authors = [
+  { name = "Ritchie Vink", email = "ritchie46@gmail.com" },
+]
+version = "1.34.0b2"
+license = { file = "LICENSE" }
+requires-python = ">=3.9"
+
+keywords = ["dataframe", "arrow", "out-of-core"]
+classifiers = [
+  "Development Status :: 5 - Production/Stable",
+  "Environment :: Console",
+  "Intended Audience :: Science/Research",
+  "License :: OSI Approved :: MIT License",
+  "Operating System :: OS Independent",
+  "Programming Language :: Python",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3 :: Only",
+  "Programming Language :: Python :: 3.9",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Programming Language :: Python :: 3.13",
+  "Programming Language :: Rust",
+  "Topic :: Scientific/Engineering",
+  "Typing :: Typed",
+]
+dependencies = ["polars-runtime-32 == 1.34.0b2"]
+
+[project.urls]
+Homepage = "https://www.pola.rs/"
+Documentation = "https://docs.pola.rs/api/python/stable/reference/index.html"
+Repository = "https://github.com/pola-rs/polars"
+Changelog = "https://github.com/pola-rs/polars/releases"
+
+[project.optional-dependencies]
+# Runtimes
+rt64 = ["polars-runtime-64 == 1.34.0b2"]
+rtcompat = ["polars-runtime-compat == 1.34.0b2"]
+
+# NOTE: keep this list in sync with show_versions() and requirements-dev.txt
+polars_cloud = ["polars_cloud >= 0.0.1a1"]
+# Interop
+numpy = ["numpy >= 1.16.0"]
+pandas = ["pandas", "polars[pyarrow]"]
+pyarrow = ["pyarrow >= 7.0.0"]
+pydantic = ["pydantic"]
+
+# Excel
+calamine = ["fastexcel >= 0.9"]
+openpyxl = ["openpyxl >= 3.0.0"]
+xlsx2csv = ["xlsx2csv >= 0.8.0"]
+xlsxwriter = ["xlsxwriter"]
+excel = ["polars[calamine,openpyxl,xlsx2csv,xlsxwriter]"]
+
+# Database
+adbc = ["adbc-driver-manager[dbapi]", "adbc-driver-sqlite[dbapi]"]
+connectorx = ["connectorx >= 0.3.2"]
+sqlalchemy = ["sqlalchemy", "polars[pandas]"]
+database = ["polars[adbc,connectorx,sqlalchemy]"]
+
+# Cloud
+fsspec = ["fsspec"]
+
+# Other I/O
+deltalake = ["deltalake >= 1.0.0"]
+iceberg = ["pyiceberg >= 0.7.1"]
+
+# Other
+async = ["gevent"]
+cloudpickle = ["cloudpickle"]
+graph = ["matplotlib"]
+plot = ["altair >= 5.4.0"]
+style = ["great-tables >= 0.8.0"]
+timezone = ["tzdata; platform_system == 'Windows'"]
+
+# GPU Engine
+gpu = ["cudf-polars-cu12"]
+
+# All
+all = [
+  "polars[async,cloudpickle,database,deltalake,excel,fsspec,graph,iceberg,numpy,pandas,plot,pyarrow,pydantic,style,timezone]",
+]
+
+[[tool.uv.index]]
+name = "pytorch"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
+[tool.uv.sources]
+torch = [
+  { index = "pytorch" },
+]