polars-runtime-compat 1.34.0b2__cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

polars/_utils/nest_asyncio.py

@@ -0,0 +1,264 @@
+#
+# Originally vendored from https://github.com/erdewit/nest_asyncio
+#
+
+# BSD 2-Clause License
+
+# Copyright (c) 2025 Ritchie Vink
+# Copyright (c) 2018-2020, Ewald de Wit
+# All rights reserved.
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# * Redistributions of source code must retain the above copyright notice, this
+#   list of conditions and the following disclaimer.
+
+# * Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+# Ignore all lints, file is currently copied.
+# ruff: noqa
+# type: ignore
+
+"""Patch asyncio to allow nested event loops."""
+
+import asyncio
+import asyncio.events as events
+import os
+import sys
+import threading
+from contextlib import contextmanager, suppress
+from heapq import heappop
+
+
+def apply(loop=None):
+    """Patch asyncio to make its event loop reentrant."""
+    _patch_asyncio()
+    _patch_policy()
+    _patch_tornado()
+
+    loop = loop or asyncio.get_event_loop()
+    _patch_loop(loop)
+
+
+def _patch_asyncio():
+    """Patch asyncio module to use pure Python tasks and futures."""
+
+    def run(main, *, debug=False):
+        loop = asyncio.get_event_loop()
+        loop.set_debug(debug)
+        task = asyncio.ensure_future(main)
+        try:
+            return loop.run_until_complete(task)
+        finally:
+            if not task.done():
+                task.cancel()
+                with suppress(asyncio.CancelledError):
+                    loop.run_until_complete(task)
+
+    def _get_event_loop(stacklevel=3):
+        loop = events._get_running_loop()
+        if loop is None:
+            loop = events.get_event_loop_policy().get_event_loop()
+        return loop
+
+    # Use module level _current_tasks, all_tasks and patch run method.
+    if hasattr(asyncio, "_nest_patched"):
+        return
+    if sys.version_info >= (3, 6, 0):
+        asyncio.Task = asyncio.tasks._CTask = asyncio.tasks.Task = asyncio.tasks._PyTask
+        asyncio.Future = asyncio.futures._CFuture = asyncio.futures.Future = (
+            asyncio.futures._PyFuture
+        )
+    if sys.version_info < (3, 7, 0):
+        asyncio.tasks._current_tasks = asyncio.tasks.Task._current_tasks
+        asyncio.all_tasks = asyncio.tasks.Task.all_tasks
+    if sys.version_info >= (3, 9, 0):
+        events._get_event_loop = events.get_event_loop = asyncio.get_event_loop = (
+            _get_event_loop
+        )
+    asyncio.run = run
+    asyncio._nest_patched = True
+
+
+def _patch_policy():
+    """Patch the policy to always return a patched loop."""
+
+    def get_event_loop(self):
+        if self._local._loop is None:
+            loop = self.new_event_loop()
+            _patch_loop(loop)
+            self.set_event_loop(loop)
+        return self._local._loop
+
+    policy = events.get_event_loop_policy()
+    policy.__class__.get_event_loop = get_event_loop
+
+
+def _patch_loop(loop):
+    """Patch loop to make it reentrant."""
+
+    def run_forever(self):
+        with manage_run(self), manage_asyncgens(self):
+            while True:
+                self._run_once()
+                if self._stopping:
+                    break
+        self._stopping = False
+
+    def run_until_complete(self, future):
+        with manage_run(self):
+            f = asyncio.ensure_future(future, loop=self)
+            if f is not future:
+                f._log_destroy_pending = False
+            while not f.done():
+                self._run_once()
+                if self._stopping:
+                    break
+            if not f.done():
+                raise RuntimeError("Event loop stopped before Future completed.")
+            return f.result()
+
+    def _run_once(self):
+        """
+        Simplified re-implementation of asyncio's _run_once that
+        runs handles as they become ready.
+        """
+        ready = self._ready
+        scheduled = self._scheduled
+        while scheduled and scheduled[0]._cancelled:
+            heappop(scheduled)
+
+        timeout = (
+            0
+            if ready or self._stopping
+            else min(max(scheduled[0]._when - self.time(), 0), 86400)
+            if scheduled
+            else None
+        )
+        event_list = self._selector.select(timeout)
+        self._process_events(event_list)
+
+        end_time = self.time() + self._clock_resolution
+        while scheduled and scheduled[0]._when < end_time:
+            handle = heappop(scheduled)
+            ready.append(handle)
+
+        for _ in range(len(ready)):
+            if not ready:
+                break
+            handle = ready.popleft()
+            if not handle._cancelled:
+                # preempt the current task so that that checks in
+                # Task.__step do not raise
+                curr_task = curr_tasks.pop(self, None)
+
+                try:
+                    handle._run()
+                finally:
+                    # restore the current task
+                    if curr_task is not None:
+                        curr_tasks[self] = curr_task
+
+        handle = None
+
+    @contextmanager
+    def manage_run(self):
+        """Set up the loop for running."""
+        self._check_closed()
+        old_thread_id = self._thread_id
+        old_running_loop = events._get_running_loop()
+        try:
+            self._thread_id = threading.get_ident()
+            events._set_running_loop(self)
+            self._num_runs_pending += 1
+            if self._is_proactorloop:
+                if self._self_reading_future is None:
+                    self.call_soon(self._loop_self_reading)
+            yield
+        finally:
+            self._thread_id = old_thread_id
+            events._set_running_loop(old_running_loop)
+            self._num_runs_pending -= 1
+            if self._is_proactorloop:
+                if (
+                    self._num_runs_pending == 0
+                    and self._self_reading_future is not None
+                ):
+                    ov = self._self_reading_future._ov
+                    self._self_reading_future.cancel()
+                    if ov is not None:
+                        self._proactor._unregister(ov)
+                    self._self_reading_future = None
+
+    @contextmanager
+    def manage_asyncgens(self):
+        if not hasattr(sys, "get_asyncgen_hooks"):
+            # Python version is too old.
+            return
+        old_agen_hooks = sys.get_asyncgen_hooks()
+        try:
+            self._set_coroutine_origin_tracking(self._debug)
+            if self._asyncgens is not None:
+                sys.set_asyncgen_hooks(
+                    firstiter=self._asyncgen_firstiter_hook,
+                    finalizer=self._asyncgen_finalizer_hook,
+                )
+            yield
+        finally:
+            self._set_coroutine_origin_tracking(False)
+            if self._asyncgens is not None:
+                sys.set_asyncgen_hooks(*old_agen_hooks)
+
+    def _check_running(self):
+        """Do not throw exception if loop is already running."""
+        pass
+
+    if hasattr(loop, "_nest_patched"):
+        return
+    if not isinstance(loop, asyncio.BaseEventLoop):
+        raise ValueError("Can't patch loop of type %s" % type(loop))
+    cls = loop.__class__
+    cls.run_forever = run_forever
+    cls.run_until_complete = run_until_complete
+    cls._run_once = _run_once
+    cls._check_running = _check_running
+    cls._check_runnung = _check_running  # typo in Python 3.7 source
+    cls._num_runs_pending = 1 if loop.is_running() else 0
+    cls._is_proactorloop = os.name == "nt" and issubclass(
+        cls, asyncio.ProactorEventLoop
+    )
+    if sys.version_info < (3, 7, 0):
+        cls._set_coroutine_origin_tracking = cls._set_coroutine_wrapper
+    curr_tasks = (
+        asyncio.tasks._current_tasks
+        if sys.version_info >= (3, 7, 0)
+        else asyncio.Task._current_tasks
+    )
+    cls._nest_patched = True
+
+
+def _patch_tornado():
+    """
+    If tornado is imported before nest_asyncio, make tornado aware of
+    the pure-Python asyncio Future.
+    """
+    if "tornado" in sys.modules:
+        import tornado.concurrent as tc  # type: ignore
+
+        tc.Future = asyncio.Future
+        if asyncio.Future not in tc.FUTURES:
+            tc.FUTURES += (asyncio.Future,)

polars/_utils/parquet.py

@@ -0,0 +1,15 @@
+from typing import Any, Callable
+
+from polars._typing import ParquetMetadataContext, ParquetMetadataFn
+
+
+def wrap_parquet_metadata_callback(
+    fn: ParquetMetadataFn,
+) -> Callable[[Any], list[tuple[str, str]]]:
+    def pyo3_compatible_callback(ctx: Any) -> list[tuple[str, str]]:
+        ctx_py = ParquetMetadataContext(
+            arrow_schema=ctx.arrow_schema,
+        )
+        return list(fn(ctx_py).items())
+
+    return pyo3_compatible_callback

polars/_utils/parse/__init__.py

@@ -0,0 +1,12 @@
+from polars._utils.parse.expr import (
+    parse_into_expression,
+    parse_into_list_of_expressions,
+    parse_predicates_constraints_into_expression,
+)
+
+__all__ = [
+    # expr
+    "parse_into_expression",
+    "parse_into_list_of_expressions",
+    "parse_predicates_constraints_into_expression",
+]

polars/_utils/parse/expr.py

@@ -0,0 +1,242 @@
+from __future__ import annotations
+
+import contextlib
+from collections.abc import Collection, Iterable, Mapping
+from typing import TYPE_CHECKING, Any
+
+import polars._reexport as pl
+from polars import functions as F
+from polars._utils.various import qualified_type_name
+from polars.exceptions import ComputeError
+
+with contextlib.suppress(ImportError):  # Module not available when building docs
+    import polars._plr as plr
+
+if TYPE_CHECKING:
+    from polars import Expr
+    from polars._plr import PyExpr
+    from polars._typing import ColumnNameOrSelector, IntoExpr, PolarsDataType
+
+
+def parse_into_expression(
+    input: IntoExpr,
+    *,
+    str_as_lit: bool = False,
+    list_as_series: bool = False,
+    structify: bool = False,
+    dtype: PolarsDataType | None = None,
+) -> PyExpr:
+    """
+    Parse a single input into an expression.
+
+    Parameters
+    ----------
+    input
+        The input to be parsed as an expression.
+    str_as_lit
+        Interpret string input as a string literal. If set to `False` (default),
+        strings are parsed as column names.
+    list_as_series
+        Interpret list input as a Series literal. If set to `False` (default),
+        lists are parsed as list literals.
+    structify
+        Convert multi-column expressions to a single struct expression.
+    dtype
+        If the input is expected to resolve to a literal with a known dtype, pass
+        this to the `lit` constructor.
+
+    Returns
+    -------
+    PyExpr
+    """
+    if isinstance(input, pl.Expr):
+        expr = input
+        if structify:
+            expr = _structify_expression(expr)
+
+    elif isinstance(input, str) and not str_as_lit:
+        expr = F.col(input)
+    elif isinstance(input, list) and list_as_series:
+        expr = F.lit(pl.Series(input), dtype=dtype)
+    else:
+        expr = F.lit(input, dtype=dtype)
+
+    return expr._pyexpr
+
+
+def _structify_expression(expr: Expr) -> Expr:
+    unaliased_expr = expr.meta.undo_aliases()
+    if unaliased_expr.meta.has_multiple_outputs():
+        try:
+            expr_name = expr.meta.output_name()
+        except ComputeError:
+            expr = F.struct(expr)
+        else:
+            expr = F.struct(unaliased_expr).alias(expr_name)
+    return expr
+
+
+def parse_into_list_of_expressions(
+    *inputs: IntoExpr | Iterable[IntoExpr],
+    __structify: bool = False,
+    **named_inputs: IntoExpr,
+) -> list[PyExpr]:
+    """
+    Parse multiple inputs into a list of expressions.
+
+    Parameters
+    ----------
+    *inputs
+        Inputs to be parsed as expressions, specified as positional arguments.
+    **named_inputs
+        Additional inputs to be parsed as expressions, specified as keyword arguments.
+        The expressions will be renamed to the keyword used.
+    __structify
+        Convert multi-column expressions to a single struct expression.
+
+    Returns
+    -------
+    list of PyExpr
+    """
+    exprs = _parse_positional_inputs(inputs, structify=__structify)  # type: ignore[arg-type]
+    if named_inputs:
+        named_exprs = _parse_named_inputs(named_inputs, structify=__structify)
+        exprs.extend(named_exprs)
+
+    return exprs
+
+
+def parse_into_selector(
+    i: ColumnNameOrSelector,
+    *,
+    strict: bool = True,
+) -> pl.Selector:
+    if isinstance(i, str):
+        import polars.selectors as cs
+
+        return cs.by_name([i], require_all=strict)
+    elif isinstance(i, pl.Selector):
+        return i
+    elif isinstance(i, pl.Expr):
+        return i.meta.as_selector()
+    else:
+        msg = f"cannot turn {qualified_type_name(i)!r} into selector"
+        raise TypeError(msg)
+
+
+def parse_list_into_selector(
+    inputs: ColumnNameOrSelector | Collection[ColumnNameOrSelector],
+    *,
+    strict: bool = True,
+) -> pl.Selector:
+    if isinstance(inputs, Collection) and not isinstance(inputs, str):
+        import polars.selectors as cs
+
+        columns = list(filter(lambda i: isinstance(i, str), inputs))
+        selector = cs.by_name(columns, require_all=strict)  # type: ignore[arg-type]
+
+        if len(columns) == len(inputs):
+            return selector
+
+        # A bit cleaner
+        if len(columns) == 0:
+            selector = cs.empty()
+
+        for i in inputs:
+            selector |= parse_into_selector(i, strict=strict)
+        return selector
+    else:
+        return parse_into_selector(inputs, strict=strict)
+
+
+def _parse_positional_inputs(
+    inputs: tuple[IntoExpr, ...] | tuple[Iterable[IntoExpr]],
+    *,
+    structify: bool = False,
+) -> list[PyExpr]:
+    inputs_iter = _parse_inputs_as_iterable(inputs)
+    return [parse_into_expression(e, structify=structify) for e in inputs_iter]
+
+
+def _parse_inputs_as_iterable(
+    inputs: tuple[Any, ...] | tuple[Iterable[Any]],
+) -> Iterable[Any]:
+    if not inputs:
+        return []
+
+    # Ensures that the outermost element cannot be a Dictionary (as an iterable)
+    if len(inputs) == 1 and isinstance(inputs[0], Mapping):
+        msg = (
+            "Cannot pass a dictionary as a single positional argument.\n"
+            "If you merely want the *keys*, use:\n"
+            "  • df.method(*your_dict.keys())\n"
+            "If you need the key value pairs, use one of:\n"
+            "  • unpack as keywords:    df.method(**your_dict)\n"
+            "  • build expressions:     df.method(expr.alias(k) for k, expr in your_dict.items())"
+        )
+        raise TypeError(msg)
+
+    # Treat elements of a single iterable as separate inputs
+    if len(inputs) == 1 and _is_iterable(inputs[0]):
+        return inputs[0]
+
+    return inputs
+
+
+def _is_iterable(input: Any | Iterable[Any]) -> bool:
+    return isinstance(input, Iterable) and not isinstance(
+        input, (str, bytes, pl.Series)
+    )
+
+
+def _parse_named_inputs(
+    named_inputs: dict[str, IntoExpr], *, structify: bool = False
+) -> Iterable[PyExpr]:
+    for name, input in named_inputs.items():
+        yield parse_into_expression(input, structify=structify).alias(name)
+
+
+def parse_predicates_constraints_into_expression(
+    *predicates: IntoExpr | Iterable[IntoExpr],
+    **constraints: Any,
+) -> PyExpr:
+    """
+    Parse predicates and constraints into a single expression.
+
+    The result is an AND-reduction of all inputs.
+
+    Parameters
+    ----------
+    *predicates
+        Predicates to be parsed, specified as positional arguments.
+    **constraints
+        Constraints to be parsed, specified as keyword arguments.
+        These will be converted to predicates of the form "keyword equals input value".
+
+    Returns
+    -------
+    PyExpr
+    """
+    all_predicates = _parse_positional_inputs(predicates)  # type: ignore[arg-type]
+
+    if constraints:
+        constraint_predicates = _parse_constraints(constraints)
+        all_predicates.extend(constraint_predicates)
+
+    return _combine_predicates(all_predicates)
+
+
+def _parse_constraints(constraints: dict[str, IntoExpr]) -> Iterable[PyExpr]:
+    for name, value in constraints.items():
+        yield F.col(name).eq(value)._pyexpr
+
+
+def _combine_predicates(predicates: list[PyExpr]) -> PyExpr:
+    if not predicates:
+        msg = "at least one predicate or constraint must be provided"
+        raise TypeError(msg)
+
+    if len(predicates) == 1:
+        return predicates[0]
+
+    return plr.all_horizontal(predicates)

polars/_utils/polars_version.py

@@ -0,0 +1,19 @@
+try:
+    import polars._plr as plr
+
+    _POLARS_VERSION = plr.__version__
+except ImportError:
+    # This is only useful for documentation
+    import warnings
+
+    warnings.warn("Polars binary is missing!", stacklevel=2)
+    _POLARS_VERSION = ""
+
+
+def get_polars_version() -> str:
+    """
+    Return the version of the Python Polars package as a string.
+
+    If the Polars binary is missing, returns an empty string.
+    """
+    return _POLARS_VERSION

polars/_utils/pycapsule.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import contextlib
+from typing import TYPE_CHECKING, Any
+
+from polars._utils.construction.dataframe import dataframe_to_pydf
+from polars._utils.wrap import wrap_df, wrap_s
+
+with contextlib.suppress(ImportError):
+    from polars._plr import PySeries
+
+if TYPE_CHECKING:
+    from polars import DataFrame
+    from polars._typing import SchemaDefinition, SchemaDict
+
+
+def is_pycapsule(obj: Any) -> bool:
+    """Check if object supports the PyCapsule interface."""
+    return hasattr(obj, "__arrow_c_stream__") or hasattr(obj, "__arrow_c_array__")
+
+
+def pycapsule_to_frame(
+    obj: Any,
+    *,
+    schema: SchemaDefinition | None = None,
+    schema_overrides: SchemaDict | None = None,
+    rechunk: bool = False,
+) -> DataFrame:
+    """Convert PyCapsule object to DataFrame."""
+    if hasattr(obj, "__arrow_c_array__"):
+        # This uses the fact that PySeries.from_arrow_c_array will create a
+        # struct-typed Series. Then we unpack that to a DataFrame.
+        tmp_col_name = ""
+        s = wrap_s(PySeries.from_arrow_c_array(obj))
+        df = s.to_frame(tmp_col_name).unnest(tmp_col_name)
+
+    elif hasattr(obj, "__arrow_c_stream__"):
+        # This uses the fact that PySeries.from_arrow_c_stream will create a
+        # struct-typed Series. Then we unpack that to a DataFrame.
+        tmp_col_name = ""
+        s = wrap_s(PySeries.from_arrow_c_stream(obj))
+        df = s.to_frame(tmp_col_name).unnest(tmp_col_name)
+    else:
+        msg = f"object does not support PyCapsule interface; found {obj!r} "
+        raise TypeError(msg)
+
+    if rechunk:
+        df = df.rechunk()
+    if schema or schema_overrides:
+        df = wrap_df(
+            dataframe_to_pydf(df, schema=schema, schema_overrides=schema_overrides)
+        )
+    return df

polars/_utils/scan.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from polars import DataFrame
+
+
+def _execute_from_rust(
+    function: Any, with_columns: list[str] | None, *args: Any
+) -> DataFrame:
+    """
+    Deserialize and execute the given function for the projected columns.
+
+    Called from polars-lazy. Polars-lazy provides the bytes of the pickled function and
+    the projected columns.
+
+    Parameters
+    ----------
+    function
+        function object
+    with_columns
+        Columns that are projected
+    *args
+        Additional function arguments.
+    """
+    return function(with_columns, *args)

polars/_utils/serde.py

@@ -0,0 +1,63 @@
+"""Utility for serializing Polars objects."""
+
+from __future__ import annotations
+
+from io import BytesIO, StringIO
+from pathlib import Path
+from typing import TYPE_CHECKING, Callable, Literal, overload
+
+from polars._utils.various import normalize_filepath
+
+if TYPE_CHECKING:
+    from io import IOBase
+
+    from polars._typing import SerializationFormat
+
+
+@overload
+def serialize_polars_object(
+    serializer: Callable[[IOBase | str], None], file: None, format: Literal["binary"]
+) -> bytes: ...
+@overload
+def serialize_polars_object(
+    serializer: Callable[[IOBase | str], None], file: None, format: Literal["json"]
+) -> str: ...
+@overload
+def serialize_polars_object(
+    serializer: Callable[[IOBase | str], None],
+    file: IOBase | str | Path,
+    format: SerializationFormat,
+) -> None: ...
+
+
+def serialize_polars_object(
+    serializer: Callable[[IOBase | str], None],
+    file: IOBase | str | Path | None,
+    format: SerializationFormat,
+) -> bytes | str | None:
+    """Serialize a Polars object (DataFrame/LazyFrame/Expr)."""
+
+    def serialize_to_bytes() -> bytes:
+        with BytesIO() as buf:
+            serializer(buf)
+            serialized = buf.getvalue()
+        return serialized
+
+    if file is None:
+        serialized = serialize_to_bytes()
+        return serialized.decode() if format == "json" else serialized
+    elif isinstance(file, StringIO):
+        serialized_str = serialize_to_bytes().decode()
+        file.write(serialized_str)
+        return None
+    elif isinstance(file, BytesIO):
+        serialized = serialize_to_bytes()
+        file.write(serialized)
+        return None
+    elif isinstance(file, (str, Path)):
+        file = normalize_filepath(file)
+        serializer(file)
+        return None
+    else:
+        serializer(file)
+        return None