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

polars/sql/functions.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal, overload
+
+if TYPE_CHECKING:
+    from polars.dataframe import DataFrame
+    from polars.lazyframe import LazyFrame
+
+
+__all__ = ["sql"]
+
+
+@overload
+def sql(query: str, *, eager: Literal[False] = False) -> LazyFrame: ...
+
+
+@overload
+def sql(query: str, *, eager: Literal[True]) -> DataFrame: ...
+
+
+def sql(query: str, *, eager: bool = False) -> DataFrame | LazyFrame:
+    """
+    Execute a SQL query against frames in the global namespace.
+
+    .. versionadded:: 0.20.31
+
+    Parameters
+    ----------
+    query
+        SQL query to execute.
+    eager
+        Automatically collect the result and return a DataFrame instead of a LazyFrame.
+
+    Notes
+    -----
+    * The Polars SQL engine can operate against Polars DataFrame, LazyFrame, and Series
+      objects, as well as Pandas DataFrame and Series, PyArrow Table and RecordBatch.
+    * Additional control over registration and execution behaviour is available
+      with the :class:`SQLContext` object.
+
+    See Also
+    --------
+    SQLContext
+
+    Examples
+    --------
+    >>> lf1 = pl.LazyFrame({"a": [1, 2, 3], "b": [6, 7, 8], "c": ["z", "y", "x"]})
+    >>> lf2 = pl.LazyFrame({"a": [3, 2, 1], "d": [125, -654, 888]})
+
+    Query the LazyFrame using SQL:
+
+    >>> lf1.sql("SELECT c, b FROM self WHERE a > 1").collect()
+    shape: (2, 2)
+    ┌─────┬─────┐
+    │ c   ┆ b   │
+    │ --- ┆ --- │
+    │ str ┆ i64 │
+    ╞═════╪═════╡
+    │ y   ┆ 7   │
+    │ x   ┆ 8   │
+    └─────┴─────┘
+
+    Join two LazyFrames:
+
+    >>> pl.sql(
+    ...     '''
+    ...     SELECT lf1.*, d
+    ...     FROM lf1
+    ...     INNER JOIN lf2 USING (a)
+    ...     WHERE a > 1 AND b < 8
+    ...     '''
+    ... ).collect()
+    shape: (1, 4)
+    ┌─────┬─────┬─────┬──────┐
+    │ a   ┆ b   ┆ c   ┆ d    │
+    │ --- ┆ --- ┆ --- ┆ ---  │
+    │ i64 ┆ i64 ┆ str ┆ i64  │
+    ╞═════╪═════╪═════╪══════╡
+    │ 2   ┆ 7   ┆ y   ┆ -654 │
+    └─────┴─────┴─────┴──────┘
+
+    Apply SQL transforms and subsequently filter natively (you can freely mix SQL and
+    native operations):
+
+    >>> pl.sql(
+    ...     query='''
+    ...         SELECT
+    ...             a,
+    ...             (a % 2 == 0) AS a_is_even,
+    ...             (b::float4 / 2) AS "b/2",
+    ...             CONCAT_WS(':', c, c, c) AS c_c_c
+    ...         FROM lf1
+    ...         ORDER BY a
+    ...     ''',
+    ... ).filter(~pl.col("c_c_c").str.starts_with("x")).collect()
+    shape: (2, 4)
+    ┌─────┬───────────┬─────┬───────┐
+    │ a   ┆ a_is_even ┆ b/2 ┆ c_c_c │
+    │ --- ┆ ---       ┆ --- ┆ ---   │
+    │ i64 ┆ bool      ┆ f32 ┆ str   │
+    ╞═════╪═══════════╪═════╪═══════╡
+    │ 1   ┆ false     ┆ 3.0 ┆ z:z:z │
+    │ 2   ┆ true      ┆ 3.5 ┆ y:y:y │
+    └─────┴───────────┴─────┴───────┘
+
+    Join polars LazyFrame with a pandas DataFrame and a pyarrow Table:
+
+    >>> import pandas as pd
+    >>> import pyarrow as pa
+    >>> pl_frame = lf1
+    >>> pd_frame = pd.DataFrame({"a": [2, 3, 4], "d": [-0.5, 0.0, 0.5]})
+    >>> pa_table = pa.Table.from_arrays(
+    ...     [pa.array([1, 2, 3]), pa.array(["x", "y", "z"])],
+    ...     names=["a", "e"],
+    ... )
+    >>> pl.sql(
+    ...     query='''
+    ...         SELECT pl_frame.*, d, e
+    ...         FROM pl_frame
+    ...         JOIN pd_frame USING(a)
+    ...         JOIN pa_table USING(a)
+    ...     ''',
+    ... ).collect()
+    shape: (2, 5)
+    ┌─────┬─────┬─────┬──────┬─────┐
+    │ a   ┆ b   ┆ c   ┆ d    ┆ e   │
+    │ --- ┆ --- ┆ --- ┆ ---  ┆ --- │
+    │ i64 ┆ i64 ┆ str ┆ f64  ┆ str │
+    ╞═════╪═════╪═════╪══════╪═════╡
+    │ 2   ┆ 7   ┆ y   ┆ -0.5 ┆ y   │
+    │ 3   ┆ 8   ┆ x   ┆ 0.0  ┆ z   │
+    └─────┴─────┴─────┴──────┴─────┘
+    """
+    from polars.sql import SQLContext
+
+    return SQLContext.execute_global(
+        query=query,
+        eager=eager,
+    )

polars/string_cache.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+import contextlib
+from typing import TYPE_CHECKING
+
+with contextlib.suppress(ImportError):  # Module not available when building docs
+    import polars._plr as plr
+    from polars._plr import PyStringCacheHolder
+
+if TYPE_CHECKING:
+    import sys
+    from types import TracebackType
+
+    if sys.version_info >= (3, 11):
+        from typing import Self
+    else:
+        from typing_extensions import Self
+
+
+__all__ = [
+    "StringCache",
+    "disable_string_cache",
+    "enable_string_cache",
+    "using_string_cache",
+]
+
+
+class StringCache(contextlib.ContextDecorator):
+    """
+    Context manager for enabling and disabling the global string cache.
+
+    :class:`Categorical` columns created under the same global string cache have
+    the same underlying physical value when string values are equal. This allows the
+    columns to be concatenated or used in a join operation, for example.
+
+    Notes
+    -----
+    Enabling the global string cache introduces some overhead.
+    The amount of overhead depends on the number of categories in your data.
+    It is advised to enable the global string cache only when strictly necessary.
+
+    If `StringCache` calls are nested, the global string cache will only be disabled
+    and cleared when the outermost context exits.
+
+    Examples
+    --------
+    Construct two Series using the same global string cache.
+
+    >>> with pl.StringCache():
+    ...     s1 = pl.Series("color", ["red", "green", "red"], dtype=pl.Categorical)
+    ...     s2 = pl.Series("color", ["blue", "red", "green"], dtype=pl.Categorical)
+
+    As both Series are constructed under the same global string cache,
+    they can be concatenated.
+
+    >>> pl.concat([s1, s2])
+    shape: (6,)
+    Series: 'color' [cat]
+    [
+            "red"
+            "green"
+            "red"
+            "blue"
+            "red"
+            "green"
+    ]
+
+    The class can also be used as a function decorator, in which case the string cache
+    is enabled during function execution, and disabled afterwards.
+
+    >>> @pl.StringCache()
+    ... def construct_categoricals() -> pl.Series:
+    ...     s1 = pl.Series("color", ["red", "green", "red"], dtype=pl.Categorical)
+    ...     s2 = pl.Series("color", ["blue", "red", "green"], dtype=pl.Categorical)
+    ...     return pl.concat([s1, s2])
+    """
+
+    def __enter__(self) -> Self:
+        self._string_cache = PyStringCacheHolder()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        del self._string_cache
+
+
+def enable_string_cache() -> None:
+    """
+    Enable the global string cache.
+
+    :class:`Categorical` columns created under the same global string cache have
+    the same underlying physical value when string values are equal. This allows the
+    columns to be concatenated or used in a join operation, for example.
+
+    See Also
+    --------
+    StringCache : Context manager for enabling and disabling the string cache.
+    disable_string_cache : Function to disable the string cache.
+
+    Notes
+    -----
+    Enabling the global string cache introduces some overhead.
+    The amount of overhead depends on the number of categories in your data.
+    It is advised to enable the global string cache only when strictly necessary.
+
+    Consider using the :class:`StringCache` context manager for a more reliable way of
+    enabling and disabling the string cache.
+
+    Examples
+    --------
+    Construct two Series using the same global string cache.
+
+    >>> pl.enable_string_cache()
+    >>> s1 = pl.Series("color", ["red", "green", "red"], dtype=pl.Categorical)
+    >>> s2 = pl.Series("color", ["blue", "red", "green"], dtype=pl.Categorical)
+    >>> pl.disable_string_cache()
+
+    As both Series are constructed under the same global string cache,
+    they can be concatenated.
+
+    >>> pl.concat([s1, s2])
+    shape: (6,)
+    Series: 'color' [cat]
+    [
+            "red"
+            "green"
+            "red"
+            "blue"
+            "red"
+            "green"
+    ]
+    """
+    plr.enable_string_cache()
+
+
+def disable_string_cache() -> None:
+    """
+    Disable and clear the global string cache.
+
+    See Also
+    --------
+    enable_string_cache : Function to enable the string cache.
+    StringCache : Context manager for enabling and disabling the string cache.
+
+    Notes
+    -----
+    Consider using the :class:`StringCache` context manager for a more reliable way of
+    enabling and disabling the string cache.
+
+    When used in conjunction with the :class:`StringCache` context manager, the string
+    cache will not be disabled until the context manager exits.
+
+    Examples
+    --------
+    Construct two Series using the same global string cache.
+
+    >>> pl.enable_string_cache()
+    >>> s1 = pl.Series("color", ["red", "green", "red"], dtype=pl.Categorical)
+    >>> s2 = pl.Series("color", ["blue", "red", "green"], dtype=pl.Categorical)
+    >>> pl.disable_string_cache()
+
+    As both Series are constructed under the same global string cache,
+    they can be concatenated.
+
+    >>> pl.concat([s1, s2])
+    shape: (6,)
+    Series: 'color' [cat]
+    [
+            "red"
+            "green"
+            "red"
+            "blue"
+            "red"
+            "green"
+    ]
+    """
+
+
+def using_string_cache() -> bool:
+    """Check whether the global string cache is enabled."""
+    return plr.using_string_cache()

polars/testing/__init__.py

@@ -0,0 +1,13 @@
+from polars.testing.asserts import (
+    assert_frame_equal,
+    assert_frame_not_equal,
+    assert_series_equal,
+    assert_series_not_equal,
+)
+
+__all__ = [
+    "assert_frame_equal",
+    "assert_frame_not_equal",
+    "assert_series_equal",
+    "assert_series_not_equal",
+]

polars/testing/asserts/__init__.py

@@ -0,0 +1,9 @@
+from polars.testing.asserts.frame import assert_frame_equal, assert_frame_not_equal
+from polars.testing.asserts.series import assert_series_equal, assert_series_not_equal
+
+__all__ = [
+    "assert_frame_equal",
+    "assert_frame_not_equal",
+    "assert_series_equal",
+    "assert_series_not_equal",
+]

polars/testing/asserts/frame.py

@@ -0,0 +1,231 @@
+from __future__ import annotations
+
+import contextlib
+from typing import cast
+
+from polars._utils.deprecation import deprecate_renamed_parameter
+from polars.dataframe import DataFrame
+from polars.lazyframe import LazyFrame
+from polars.testing.asserts.utils import raise_assertion_error
+
+with contextlib.suppress(ImportError):  # Module not available when building docs
+    from polars._plr import assert_dataframe_equal_py
+
+
+def _assert_correct_input_type(
+    left: DataFrame | LazyFrame, right: DataFrame | LazyFrame
+) -> bool:
+    __tracebackhide__ = True
+
+    if isinstance(left, DataFrame) and isinstance(right, DataFrame):
+        return False
+    elif isinstance(left, LazyFrame) and isinstance(right, LazyFrame):
+        return True
+    else:
+        raise_assertion_error(
+            "inputs",
+            "unexpected input types",
+            type(left).__name__,
+            type(right).__name__,
+        )
+
+
+@deprecate_renamed_parameter("check_dtype", "check_dtypes", version="0.20.31")
+@deprecate_renamed_parameter("rtol", "rel_tol", version="1.32.3")
+@deprecate_renamed_parameter("atol", "abs_tol", version="1.32.3")
+def assert_frame_equal(
+    left: DataFrame | LazyFrame,
+    right: DataFrame | LazyFrame,
+    *,
+    check_row_order: bool = True,
+    check_column_order: bool = True,
+    check_dtypes: bool = True,
+    check_exact: bool = False,
+    rel_tol: float = 1e-5,
+    abs_tol: float = 1e-8,
+    categorical_as_str: bool = False,
+) -> None:
+    """
+    Assert that the left and right frame are equal.
+
+    Raises a detailed `AssertionError` if the frames differ.
+    This function is intended for use in unit tests.
+
+    .. versionchanged:: 0.20.31
+        The `check_dtype` parameter was renamed `check_dtypes`.
+
+    .. versionchanged:: 1.32.3
+        The `rtol` and `atol` parameters were renamed to `rel_tol` and `abs_tol`,
+        respectively.
+
+    Parameters
+    ----------
+    left
+        The first DataFrame or LazyFrame to compare.
+    right
+        The second DataFrame or LazyFrame to compare.
+    check_row_order
+        Requires row order to match.
+    check_column_order
+        Requires column order to match.
+    check_dtypes
+        Requires data types to match.
+    check_exact
+        Requires float values to match exactly. If set to `False`, values are considered
+        equal when within tolerance of each other (see `rel_tol` and `abs_tol`).
+        Only affects columns with a Float data type.
+    rel_tol
+        Relative tolerance for inexact checking. Fraction of values in `right`.
+    abs_tol
+        Absolute tolerance for inexact checking.
+    categorical_as_str
+        Cast categorical columns to string before comparing. Enabling this helps
+        compare columns that do not share the same string cache.
+
+    See Also
+    --------
+    assert_series_equal
+    assert_frame_not_equal
+
+    Notes
+    -----
+    When using pytest, it may be worthwhile to shorten Python traceback printing
+    by passing `--tb=short`. The default mode tends to be unhelpfully verbose.
+    More information in the
+    `pytest docs <https://docs.pytest.org/en/latest/how-to/output.html#modifying-python-traceback-printing>`_.
+
+    Examples
+    --------
+    >>> from polars.testing import assert_frame_equal
+    >>> df1 = pl.DataFrame({"a": [1, 2, 3]})
+    >>> df2 = pl.DataFrame({"a": [1, 5, 3]})
+    >>> assert_frame_equal(df1, df2)
+    Traceback (most recent call last):
+    ...
+    AssertionError: DataFrames are different (value mismatch for column "a")
+    [left]: shape: (3,)
+    Series: 'a' [i64]
+    [
+        1
+        2
+        3
+    ]
+    [right]: shape: (3,)
+    Series: 'a' [i64]
+    [
+        1
+        5
+        3
+    ]
+    """
+    __tracebackhide__ = True
+
+    lazy = _assert_correct_input_type(left, right)
+
+    # Rust back-end function expects DataFrames so LazyFrames must be collected
+    if lazy:
+        left, right = left.collect(), right.collect()  # type: ignore[union-attr]
+
+    # Tell type checker these are now DataFrames to prevent type errors
+    left, right = cast("DataFrame", left), cast("DataFrame", right)
+
+    assert_dataframe_equal_py(
+        left._df,
+        right._df,
+        check_row_order=check_row_order,
+        check_column_order=check_column_order,
+        check_dtypes=check_dtypes,
+        check_exact=check_exact,
+        rel_tol=rel_tol,
+        abs_tol=abs_tol,
+        categorical_as_str=categorical_as_str,
+    )
+
+
+@deprecate_renamed_parameter("check_dtype", "check_dtypes", version="0.20.31")
+@deprecate_renamed_parameter("rtol", "rel_tol", version="1.32.3")
+@deprecate_renamed_parameter("atol", "abs_tol", version="1.32.3")
+def assert_frame_not_equal(
+    left: DataFrame | LazyFrame,
+    right: DataFrame | LazyFrame,
+    *,
+    check_row_order: bool = True,
+    check_column_order: bool = True,
+    check_dtypes: bool = True,
+    check_exact: bool = False,
+    rel_tol: float = 1e-5,
+    abs_tol: float = 1e-8,
+    categorical_as_str: bool = False,
+) -> None:
+    """
+    Assert that the left and right frame are **not** equal.
+
+    This function is intended for use in unit tests.
+
+    .. versionchanged:: 0.20.31
+        The `check_dtype` parameter was renamed `check_dtypes`.
+
+    .. versionchanged:: 1.32.3
+        The `rtol` and `atol` parameters were renamed to `rel_tol` and `abs_tol`,
+        respectively.
+
+    Parameters
+    ----------
+    left
+        The first DataFrame or LazyFrame to compare.
+    right
+        The second DataFrame or LazyFrame to compare.
+    check_row_order
+        Requires row order to match.
+    check_column_order
+        Requires column order to match.
+    check_dtypes
+        Requires data types to match.
+    check_exact
+        Requires float values to match exactly. If set to `False`, values are considered
+        equal when within tolerance of each other (see `rel_tol` and `abs_tol`).
+        Only affects columns with a Float data type.
+    rel_tol
+        Relative tolerance for inexact checking. Fraction of values in `right`.
+    abs_tol
+        Absolute tolerance for inexact checking.
+    categorical_as_str
+        Cast categorical columns to string before comparing. Enabling this helps
+        compare columns that do not share the same string cache.
+
+    See Also
+    --------
+    assert_frame_equal
+    assert_series_not_equal
+
+    Examples
+    --------
+    >>> from polars.testing import assert_frame_not_equal
+    >>> df1 = pl.DataFrame({"a": [1, 2, 3]})
+    >>> df2 = pl.DataFrame({"a": [1, 2, 3]})
+    >>> assert_frame_not_equal(df1, df2)
+    Traceback (most recent call last):
+    ...
+    AssertionError: DataFrames are equal (but are expected not to be)
+    """
+    __tracebackhide__ = True
+
+    lazy = _assert_correct_input_type(left, right)
+    try:
+        assert_frame_equal(
+            left=left,
+            right=right,
+            check_column_order=check_column_order,
+            check_row_order=check_row_order,
+            check_dtypes=check_dtypes,
+            check_exact=check_exact,
+            rel_tol=rel_tol,
+            abs_tol=abs_tol,
+            categorical_as_str=categorical_as_str,
+        )
+    except AssertionError:
+        return
+    else:
+        objects = "LazyFrames" if lazy else "DataFrames"
+        msg = f"{objects} are equal (but are expected not to be)"
+        raise AssertionError(msg)