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

polars/io/partition.py

@@ -0,0 +1,491 @@
+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

@@ -0,0 +1,187 @@
+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

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