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

polars/io/partition.py

@@ -1,491 +0,0 @@
-from __future__ import annotations
-
-import contextlib
-from collections.abc import Iterable, Mapping, Sequence
-from pathlib import Path
-from typing import TYPE_CHECKING
-
-from polars import DataFrame, col
-from polars._typing import PartitioningScheme
-from polars._utils.unstable import issue_unstable_warning
-from polars.expr import Expr
-
-if TYPE_CHECKING:
-    with contextlib.suppress(ImportError):  # Module not available when building docs
-        from polars._plr import PyDataFrame, PyExpr
-
-    from typing import IO, Any, Callable
-
-with contextlib.suppress(ImportError):  # Module not available when building docs
-    from polars._plr import PyPartitioning
-
-
-class KeyedPartition:
-    """
-    A key-value pair for a partition.
-
-    .. warning::
-        This functionality is currently considered **unstable**. It may be
-        changed at any point without it being considered a breaking change.
-
-    See Also
-    --------
-    PartitionByKey
-    PartitionParted
-    KeyedPartitionContext
-    """
-
-    def __init__(self, name: str, str_value: str, raw_value: Any) -> None:
-        self.name = name
-        self.str_value = str_value
-        self.raw_value = raw_value
-
-    name: str  #: Name of the key column.
-    str_value: str  #: Value of the key as a path and URL safe string.
-    raw_value: Any  #: Value of the key for this partition.
-
-    def hive_name(self) -> str:
-        """Get the `key=value`."""
-        return f"{self.name}={self.str_value}"
-
-
-class KeyedPartitionContext:
-    """
-    Callback context for a partition creation using keys.
-
-    .. warning::
-        This functionality is currently considered **unstable**. It may be
-        changed at any point without it being considered a breaking change.
-
-    See Also
-    --------
-    PartitionByKey
-    PartitionParted
-    """
-
-    def __init__(
-        self,
-        file_idx: int,
-        part_idx: int,
-        in_part_idx: int,
-        keys: list[KeyedPartition],
-        file_path: Path,
-        full_path: Path,
-    ) -> None:
-        self.file_idx = file_idx
-        self.part_idx = part_idx
-        self.in_part_idx = in_part_idx
-        self.keys = keys
-        self.file_path = file_path
-        self.full_path = full_path
-
-    file_idx: int  #: The index of the created file starting from zero.
-    part_idx: int  #: The index of the created partition starting from zero.
-    in_part_idx: int  #: The index of the file within this partition starting from zero.
-    keys: list[KeyedPartition]  #: All the key names and values used for this partition.
-    file_path: Path  #: The chosen output path before the callback was called without `base_path`.
-    full_path: (
-        Path  #: The chosen output path before the callback was called with `base_path`.
-    )
-
-    def hive_dirs(self) -> Path:
-        """The keys mapped to hive directories."""
-        assert len(self.keys) > 0
-        p = Path(self.keys[0].hive_name())
-        for key in self.keys[1:]:
-            p /= Path(key.hive_name())
-        return p
-
-
-class BasePartitionContext:
-    """
-    Callback context for a partition creation.
-
-    .. warning::
-        This functionality is currently considered **unstable**. It may be
-        changed at any point without it being considered a breaking change.
-
-    See Also
-    --------
-    PartitionMaxSize
-    """
-
-    def __init__(self, file_idx: int, file_path: Path, full_path: Path) -> None:
-        self.file_idx = file_idx
-        self.file_path = file_path
-        self.full_path = full_path
-
-    file_idx: int  #: The index of the created file starting from zero.
-    file_path: Path  #: The chosen output path before the callback was called without `base_path`.
-    full_path: (
-        Path  #: The chosen output path before the callback was called with `base_path`.
-    )
-
-
-def _cast_base_file_path_cb(
-    file_path_cb: Callable[[BasePartitionContext], Path | str | IO[bytes] | IO[str]]
-    | None,
-) -> Callable[[BasePartitionContext], Path | str | IO[bytes] | IO[str]] | None:
-    if file_path_cb is None:
-        return None
-    return lambda ctx: file_path_cb(
-        BasePartitionContext(
-            file_idx=ctx.file_idx,
-            file_path=Path(ctx.file_path),
-            full_path=Path(ctx.full_path),
-        )
-    )
-
-
-def _cast_keyed_file_path_cb(
-    file_path_cb: Callable[[KeyedPartitionContext], Path | str | IO[bytes] | IO[str]]
-    | None,
-) -> Callable[[KeyedPartitionContext], Path | str | IO[bytes] | IO[str]] | None:
-    if file_path_cb is None:
-        return None
-    return lambda ctx: file_path_cb(
-        KeyedPartitionContext(
-            file_idx=ctx.file_idx,
-            part_idx=ctx.part_idx,
-            in_part_idx=ctx.in_part_idx,
-            keys=[
-                KeyedPartition(
-                    name=kv.name, str_value=kv.str_value, raw_value=kv.raw_value
-                )
-                for kv in ctx.keys
-            ],
-            file_path=Path(ctx.file_path),
-            full_path=Path(ctx.full_path),
-        )
-    )
-
-
-def _prepare_per_partition_sort_by(
-    e: str | Expr | Iterable[str | Expr] | None,
-) -> list[PyExpr] | None:
-    def prepare_one(v: str | Expr) -> PyExpr:
-        if isinstance(v, str):
-            return col(v)._pyexpr
-        elif isinstance(v, Expr):
-            return v._pyexpr
-        else:
-            msg = f"cannot do a per partition sort by for {v!r}"
-            raise TypeError(msg)
-
-    if e is None:
-        return None
-    elif isinstance(e, str):
-        return [col(e)._pyexpr]
-    elif isinstance(e, Expr):
-        return [e._pyexpr]
-    elif isinstance(e, Iterable):
-        return [prepare_one(v) for v in e]
-    else:
-        msg = f"cannot do a per partition sort by for {e!r}"
-        raise TypeError(msg)
-
-
-def _prepare_finish_callback(
-    f: Callable[[DataFrame], None] | None,
-) -> Callable[[PyDataFrame], None] | None:
-    if f is None:
-        return None
-
-    def cb(pydf: PyDataFrame) -> None:
-        nonlocal f
-        f(DataFrame._from_pydf(pydf))
-
-    return cb
-
-
-class PartitionMaxSize(PartitioningScheme):
-    """
-    Partitioning scheme to write files with a maximum size.
-
-    This partitioning scheme generates files that have a given maximum size. If
-    the size reaches the maximum size, it is closed and a new file is opened.
-
-    .. warning::
-        This functionality is currently considered **unstable**. It may be
-        changed at any point without it being considered a breaking change.
-
-    Parameters
-    ----------
-    base_path
-        The base path for the output files.
-    file_path
-        A callback to register or modify the output path for each partition
-        relative to the `base_path`. The callback provides a
-        :class:`polars.io.partition.BasePartitionContext` that contains information
-        about the partition.
-
-        If no callback is given, it defaults to `{ctx.file_idx}.{EXT}`.
-    max_size : int
-        The maximum size in rows of each of the generated files.
-    per_partition_sort_by
-        Columns or expressions to sort over within each partition.
-
-        Note that this might increase the memory consumption needed for each partition.
-    finish_callback
-        A callback that gets called when the query finishes successfully.
-
-        For parquet files, the callback is given a dataframe with metrics about all
-        files written files.
-
-    Examples
-    --------
-    Split a parquet file by over smaller CSV files with 100 000 rows each:
-
-    >>> pl.scan_parquet("/path/to/file.parquet").sink_csv(
-    ...     PartitionMax("./out", max_size=100_000),
-    ... )  # doctest: +SKIP
-
-    See Also
-    --------
-    PartitionByKey
-    PartitionParted
-    polars.io.partition.BasePartitionContext
-    """
-
-    def __init__(
-        self,
-        base_path: str | Path,
-        *,
-        file_path: Callable[[BasePartitionContext], Path | str | IO[bytes] | IO[str]]
-        | None = None,
-        max_size: int,
-        per_partition_sort_by: str | Expr | Iterable[str | Expr] | None = None,
-        finish_callback: Callable[[DataFrame], None] | None = None,
-    ) -> None:
-        issue_unstable_warning("partitioning strategies are considered unstable.")
-        super().__init__(
-            PyPartitioning.new_max_size(
-                base_path=base_path,
-                file_path_cb=_cast_base_file_path_cb(file_path),
-                max_size=max_size,
-                per_partition_sort_by=_prepare_per_partition_sort_by(
-                    per_partition_sort_by
-                ),
-                finish_callback=_prepare_finish_callback(finish_callback),
-            )
-        )
-
-
-def _lower_by(
-    by: str | Expr | Sequence[str | Expr] | Mapping[str, Expr],
-) -> list[PyExpr]:
-    def to_expr(i: str | Expr) -> Expr:
-        if isinstance(i, str):
-            return col(i)
-        else:
-            return i
-
-    lowered_by: list[PyExpr]
-    if isinstance(by, str):
-        lowered_by = [col(by)._pyexpr]
-    elif isinstance(by, Expr):
-        lowered_by = [by._pyexpr]
-    elif isinstance(by, Sequence):
-        lowered_by = [to_expr(e)._pyexpr for e in by]
-    elif isinstance(by, Mapping):
-        lowered_by = [e.alias(n)._pyexpr for n, e in by.items()]
-    else:
-        msg = "invalid `by` type"
-        raise TypeError(msg)
-
-    return lowered_by
-
-
-class PartitionByKey(PartitioningScheme):
-    """
-    Partitioning scheme to write files split by the values of keys.
-
-    This partitioning scheme generates an arbitrary amount of files splitting
-    the data depending on what the value is of key expressions.
-
-    The amount of files that can be written is not limited. However, when
-    writing beyond a certain amount of files, the data for the remaining
-    partitions is buffered before writing to the file.
-
-    .. warning::
-        This functionality is currently considered **unstable**. It may be
-        changed at any point without it being considered a breaking change.
-
-    Parameters
-    ----------
-    base_path
-        The base path for the output files.
-
-        Use the `mkdir` option on the `sink_*` methods to ensure directories in
-        the path are created.
-    file_path
-        A callback to register or modify the output path for each partition
-        relative to the `base_path`. The callback provides a
-        :class:`polars.io.partition.KeyedPartitionContext` that contains information
-        about the partition.
-
-        If no callback is given, it defaults to
-        `{ctx.keys.hive_dirs()}/{ctx.in_part_idx}.{EXT}`.
-    by
-        The expressions to partition by.
-    include_key : bool
-        Whether to include the key columns in the output files.
-    per_partition_sort_by
-        Columns or expressions to sort over within each partition.
-
-        Note that this might increase the memory consumption needed for each partition.
-    finish_callback
-        A callback that gets called when the query finishes successfully.
-
-        For parquet files, the callback is given a dataframe with metrics about all
-        files written files.
-
-    Examples
-    --------
-    Split into a hive-partitioning style partition:
-
-    >>> (
-    ...     pl.DataFrame({"a": [1, 2, 3], "b": [5, 7, 9], "c": ["A", "B", "C"]})
-    ...     .lazy()
-    ...     .sink_parquet(
-    ...         PartitionByKey(
-    ...             "./out",
-    ...             by=[pl.col.a, pl.col.b],
-    ...             include_key=False,
-    ...         ),
-    ...         mkdir=True,
-    ...     )
-    ... )  # doctest: +SKIP
-
-    Split a parquet file by a column `year` into CSV files:
-
-    >>> pl.scan_parquet("/path/to/file.parquet").sink_csv(
-    ...     PartitionByKey(
-    ...         "./out/",
-    ...         file_path=lambda ctx: f"year={ctx.keys[0].str_value}.csv",
-    ...         by="year",
-    ...     ),
-    ... )  # doctest: +SKIP
-
-    See Also
-    --------
-    PartitionMaxSize
-    PartitionParted
-    polars.io.partition.KeyedPartitionContext
-    """
-
-    def __init__(
-        self,
-        base_path: str | Path,
-        *,
-        file_path: Callable[[KeyedPartitionContext], Path | str | IO[bytes] | IO[str]]
-        | None = None,
-        by: str | Expr | Sequence[str | Expr] | Mapping[str, Expr],
-        include_key: bool = True,
-        per_partition_sort_by: str | Expr | Iterable[str | Expr] | None = None,
-        finish_callback: Callable[[DataFrame], None] | None = None,
-    ) -> None:
-        issue_unstable_warning("partitioning strategies are considered unstable.")
-
-        lowered_by = _lower_by(by)
-        super().__init__(
-            PyPartitioning.new_by_key(
-                base_path=base_path,
-                file_path_cb=_cast_keyed_file_path_cb(file_path),
-                by=lowered_by,
-                include_key=include_key,
-                per_partition_sort_by=_prepare_per_partition_sort_by(
-                    per_partition_sort_by
-                ),
-                finish_callback=_prepare_finish_callback(finish_callback),
-            )
-        )
-
-
-class PartitionParted(PartitioningScheme):
-    """
-    Partitioning scheme to split parted dataframes.
-
-    This is a specialized version of :class:`PartitionByKey`. Where as
-    :class:`PartitionByKey` accepts data in any order, this scheme expects the input
-    data to be pre-grouped or pre-sorted. This scheme suffers a lot less overhead than
-    :class:`PartitionByKey`, but may not be always applicable.
-
-    Each new value of the key expressions starts a new partition, therefore repeating
-    the same value multiple times may overwrite previous partitions.
-
-    .. warning::
-        This functionality is currently considered **unstable**. It may be
-        changed at any point without it being considered a breaking change.
-
-    Parameters
-    ----------
-    base_path
-        The base path for the output files.
-
-        Use the `mkdir` option on the `sink_*` methods to ensure directories in
-        the path are created.
-    file_path
-        A callback to register or modify the output path for each partition
-        relative to the `base_path`.The callback provides a
-        :class:`polars.io.partition.KeyedPartitionContext` that contains information
-        about the partition.
-
-        If no callback is given, it defaults to
-        `{ctx.keys.hive_dirs()}/{ctx.in_part_idx}.{EXT}`.
-    by
-        The expressions to partition by.
-    include_key : bool
-        Whether to include the key columns in the output files.
-    per_partition_sort_by
-        Columns or expressions to sort over within each partition.
-
-        Note that this might increase the memory consumption needed for each partition.
-    finish_callback
-        A callback that gets called when the query finishes successfully.
-
-        For parquet files, the callback is given a dataframe with metrics about all
-        files written files.
-
-    Examples
-    --------
-    Split a parquet file by a column `year` into CSV files:
-
-    >>> pl.scan_parquet("/path/to/file.parquet").sink_csv(
-    ...     PartitionParted("./out", by="year"),
-    ...     mkdir=True,
-    ... )  # doctest: +SKIP
-
-    See Also
-    --------
-    PartitionMaxSize
-    PartitionByKey
-    polars.io.partition.KeyedPartitionContext
-    """
-
-    def __init__(
-        self,
-        base_path: str | Path,
-        *,
-        file_path: Callable[[KeyedPartitionContext], Path | str | IO[bytes] | IO[str]]
-        | None = None,
-        by: str | Expr | Sequence[str | Expr] | Mapping[str, Expr],
-        include_key: bool = True,
-        per_partition_sort_by: str | Expr | Iterable[str | Expr] | None = None,
-        finish_callback: Callable[[DataFrame], None] | None = None,
-    ) -> None:
-        issue_unstable_warning("partitioning strategies are considered unstable.")
-
-        lowered_by = _lower_by(by)
-        super().__init__(
-            PyPartitioning.new_by_key(
-                base_path=base_path,
-                file_path_cb=_cast_keyed_file_path_cb(file_path),
-                by=lowered_by,
-                include_key=include_key,
-                per_partition_sort_by=_prepare_per_partition_sort_by(
-                    per_partition_sort_by
-                ),
-                finish_callback=_prepare_finish_callback(finish_callback),
-            )
-        )

polars/io/plugins.py

@@ -1,187 +0,0 @@
-from __future__ import annotations
-
-import os
-import sys
-from collections.abc import Iterator
-from typing import TYPE_CHECKING, Callable
-
-import polars._reexport as pl
-from polars._utils.unstable import unstable
-
-if TYPE_CHECKING:
-    from collections.abc import Iterator
-    from typing import Callable
-
-    from polars import DataFrame, Expr, LazyFrame
-    from polars._typing import SchemaDict
-
-
-@unstable()
-def register_io_source(
-    io_source: Callable[
-        [list[str] | None, Expr | None, int | None, int | None], Iterator[DataFrame]
-    ],
-    *,
-    schema: Callable[[], SchemaDict] | SchemaDict,
-    validate_schema: bool = False,
-    is_pure: bool = False,
-) -> LazyFrame:
-    """
-    Register your IO plugin and initialize a LazyFrame.
-
-    See the `user guide <https://docs.pola.rs/user-guide/plugins/io_plugins>`_
-    for more information about plugins.
-
-    .. warning::
-        This functionality is considered **unstable**. It may be changed
-        at any point without it being considered a breaking change.
-
-
-    Parameters
-    ----------
-    io_source
-        Function that accepts the following arguments:
-            with_columns
-                Columns that are projected. The reader must
-                project these columns if applied
-            predicate
-                Polars expression. The reader must filter
-                their rows accordingly.
-            n_rows
-                Materialize only n rows from the source.
-                The reader can stop when `n_rows` are read.
-            batch_size
-                A hint of the ideal batch size the reader's
-                generator must produce.
-
-        The function should return a an iterator/generator
-        that produces DataFrames.
-    schema
-        Schema or function that when called produces the schema that the reader
-        will produce before projection pushdown.
-    validate_schema
-        Whether the engine should validate if the batches generated match
-        the given schema. It's an implementation error if this isn't
-        the case and can lead to bugs that are hard to solve.
-    is_pure
-        Whether the IO source is pure. Repeated occurrences of same IO source in
-        a LazyFrame plan can be de-duplicated during optimization if they are
-        pure.
-
-    Returns
-    -------
-    LazyFrame
-    """
-
-    def wrap(
-        with_columns: list[str] | None,
-        predicate: bytes | None,
-        n_rows: int | None,
-        batch_size: int | None,
-    ) -> tuple[Iterator[DataFrame], bool]:
-        parsed_predicate_success = True
-        parsed_predicate = None
-        if predicate:
-            try:
-                parsed_predicate = pl.Expr.deserialize(predicate)
-            except Exception as e:
-                if os.environ.get("POLARS_VERBOSE"):
-                    print(
-                        f"failed parsing IO plugin expression\n\nfilter will be handled on Polars' side: {e}",
-                        file=sys.stderr,
-                    )
-                parsed_predicate_success = False
-
-        return io_source(
-            with_columns, parsed_predicate, n_rows, batch_size
-        ), parsed_predicate_success
-
-    return pl.LazyFrame._scan_python_function(
-        schema=schema,
-        scan_fn=wrap,
-        pyarrow=False,
-        validate_schema=validate_schema,
-        is_pure=is_pure,
-    )
-
-
-@unstable()
-def _defer(
-    function: Callable[[], DataFrame],
-    *,
-    schema: SchemaDict | Callable[[], SchemaDict],
-    validate_schema: bool = True,
-) -> LazyFrame:
-    """
-    Deferred execution.
-
-    Takes a function that produces a `DataFrame` but defers execution until the
-    `LazyFrame` is collected.
-
-    Parameters
-    ----------
-    function
-        Function that takes no arguments and produces a `DataFrame`.
-    schema
-        Schema of the `DataFrame` the deferred function will return.
-        The caller must ensure this schema is correct.
-    validate_schema
-        Whether the engine should validate if the batches generated match
-        the given schema. It's an implementation error if this isn't
-        the case and can lead to bugs that are hard to solve.
-
-    Examples
-    --------
-    Delay DataFrame execution until query is executed.
-
-    >>> import numpy as np
-    >>> np.random.seed(0)
-    >>> lf = pl.defer(
-    ...     lambda: pl.DataFrame({"a": np.random.randn(3)}), schema={"a": pl.Float64}
-    ... )
-    >>> lf.collect()
-    shape: (3, 1)
-    ┌──────────┐
-    │ a        │
-    │ ---      │
-    │ f64      │
-    ╞══════════╡
-    │ 1.764052 │
-    │ 0.400157 │
-    │ 0.978738 │
-    └──────────┘
-
-     Run an eager source in Polars Cloud
-
-    >>> (
-    ...     pl.defer(
-    ...         lambda: pl.read_database("select * from tbl"),
-    ...         schema={"a": pl.Float64, "b": pl.Boolean},
-    ...     )
-    ...     .filter("b")
-    ...     .sum("a")
-    ...     .remote()
-    ...     .collect()
-    ... )  # doctest: +SKIP
-
-
-    """
-
-    def source(
-        with_columns: list[str] | None,
-        predicate: Expr | None,
-        n_rows: int | None,
-        batch_size: int | None,
-    ) -> Iterator[DataFrame]:
-        lf = function().lazy()
-        if with_columns is not None:
-            lf = lf.select(with_columns)
-        if predicate is not None:
-            lf = lf.filter(predicate)
-        if n_rows is not None:
-            lf = lf.limit(n_rows)
-        yield lf.collect()
-
-    return register_io_source(
-        io_source=source, schema=schema, validate_schema=validate_schema
-    )

polars/io/pyarrow_dataset/__init__.py

@@ -1,5 +0,0 @@
-from polars.io.pyarrow_dataset.functions import scan_pyarrow_dataset
-
-__all__ = [
-    "scan_pyarrow_dataset",
-]