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

polars/series/struct.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+import sys
+from typing import TYPE_CHECKING
+
+from polars._utils.various import (
+    BUILDING_SPHINX_DOCS,
+    qualified_type_name,
+    sphinx_accessor,
+)
+from polars._utils.wrap import wrap_df
+from polars.schema import Schema
+from polars.series.utils import expr_dispatch
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from polars import DataFrame, Series
+    from polars._plr import PySeries
+elif BUILDING_SPHINX_DOCS:
+    # note: we assign this way to work around an autocomplete issue in ipython/jedi
+    # (ref: https://github.com/davidhalter/jedi/issues/2057)
+    current_module = sys.modules[__name__]
+    current_module.property = sphinx_accessor
+
+
+@expr_dispatch
+class StructNameSpace:
+    """Series.struct namespace."""
+
+    _accessor = "struct"
+
+    def __init__(self, series: Series) -> None:
+        self._s: PySeries = series._s
+
+    def __getitem__(self, item: int | str) -> Series:
+        if isinstance(item, int):
+            return self.field(self.fields[item])
+        elif isinstance(item, str):
+            return self.field(item)
+        else:
+            msg = f"expected type 'int | str', got {qualified_type_name(item)!r}"
+            raise TypeError(msg)
+
+    def _ipython_key_completions_(self) -> list[str]:
+        return self.fields
+
+    @property
+    def fields(self) -> list[str]:
+        """
+        Get the names of the fields.
+
+        Examples
+        --------
+        >>> s = pl.Series([{"a": 1, "b": 2}, {"a": 3, "b": 4}])
+        >>> s.struct.fields
+        ['a', 'b']
+        """
+        if getattr(self, "_s", None) is None:
+            return []
+        return self._s.struct_fields()
+
+    def field(self, name: str) -> Series:
+        """
+        Retrieve one of the fields of this `Struct` as a new Series.
+
+        Parameters
+        ----------
+        name
+            Name of the field.
+
+        Examples
+        --------
+        >>> s = pl.Series([{"a": 1, "b": 2}, {"a": 3, "b": 4}])
+        >>> s.struct.field("a")
+        shape: (2,)
+        Series: 'a' [i64]
+        [
+            1
+            3
+        ]
+        """
+
+    def rename_fields(self, names: Sequence[str]) -> Series:
+        """
+        Rename the fields of the struct.
+
+        Parameters
+        ----------
+        names
+            New names in the order of the struct's fields.
+
+        Examples
+        --------
+        >>> s = pl.Series([{"a": 1, "b": 2}, {"a": 3, "b": 4}])
+        >>> s.struct.fields
+        ['a', 'b']
+        >>> s = s.struct.rename_fields(["c", "d"])
+        >>> s.struct.fields
+        ['c', 'd']
+        """
+
+    @property
+    def schema(self) -> Schema:
+        """
+        Get the struct definition as a name/dtype schema dict.
+
+        Examples
+        --------
+        >>> s = pl.Series([{"a": 1, "b": 2}, {"a": 3, "b": 4}])
+        >>> s.struct.schema
+        Schema({'a': Int64, 'b': Int64})
+        """
+        if getattr(self, "_s", None) is None:
+            return Schema({})
+
+        schema = self._s.dtype().to_schema()
+        return Schema(schema, check_dtypes=False)
+
+    def unnest(self) -> DataFrame:
+        """
+        Convert this struct Series to a DataFrame with a separate column for each field.
+
+        Examples
+        --------
+        >>> s = pl.Series([{"a": 1, "b": 2}, {"a": 3, "b": 4}])
+        >>> s.struct.unnest()
+        shape: (2, 2)
+        ┌─────┬─────┐
+        │ a   ┆ b   │
+        │ --- ┆ --- │
+        │ i64 ┆ i64 │
+        ╞═════╪═════╡
+        │ 1   ┆ 2   │
+        │ 3   ┆ 4   │
+        └─────┴─────┘
+        """
+        return wrap_df(self._s.struct_unnest())
+
+    def json_encode(self) -> Series:
+        """
+        Convert this struct to a string column with json values.
+
+        Examples
+        --------
+        >>> s = pl.Series("a", [{"a": [1, 2], "b": [45]}, {"a": [9, 1, 3], "b": None}])
+        >>> s.struct.json_encode()
+        shape: (2,)
+        Series: 'a' [str]
+        [
+            "{"a":[1,2],"b":[45]}"
+            "{"a":[9,1,3],"b":null}"
+        ]
+        """

polars/series/utils.py

@@ -0,0 +1,191 @@
+from __future__ import annotations
+
+import inspect
+import sys
+from functools import wraps
+from typing import TYPE_CHECKING, Any, Callable, TypeVar
+
+import polars._reexport as pl
+from polars import functions as F
+from polars._utils.wrap import wrap_s
+from polars.datatypes import dtype_to_ffiname
+
+if TYPE_CHECKING:
+    from polars import Series
+    from polars._plr import PySeries
+    from polars._typing import PolarsDataType
+
+    if sys.version_info >= (3, 10):
+        from typing import ParamSpec
+    else:
+        from typing_extensions import ParamSpec
+
+    T = TypeVar("T")
+    P = ParamSpec("P")
+    SeriesMethod = Callable[..., Series]
+
+
+def expr_dispatch(cls: type[T]) -> type[T]:
+    """
+    Series/NameSpace class decorator that sets up expression dispatch.
+
+    * Applied to the Series class, and/or any Series 'NameSpace' classes.
+    * Walks the class attributes, looking for methods that have empty function
+      bodies, with signatures compatible with an existing Expr function.
+    * IFF both conditions are met, the empty method is decorated with @call_expr.
+    """
+    # create lookup of expression functions in this namespace
+    namespace = getattr(cls, "_accessor", None)
+    expr_lookup = _expr_lookup(namespace)
+
+    for name in dir(cls):
+        if (
+            # private
+            not name.startswith("_")
+            # Avoid error when building docs
+            # https://github.com/pola-rs/polars/pull/13238#discussion_r1438787093
+            # TODO: is there a better way to do this?
+            and name != "plot"
+        ):
+            attr = getattr(cls, name)
+            if callable(attr):
+                attr = _undecorated(attr)
+                # note: `co_varnames` starts with the function args, but needs to be
+                # constrained by `co_argcount` as it also includes function-level consts
+                args = attr.__code__.co_varnames[: attr.__code__.co_argcount]
+                # if an expression method with compatible method exists, further check
+                # that the series implementation has an empty function body
+                if (namespace, name, args) in expr_lookup and _is_empty_method(attr):
+                    setattr(cls, name, call_expr(attr))
+    return cls
+
+
+def _expr_lookup(namespace: str | None) -> set[tuple[str | None, str, tuple[str, ...]]]:
+    """Create lookup of potential Expr methods (in the given namespace)."""
+    # dummy Expr object that we can introspect
+    expr = pl.Expr()
+    expr._pyexpr = None  # type: ignore[assignment]
+
+    # optional indirection to "expr.str", "expr.dt", etc
+    if namespace is not None:
+        expr = getattr(expr, namespace)
+
+    lookup = set()
+    for name in dir(expr):
+        if not name.startswith("_"):
+            try:
+                m = getattr(expr, name)
+            except AttributeError:  # may raise for @property methods
+                continue
+            if callable(m):
+                # add function signature (argument names only) to the lookup
+                # as a _possible_ candidate for expression-dispatch
+                m = _undecorated(m)
+                args = m.__code__.co_varnames[: m.__code__.co_argcount]
+                lookup.add((namespace, name, args))
+    return lookup
+
+
+def _undecorated(function: Callable[P, T]) -> Callable[P, T]:
+    """Return the given function without any decorators."""
+    while hasattr(function, "__wrapped__"):
+        function = function.__wrapped__
+    return function
+
+
+def call_expr(func: SeriesMethod) -> SeriesMethod:
+    """Dispatch Series method to an expression implementation."""
+
+    @wraps(func)
+    def wrapper(self: Any, *args: P.args, **kwargs: P.kwargs) -> Series:
+        s = wrap_s(self._s)
+        expr = F.col(s.name)
+        if (namespace := getattr(self, "_accessor", None)) is not None:
+            expr = getattr(expr, namespace)
+        f = getattr(expr, func.__name__)
+        return s.to_frame().select_seq(f(*args, **kwargs)).to_series()
+
+    # note: applying explicit '__signature__' helps IDEs (especially PyCharm)
+    # with proper autocomplete, in addition to what @functools.wraps does
+    setattr(wrapper, "__signature__", inspect.signature(func))  # noqa: B010
+    return wrapper
+
+
+def _is_empty_method(func: SeriesMethod) -> bool:
+    """
+    Confirm that the given function has no implementation.
+
+    Definitions of empty:
+
+    - only has a docstring (body is empty)
+    - has no docstring and just contains 'pass' (or equivalent)
+    """
+    fc = func.__code__
+    return (fc.co_code in _EMPTY_BYTECODE) and (
+        (len(fc.co_consts) == 2 and fc.co_consts[1] is None)
+        # account for optimized-out docstrings (eg: running 'python -OO')
+        or (sys.flags.optimize == 2 and fc.co_consts == (None,))
+    )
+
+
+class _EmptyBytecodeHelper:
+    def __init__(self) -> None:
+        # generate bytecode for empty functions with/without a docstring
+        def _empty_with_docstring() -> None:
+            """"""  # noqa: D419
+
+        def _empty_without_docstring() -> None:
+            pass
+
+        self.empty_bytecode = (
+            _empty_with_docstring.__code__.co_code,
+            _empty_without_docstring.__code__.co_code,
+        )
+
+    def __contains__(self, item: bytes) -> bool:
+        return item in self.empty_bytecode
+
+
+_EMPTY_BYTECODE = _EmptyBytecodeHelper()
+
+
+def get_ffi_func(
+    name: str, dtype: PolarsDataType, obj: PySeries
+) -> Callable[..., Any] | None:
+    """
+    Dynamically obtain the proper FFI function/ method.
+
+    Parameters
+    ----------
+    name
+        function or method name where dtype is replaced by <>
+        for example
+            "call_foo_<>"
+    dtype
+        polars dtype.
+    obj
+        Object to find the method for.
+
+    Returns
+    -------
+    callable or None
+        FFI function, or None if not found.
+    """
+    ffi_name = dtype_to_ffiname(dtype)
+    fname = name.replace("<>", ffi_name)
+    return getattr(obj, fname, None)
+
+
+def _with_no_check_length(func: Callable[..., Any]) -> Any:
+    from polars._plr import check_length
+
+    # Catch any error so that we can be sure that we always restore length checks
+    try:
+        check_length(False)
+        result = func()
+        check_length(True)
+    except Exception:
+        check_length(True)
+        raise
+    else:
+        return result

polars/sql/__init__.py

@@ -0,0 +1,7 @@
+from polars.sql.context import SQLContext
+from polars.sql.functions import sql
+
+__all__ = [
+    "SQLContext",
+    "sql",
+]