polars-runtime-compat 1.34.0b3__cp39-abi3-manylinux_2_24_aarch64.whl → 1.34.0b5__cp39-abi3-manylinux_2_24_aarch64.whl

polars/sql/context.py

@@ -1,677 +0,0 @@
-from __future__ import annotations
-
-import contextlib
-import re
-from typing import (
-    TYPE_CHECKING,
-    Callable,
-    Generic,
-    Union,
-    overload,
-)
-
-from polars._dependencies import _check_for_pandas, _check_for_pyarrow
-from polars._dependencies import pandas as pd
-from polars._dependencies import pyarrow as pa
-from polars._typing import FrameType
-from polars._utils.deprecation import deprecate_renamed_parameter
-from polars._utils.pycapsule import is_pycapsule
-from polars._utils.unstable import issue_unstable_warning
-from polars._utils.various import _get_stack_locals, qualified_type_name
-from polars._utils.wrap import wrap_ldf
-from polars.convert import from_arrow, from_pandas
-from polars.dataframe import DataFrame
-from polars.lazyframe import LazyFrame
-from polars.series import Series
-
-with contextlib.suppress(ImportError):  # Module not available when building docs
-    from polars._plr import PySQLContext
-
-if TYPE_CHECKING:
-    import sys
-    from collections.abc import Collection, Mapping
-    from types import TracebackType
-    from typing import Any, Final, Literal
-
-    if sys.version_info >= (3, 10):
-        from typing import TypeAlias
-    else:
-        from typing_extensions import TypeAlias
-
-    if sys.version_info >= (3, 11):
-        from typing import Self
-    else:
-        from typing_extensions import Self
-
-    CompatibleFrameType: TypeAlias = Union[
-        DataFrame,
-        LazyFrame,
-        Series,
-        pd.DataFrame,
-        pd.Series[Any],
-        pa.Table,
-        pa.RecordBatch,
-    ]
-
-__all__ = ["SQLContext"]
-
-
-def _compatible_frame(obj: Any) -> bool:
-    """Check if the object can be converted to DataFrame."""
-    return (
-        is_pycapsule(obj)
-        or isinstance(obj, LazyFrame)
-        or (_check_for_pandas(obj) and isinstance(obj, (pd.DataFrame, pd.Series)))
-        or (_check_for_pyarrow(obj) and isinstance(obj, (pa.Table, pa.RecordBatch)))
-    )
-
-
-def _ensure_lazyframe(obj: Any) -> LazyFrame:
-    """Return LazyFrame from compatible input."""
-    if isinstance(obj, (DataFrame, LazyFrame)):
-        return obj.lazy()
-    elif isinstance(obj, Series):
-        return obj.to_frame().lazy()
-    elif _check_for_pandas(obj) and isinstance(obj, (pd.DataFrame, pd.Series)):
-        if isinstance(frame := from_pandas(obj), Series):
-            frame = frame.to_frame()
-        return frame.lazy()
-    elif is_pycapsule(obj) or (
-        _check_for_pyarrow(obj) and isinstance(obj, (pa.Table, pa.RecordBatch))
-    ):
-        return from_arrow(obj).lazy()  # type: ignore[union-attr]
-    else:
-        msg = f"unrecognised frame type: {qualified_type_name(obj)}"
-        raise ValueError(msg)
-
-
-def _get_frame_locals(
-    *,
-    all_compatible: bool,
-    n_objects: int | None = None,
-    named: str | Collection[str] | Callable[[str], bool] | None = None,
-) -> dict[str, Any]:
-    """Return compatible frame objects from the local stack."""
-    of_type = _compatible_frame if all_compatible else (DataFrame, LazyFrame, Series)
-    return _get_stack_locals(of_type=of_type, n_objects=n_objects, named=named)  # type: ignore[arg-type]
-
-
-class SQLContext(Generic[FrameType]):
-    """
-    Run SQL queries against DataFrame, LazyFrame, and Series data.
-
-    .. warning::
-        This functionality is considered **unstable**, although it is close to being
-        considered stable. It may be changed at any point without it being considered
-        a breaking change.
-    """
-
-    _ctxt: PySQLContext
-    _eager_execution: Final[bool]
-    _tables_scope_stack: list[set[str]]
-
-    # note: the type-overloaded methods are required to support accurate typing
-    # of the frame return from "execute" (which may be DataFrame or LazyFrame),
-    # as that is influenced by both the "eager" flag at init-time AND the "eager"
-    # flag at query-time (if anyone can find a lighter-weight set of annotations
-    # that successfully resolves this, please go for it... ;)
-
-    @overload
-    def __init__(
-        self: SQLContext[LazyFrame],
-        frames: Mapping[str, CompatibleFrameType | None] | None = ...,
-        *,
-        register_globals: bool | int = ...,
-        all_compatible: bool = ...,
-        eager: Literal[False] = False,
-        **named_frames: CompatibleFrameType | None,
-    ) -> None: ...
-
-    @overload
-    def __init__(
-        self: SQLContext[DataFrame],
-        frames: Mapping[str, CompatibleFrameType | None] | None = ...,
-        *,
-        register_globals: bool | int = ...,
-        all_compatible: bool = ...,
-        eager: Literal[True],
-        **named_frames: CompatibleFrameType | None,
-    ) -> None: ...
-
-    @overload
-    def __init__(
-        self: SQLContext[DataFrame],
-        frames: Mapping[str, CompatibleFrameType | None] | None = ...,
-        *,
-        register_globals: bool | int = ...,
-        all_compatible: bool = ...,
-        eager: bool,
-        **named_frames: CompatibleFrameType | None,
-    ) -> None: ...
-
-    @deprecate_renamed_parameter("eager_execution", "eager", version="0.20.31")
-    def __init__(
-        self,
-        frames: Mapping[str, CompatibleFrameType | None] | None = None,
-        *,
-        register_globals: bool | int = False,
-        eager: bool = False,
-        **named_frames: CompatibleFrameType | None,
-    ) -> None:
-        """
-        Initialize a new `SQLContext`.
-
-        .. versionchanged:: 0.20.31
-            The `eager_execution` parameter was renamed `eager`.
-
-        Parameters
-        ----------
-        frames
-            A `{name:frame, ...}` mapping which can include Polars frames *and*
-            pandas DataFrames, Series and pyarrow Table and RecordBatch objects.
-        register_globals
-            Register compatible objects (polars DataFrame, LazyFrame, and Series) found
-            in the globals, automatically mapping their variable name to a table name.
-            To register other objects (pandas/pyarrow data) pass them explicitly, or
-            call the `execute_global` classmethod. If given an integer then only the
-            most recent "n" objects found will be registered.
-        eager
-            If True, returns execution results as `DataFrame` instead of `LazyFrame`.
-            (Note that the query itself is always executed in lazy-mode; this parameter
-            impacts whether :meth:`execute` returns an eager or lazy result frame).
-        **named_frames
-            Named eager/lazy frames, provided as kwargs.
-
-        Examples
-        --------
-        >>> lf = pl.LazyFrame({"a": [1, 2, 3], "b": ["x", None, "z"]})
-        >>> res = pl.SQLContext(frame=lf).execute(
-        ...     "SELECT b, a*2 AS two_a FROM frame WHERE b IS NOT NULL"
-        ... )
-        >>> res.collect()
-        shape: (2, 2)
-        ┌─────┬───────┐
-        │ b   ┆ two_a │
-        │ --- ┆ ---   │
-        │ str ┆ i64   │
-        ╞═════╪═══════╡
-        │ x   ┆ 2     │
-        │ z   ┆ 6     │
-        └─────┴───────┘
-        """
-        issue_unstable_warning(
-            "`SQLContext` is considered **unstable**, although it is close to being considered stable."
-        )
-        self._ctxt = PySQLContext.new()
-        self._eager_execution = eager
-
-        frames = dict(frames or {})
-        if register_globals:
-            for name, obj in _get_frame_locals(
-                all_compatible=False,
-            ).items():
-                if name not in frames and name not in named_frames:
-                    named_frames[name] = obj
-
-        if frames or named_frames:
-            frames.update(named_frames)
-            self.register_many(frames)
-
-    @overload
-    @classmethod
-    def execute_global(
-        cls, query: str, *, eager: Literal[False] = False
-    ) -> LazyFrame: ...
-
-    @overload
-    @classmethod
-    def execute_global(cls, query: str, *, eager: Literal[True]) -> DataFrame: ...
-
-    @overload
-    @classmethod
-    def execute_global(cls, query: str, *, eager: bool) -> DataFrame | LazyFrame: ...
-
-    @classmethod
-    def execute_global(
-        cls, query: str, *, eager: bool = False
-    ) -> DataFrame | LazyFrame:
-        """
-        Immediately execute a SQL query, automatically registering frame globals.
-
-        Notes
-        -----
-        * This convenience method automatically registers all compatible objects in
-          the local stack that are referenced in the query, mapping their variable name
-          to a table name. Note that in addition to polars DataFrame, LazyFrame, and
-          Series this method *also* registers pandas DataFrame, Series, and pyarrow
-          Table and RecordBatch objects.
-        * Instead of calling this classmethod you should consider using `pl.sql`,
-          which will use this code internally.
-
-        Parameters
-        ----------
-        query
-            A valid SQL query string.
-        eager
-            If True, returns execution results as `DataFrame` instead of `LazyFrame`.
-            (Note that the query itself is always executed in lazy-mode).
-
-        Examples
-        --------
-        >>> import pandas as pd
-        >>> df = pl.LazyFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
-        >>> df_pandas = pd.DataFrame({"a": [2, 3, 4], "c": [7, 8, 9]})
-
-        Join a polars LazyFrame with a pandas DataFrame (note use of the preferred
-        `pl.sql` method, which is equivalent to `SQLContext.execute_global`):
-
-        >>> pl.sql("SELECT df.*, c FROM df JOIN df_pandas USING(a)").collect()
-        shape: (2, 3)
-        ┌─────┬─────┬─────┐
-        │ a   ┆ b   ┆ c   │
-        │ --- ┆ --- ┆ --- │
-        │ i64 ┆ i64 ┆ i64 │
-        ╞═════╪═════╪═════╡
-        │ 2   ┆ 5   ┆ 7   │
-        │ 3   ┆ 6   ┆ 8   │
-        └─────┴─────┴─────┘
-        """
-        # basic extraction of possible table names from the query, so we don't register
-        # unnecessary objects from the globals (ideally we shuoold look to make the
-        # underlying `sqlparser-rs` lib parse the query to identify table names)
-        q = re.split(r"\bFROM\b", query, maxsplit=1, flags=re.I)
-        possible_names = (
-            {
-                nm.strip('"')
-                for nm in re.split(r"\b", q[1])
-                if re.match(r'^("[^"]+")$', nm) or nm.isidentifier()
-            }
-            if len(q) > 1
-            else set()
-        )
-        # get compatible frame objects from the globals, constraining by possible names
-        named_frames = _get_frame_locals(all_compatible=True, named=possible_names)
-        with cls(frames=named_frames, register_globals=False) as ctx:
-            return ctx.execute(query=query, eager=eager)
-
-    def __enter__(self) -> SQLContext[FrameType]:
-        """Track currently registered tables on scope entry; supports nested scopes."""
-        self._tables_scope_stack = getattr(self, "_tables_scope_stack", [])
-        self._tables_scope_stack.append(set(self.tables()))
-        return self
-
-    def __exit__(
-        self,
-        exc_type: type[BaseException] | None,
-        exc_val: BaseException | None,
-        exc_tb: TracebackType | None,
-    ) -> None:
-        """
-        Unregister any tables created within the given scope on context exit.
-
-        See Also
-        --------
-        unregister
-        """
-        self.unregister(
-            names=(set(self.tables()) - self._tables_scope_stack.pop()),
-        )
-
-    def __repr__(self) -> str:
-        n_tables = len(self.tables())
-        return f"<SQLContext [tables:{n_tables}] at 0x{id(self):x}>"
-
-    # these overloads are necessary to cover the possible permutations
-    # of the init-time "eager" param, and the local "eager" param.
-
-    @overload
-    def execute(
-        self: SQLContext[DataFrame], query: str, *, eager: None = ...
-    ) -> DataFrame: ...
-
-    @overload
-    def execute(
-        self: SQLContext[DataFrame], query: str, *, eager: Literal[False]
-    ) -> LazyFrame: ...
-
-    @overload
-    def execute(
-        self: SQLContext[DataFrame], query: str, *, eager: Literal[True]
-    ) -> DataFrame: ...
-
-    @overload
-    def execute(
-        self: SQLContext[LazyFrame], query: str, *, eager: None = ...
-    ) -> LazyFrame: ...
-
-    @overload
-    def execute(
-        self: SQLContext[LazyFrame], query: str, *, eager: Literal[False]
-    ) -> LazyFrame: ...
-
-    @overload
-    def execute(
-        self: SQLContext[LazyFrame], query: str, *, eager: Literal[True]
-    ) -> DataFrame: ...
-
-    @overload
-    def execute(
-        self, query: str, *, eager: bool | None = ...
-    ) -> LazyFrame | DataFrame: ...
-
-    def execute(
-        self, query: str, *, eager: bool | None = None
-    ) -> LazyFrame | DataFrame:
-        """
-        Parse the given SQL query and execute it against the registered frame data.
-
-        Parameters
-        ----------
-        query
-            A valid string SQL query.
-        eager
-            Apply the query eagerly, returning `DataFrame` instead of `LazyFrame`.
-            If unset, the value of the init-time "eager" parameter will be used.
-            Note that the query itself is always executed in lazy-mode; this
-            parameter only impacts the type of the returned frame.
-
-        Examples
-        --------
-        Declare frame data and register with a SQLContext:
-
-        >>> df = pl.DataFrame(
-        ...     data=[
-        ...         ("The Godfather", 1972, 6_000_000, 134_821_952, 9.2),
-        ...         ("The Dark Knight", 2008, 185_000_000, 533_316_061, 9.0),
-        ...         ("Schindler's List", 1993, 22_000_000, 96_067_179, 8.9),
-        ...         ("Pulp Fiction", 1994, 8_000_000, 107_930_000, 8.9),
-        ...         ("The Shawshank Redemption", 1994, 25_000_000, 28_341_469, 9.3),
-        ...     ],
-        ...     schema=["title", "release_year", "budget", "gross", "imdb_score"],
-        ...     orient="row",
-        ... )
-        >>> ctx = pl.SQLContext(films=df)
-
-        Execute a SQL query against the registered frame data:
-
-        >>> ctx.execute(
-        ...     '''
-        ...     SELECT title, release_year, imdb_score
-        ...     FROM films
-        ...     WHERE release_year > 1990
-        ...     ORDER BY imdb_score DESC
-        ...     ''',
-        ...     eager=True,
-        ... )
-        shape: (4, 3)
-        ┌──────────────────────────┬──────────────┬────────────┐
-        │ title                    ┆ release_year ┆ imdb_score │
-        │ ---                      ┆ ---          ┆ ---        │
-        │ str                      ┆ i64          ┆ f64        │
-        ╞══════════════════════════╪══════════════╪════════════╡
-        │ The Shawshank Redemption ┆ 1994         ┆ 9.3        │
-        │ The Dark Knight          ┆ 2008         ┆ 9.0        │
-        │ Schindler's List         ┆ 1993         ┆ 8.9        │
-        │ Pulp Fiction             ┆ 1994         ┆ 8.9        │
-        └──────────────────────────┴──────────────┴────────────┘
-
-        Execute a GROUP BY query:
-
-        >>> ctx.execute(
-        ...     '''
-        ...     SELECT
-        ...         MAX(release_year / 10) * 10 AS decade,
-        ...         SUM(gross) AS total_gross,
-        ...         COUNT(title) AS n_films,
-        ...     FROM films
-        ...     GROUP BY (release_year / 10) -- decade
-        ...     ORDER BY total_gross DESC
-        ...     ''',
-        ...     eager=True,
-        ... )
-        shape: (3, 3)
-        ┌────────┬─────────────┬─────────┐
-        │ decade ┆ total_gross ┆ n_films │
-        │ ---    ┆ ---         ┆ ---     │
-        │ i64    ┆ i64         ┆ u32     │
-        ╞════════╪═════════════╪═════════╡
-        │ 2000   ┆ 533316061   ┆ 1       │
-        │ 1990   ┆ 232338648   ┆ 3       │
-        │ 1970   ┆ 134821952   ┆ 1       │
-        └────────┴─────────────┴─────────┘
-        """
-        res = wrap_ldf(self._ctxt.execute(query))
-        return res.collect() if (eager or self._eager_execution) else res
-
-    def register(self, name: str, frame: CompatibleFrameType | None) -> Self:
-        """
-        Register a single frame as a table, using the given name.
-
-        Parameters
-        ----------
-        name
-            Name of the table.
-        frame
-            eager/lazy frame to associate with this table name.
-
-        See Also
-        --------
-        register_globals
-        register_many
-        unregister
-
-        Examples
-        --------
-        >>> df = pl.DataFrame({"hello": ["world"]})
-        >>> ctx = pl.SQLContext()
-        >>> ctx.register("frame_data", df).execute("SELECT * FROM frame_data").collect()
-        shape: (1, 1)
-        ┌───────┐
-        │ hello │
-        │ ---   │
-        │ str   │
-        ╞═══════╡
-        │ world │
-        └───────┘
-        """
-        frame = LazyFrame() if frame is None else _ensure_lazyframe(frame)
-        self._ctxt.register(name, frame._ldf)
-        return self
-
-    def register_globals(
-        self, n: int | None = None, *, all_compatible: bool = True
-    ) -> Self:
-        """
-        Register all frames (lazy or eager) found in the current globals scope.
-
-        Automatically maps variable names to table names.
-
-        See Also
-        --------
-        register
-        register_many
-        unregister
-
-        Parameters
-        ----------
-        n
-            Register only the most recent "n" frames.
-        all_compatible
-            Control whether we *also* register pandas DataFrame, Series, and
-            pyarrow Table and RecordBatch objects. If False, only Polars
-            classes are registered with the SQL engine.
-
-        Examples
-        --------
-        >>> df1 = pl.DataFrame({"a": [1, 2, 3], "b": ["x", None, "z"]})
-        >>> df2 = pl.DataFrame({"a": [2, 3, 4], "c": ["t", "w", "v"]})
-
-        Register frames directly from variables found in the current globals scope:
-
-        >>> ctx = pl.SQLContext(register_globals=True)
-        >>> ctx.tables()
-        ['df1', 'df2']
-
-        Query using the register variable/frame names
-
-        >>> ctx.execute(
-        ...     "SELECT a, b, c FROM df1 LEFT JOIN df2 USING (a) ORDER BY a DESC"
-        ... ).collect()
-        shape: (3, 3)
-        ┌─────┬──────┬──────┐
-        │ a   ┆ b    ┆ c    │
-        │ --- ┆ ---  ┆ ---  │
-        │ i64 ┆ str  ┆ str  │
-        ╞═════╪══════╪══════╡
-        │ 3   ┆ z    ┆ w    │
-        │ 2   ┆ null ┆ t    │
-        │ 1   ┆ x    ┆ null │
-        └─────┴──────┴──────┘
-        """
-        frames = _get_frame_locals(all_compatible=all_compatible, n_objects=n)
-        return self.register_many(frames=frames)
-
-    def register_many(
-        self,
-        frames: Mapping[str, CompatibleFrameType | None] | None = None,
-        **named_frames: CompatibleFrameType | None,
-    ) -> Self:
-        """
-        Register multiple eager/lazy frames as tables, using the associated names.
-
-        Parameters
-        ----------
-        frames
-            A `{name:frame, ...}` mapping.
-        **named_frames
-            Named eager/lazy frames, provided as kwargs.
-
-        See Also
-        --------
-        register
-        register_globals
-        unregister
-
-        Examples
-        --------
-        >>> lf1 = pl.LazyFrame({"a": [1, 2, 3], "b": ["m", "n", "o"]})
-        >>> lf2 = pl.LazyFrame({"a": [2, 3, 4], "c": ["p", "q", "r"]})
-        >>> lf3 = pl.LazyFrame({"a": [3, 4, 5], "b": ["s", "t", "u"]})
-        >>> lf4 = pl.LazyFrame({"a": [4, 5, 6], "c": ["v", "w", "x"]})
-
-        Register multiple frames at once, either by passing in as a dict...
-
-        >>> ctx = pl.SQLContext().register_many({"tbl1": lf1, "tbl2": lf2})
-        >>> ctx.tables()
-        ['tbl1', 'tbl2']
-
-        ...or using keyword args:
-
-        >>> ctx.register_many(tbl3=lf3, tbl4=lf4).tables()
-        ['tbl1', 'tbl2', 'tbl3', 'tbl4']
-        """
-        frames = dict(frames or {})
-        frames.update(named_frames)
-        for name, frame in frames.items():
-            self.register(name, frame)
-        return self
-
-    def unregister(self, names: str | Collection[str]) -> Self:
-        """
-        Unregister one or more eager/lazy frames by name.
-
-        Parameters
-        ----------
-        names
-            Names of the tables to unregister.
-
-        Notes
-        -----
-        You can also control table registration lifetime by using `SQLContext` as a
-        context manager; this can often be more useful when such control is wanted:
-
-        >>> df0 = pl.DataFrame({"colx": [0, 1, 2]})
-        >>> df1 = pl.DataFrame({"colx": [1, 2, 3]})
-        >>> df2 = pl.DataFrame({"colx": [2, 3, 4]})
-
-        Frames registered in-scope are automatically unregistered on scope-exit. Note
-        that frames registered on construction will persist through subsequent scopes.
-
-        >>> # register one frame at construction time, and the other two in-scope
-        >>> with pl.SQLContext(tbl0=df0) as ctx:
-        ...     ctx.register_many(tbl1=df1, tbl2=df2).tables()
-        ['tbl0', 'tbl1', 'tbl2']
-
-        After scope exit, none of the tables registered in-scope remain:
-
-        >>> ctx.tables()
-        ['tbl0']
-
-        See Also
-        --------
-        register
-        register_globals
-        register_many
-
-        Examples
-        --------
-        >>> df0 = pl.DataFrame({"ints": [9, 8, 7, 6, 5]})
-        >>> lf1 = pl.LazyFrame({"text": ["a", "b", "c"]})
-        >>> lf2 = pl.LazyFrame({"misc": ["testing1234"]})
-
-        Register with a SQLContext object:
-
-        >>> ctx = pl.SQLContext(test1=df0, test2=lf1, test3=lf2)
-        >>> ctx.tables()
-        ['test1', 'test2', 'test3']
-
-        Unregister one or more of the tables:
-
-        >>> ctx.unregister(["test1", "test3"]).tables()
-        ['test2']
-        >>> ctx.unregister("test2").tables()
-        []
-        """
-        if isinstance(names, str):
-            names = [names]
-        for nm in names:
-            self._ctxt.unregister(nm)
-        return self
-
-    def tables(self) -> list[str]:
-        """
-        Return a list of the registered table names.
-
-        Notes
-        -----
-        The :meth:`tables` method will return the same values as the
-        "SHOW TABLES" SQL statement, but as a list instead of a frame.
-
-        Executing as SQL:
-
-        >>> frame_data = pl.DataFrame({"hello": ["world"]})
-        >>> ctx = pl.SQLContext(hello_world=frame_data)
-        >>> ctx.execute("SHOW TABLES", eager=True)
-        shape: (1, 1)
-        ┌─────────────┐
-        │ name        │
-        │ ---         │
-        │ str         │
-        ╞═════════════╡
-        │ hello_world │
-        └─────────────┘
-
-        Calling the method:
-
-        >>> ctx.tables()
-        ['hello_world']
-
-        Examples
-        --------
-        >>> df1 = pl.DataFrame({"hello": ["world"]})
-        >>> df2 = pl.DataFrame({"foo": ["bar", "baz"]})
-        >>> ctx = pl.SQLContext(hello_data=df1, foo_bar=df2)
-        >>> ctx.tables()
-        ['foo_bar', 'hello_data']
-        """
-        return sorted(self._ctxt.get_tables())