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

polars/io/pyarrow_dataset/anonymous_scan.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+from functools import partial
+from typing import TYPE_CHECKING
+
+import polars._reexport as pl
+from polars._dependencies import pyarrow as pa
+
+if TYPE_CHECKING:
+    from polars import DataFrame, LazyFrame
+
+
+def _scan_pyarrow_dataset(
+    ds: pa.dataset.Dataset,
+    *,
+    allow_pyarrow_filter: bool = True,
+    batch_size: int | None = None,
+) -> LazyFrame:
+    """
+    Pickle the partially applied function `_scan_pyarrow_dataset_impl`.
+
+    The bytes are then sent to the polars logical plan. It can be deserialized once
+    executed and ran.
+
+    Parameters
+    ----------
+    ds
+        pyarrow dataset
+    allow_pyarrow_filter
+        Allow predicates to be pushed down to pyarrow. This can lead to different
+        results if comparisons are done with null values as pyarrow handles this
+        different than polars does.
+    batch_size
+        The maximum row count for scanned pyarrow record batches.
+    """
+    func = partial(_scan_pyarrow_dataset_impl, ds, batch_size=batch_size)
+    return pl.LazyFrame._scan_python_function(
+        ds.schema, func, pyarrow=allow_pyarrow_filter
+    )
+
+
+def _scan_pyarrow_dataset_impl(
+    ds: pa.dataset.Dataset,
+    with_columns: list[str] | None,
+    predicate: str | None,
+    n_rows: int | None,
+    batch_size: int | None,
+) -> DataFrame:
+    """
+    Take the projected columns and materialize an arrow table.
+
+    Parameters
+    ----------
+    ds
+        pyarrow dataset
+    with_columns
+        Columns that are projected
+    predicate
+        pyarrow expression that can be evaluated with eval
+    n_rows:
+        Materialize only n rows from the arrow dataset
+    batch_size
+        The maximum row count for scanned pyarrow record batches.
+
+    Warnings
+    --------
+    Don't use this if you accept untrusted user inputs. Predicates will be evaluated
+    with python 'eval'. There is sanitation in place, but it is a possible attack
+    vector.
+
+    Returns
+    -------
+    DataFrame
+    """
+    from polars import from_arrow
+
+    _filter = None
+
+    if predicate:
+        from polars._utils.convert import (
+            to_py_date,
+            to_py_datetime,
+            to_py_time,
+            to_py_timedelta,
+        )
+        from polars.datatypes import Date, Datetime, Duration
+
+        _filter = eval(
+            predicate,
+            {
+                "pa": pa,
+                "Date": Date,
+                "Datetime": Datetime,
+                "Duration": Duration,
+                "to_py_date": to_py_date,
+                "to_py_datetime": to_py_datetime,
+                "to_py_time": to_py_time,
+                "to_py_timedelta": to_py_timedelta,
+            },
+        )
+
+    common_params = {"columns": with_columns, "filter": _filter}
+    if batch_size is not None:
+        common_params["batch_size"] = batch_size
+
+    if n_rows:
+        return from_arrow(ds.head(n_rows, **common_params))  # type: ignore[return-value]
+
+    return from_arrow(ds.to_table(**common_params))  # type: ignore[return-value]

polars/io/pyarrow_dataset/functions.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from polars._utils.unstable import unstable
+from polars.io.pyarrow_dataset.anonymous_scan import _scan_pyarrow_dataset
+
+if TYPE_CHECKING:
+    from polars import LazyFrame
+    from polars._dependencies import pyarrow as pa
+
+
+@unstable()
+def scan_pyarrow_dataset(
+    source: pa.dataset.Dataset,
+    *,
+    allow_pyarrow_filter: bool = True,
+    batch_size: int | None = None,
+) -> LazyFrame:
+    """
+    Scan a pyarrow dataset.
+
+    .. warning::
+        This functionality is considered **unstable**. It may be changed
+        at any point without it being considered a breaking change.
+
+    This can be useful to connect to cloud or partitioned datasets.
+
+    Parameters
+    ----------
+    source
+        Pyarrow dataset to scan.
+    allow_pyarrow_filter
+        Allow predicates to be pushed down to pyarrow. This can lead to different
+        results if comparisons are done with null values as pyarrow handles this
+        different than polars does.
+    batch_size
+        The maximum row count for scanned pyarrow record batches.
+
+    Warnings
+    --------
+    Don't use this if you accept untrusted user inputs. Predicates will be evaluated
+    with python 'eval'. There is sanitation in place, but it is a possible attack
+    vector.
+    This method can only can push down predicates that are allowed by PyArrow
+    (e.g. not the full Polars API).
+
+    If :func:`scan_parquet` works for your source, you should use that instead.
+
+    Notes
+    -----
+    When using partitioning, the appropriate `partitioning` option must be set on
+    `pyarrow.dataset.dataset` before passing to Polars or the partitioned-on column(s)
+    may not get passed to Polars.
+
+    Examples
+    --------
+    >>> import pyarrow.dataset as ds
+    >>> dset = ds.dataset("s3://my-partitioned-folder/", format="ipc")  # doctest: +SKIP
+    >>> (
+    ...     pl.scan_pyarrow_dataset(dset)
+    ...     .filter("bools")
+    ...     .select("bools", "floats", "date")
+    ...     .collect()
+    ... )  # doctest: +SKIP
+    shape: (1, 3)
+    ┌───────┬────────┬────────────┐
+    │ bools ┆ floats ┆ date       │
+    │ ---   ┆ ---    ┆ ---        │
+    │ bool  ┆ f64    ┆ date       │
+    ╞═══════╪════════╪════════════╡
+    │ true  ┆ 2.0    ┆ 1970-05-04 │
+    └───────┴────────┴────────────┘
+    """
+    return _scan_pyarrow_dataset(
+        source,
+        allow_pyarrow_filter=allow_pyarrow_filter,
+        batch_size=batch_size,
+    )

polars/io/scan_options/__init__.py

@@ -0,0 +1,5 @@
+from polars.io.scan_options.cast_options import ScanCastOptions
+
+__all__ = [
+    "ScanCastOptions",
+]

polars/io/scan_options/_options.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from polars._typing import (
+        ColumnMapping,
+        DefaultFieldValues,
+        DeletionFiles,
+        SchemaDict,
+    )
+    from polars.dataframe.frame import DataFrame
+    from polars.io.cloud.credential_provider._builder import CredentialProviderBuilder
+    from polars.io.scan_options.cast_options import ScanCastOptions
+
+from dataclasses import dataclass
+
+
+# TODO: Add `kw_only=True` after 3.9 support dropped
+@dataclass
+class ScanOptions:
+    """
+    Holds scan options that are generic over scan type.
+
+    For internal use. Most of the options will parse into `UnifiedScanArgs`.
+    """
+
+    row_index: tuple[str, int] | None = None
+    # (i64, usize)
+    pre_slice: tuple[int, int] | None = None
+    cast_options: ScanCastOptions | None = None
+    extra_columns: Literal["ignore", "raise"] = "raise"
+    missing_columns: Literal["insert", "raise"] = "raise"
+    include_file_paths: str | None = None
+
+    # For path expansion
+    glob: bool = True
+    hidden_file_prefix: Sequence[str] | None = None
+
+    # Hive
+    # Note: `None` means auto.
+    hive_partitioning: bool | None = None
+    hive_schema: SchemaDict | None = None
+    try_parse_hive_dates: bool = True
+
+    rechunk: bool = False
+    cache: bool = True
+
+    # Cloud
+    storage_options: list[tuple[str, str]] | None = None
+    credential_provider: CredentialProviderBuilder | None = None
+    retries: int = 2
+
+    column_mapping: ColumnMapping | None = None
+    default_values: DefaultFieldValues | None = None
+    deletion_files: DeletionFiles | None = None
+    table_statistics: DataFrame | None = None

polars/io/scan_options/cast_options.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
+from polars._utils.unstable import issue_unstable_warning
+
+if TYPE_CHECKING:
+    from collections.abc import Collection
+
+    from typing_extensions import TypeAlias
+
+
+FloatCastOption: TypeAlias = Literal["upcast", "downcast"]
+DatetimeCastOption: TypeAlias = Literal["nanosecond-downcast", "convert-timezone"]
+
+_DEFAULT_CAST_OPTIONS_ICEBERG: ScanCastOptions | None = None
+
+
+class ScanCastOptions:
+    """Options for scanning files."""
+
+    def __init__(
+        self,
+        *,
+        integer_cast: Literal["upcast", "forbid"] = "forbid",
+        float_cast: Literal["forbid"]
+        | FloatCastOption
+        | Collection[FloatCastOption] = "forbid",
+        datetime_cast: Literal["forbid"]
+        | DatetimeCastOption
+        | Collection[DatetimeCastOption] = "forbid",
+        missing_struct_fields: Literal["insert", "raise"] = "raise",
+        extra_struct_fields: Literal["ignore", "raise"] = "raise",
+        categorical_to_string: Literal["allow", "forbid"] = "forbid",
+        _internal_call: bool = False,
+    ) -> None:
+        """
+        Common configuration for scanning files.
+
+        .. warning::
+                This functionality is considered **unstable**. It may be changed
+                at any point without it being considered a breaking change.
+
+        Parameters
+        ----------
+        integer_cast
+            Configuration for casting from integer types:
+
+            * `upcast`: Allow lossless casting to wider integer types.
+            * `forbid`: Raises an error if dtypes do not match.
+
+        float_cast
+            Configuration for casting from float types:
+
+            * `upcast`: Allow casting to higher precision float types.
+            * `downcast`: Allow casting to lower precision float types.
+            * `forbid`: Raises an error if dtypes do not match.
+
+        datetime_cast
+            Configuration for casting from datetime types:
+
+            * `nanosecond-downcast`: Allow nanosecond precision datetime to be \
+            downcasted to any lower precision. This has a similar effect to \
+            PyArrow's `coerce_int96_timestamp_unit`.
+            * `convert-timezone`: Allow casting to a different timezone.
+            * `forbid`: Raises an error if dtypes do not match.
+
+        missing_struct_fields
+            Configuration for behavior when struct fields defined in the schema
+            are missing from the data:
+
+            * `insert`: Inserts the missing fields.
+            * `raise`: Raises an error.
+
+        extra_struct_fields
+            Configuration for behavior when extra struct fields outside of the
+            defined schema are encountered in the data:
+
+            * `ignore`: Silently ignores.
+            * `raise`: Raises an error.
+
+        categorical_to_string
+            Configuration for behavior when reading in a column whose expected
+            type is string, but type in the file is categorical.
+
+            * `allow`: Categorical is casted to string.
+            * `forbid`: Raises an error.
+
+        """
+        if not _internal_call:
+            issue_unstable_warning("ScanCastOptions is considered unstable.")
+
+        self.integer_cast = integer_cast
+        self.float_cast = float_cast
+        self.datetime_cast = datetime_cast
+        self.missing_struct_fields = missing_struct_fields
+        self.extra_struct_fields = extra_struct_fields
+        self.categorical_to_string = categorical_to_string
+
+    # Note: We don't cache this here, it's cached on the Rust-side.
+    @staticmethod
+    def _default() -> ScanCastOptions:
+        return ScanCastOptions(_internal_call=True)
+
+    @classmethod
+    def _default_iceberg(cls) -> ScanCastOptions:
+        """
+        Default options suitable for Iceberg / Deltalake.
+
+        This in general has all casting options enabled. Note: do not modify the
+        returned config object, it is a cached global object.
+        """
+        global _DEFAULT_CAST_OPTIONS_ICEBERG
+
+        if _DEFAULT_CAST_OPTIONS_ICEBERG is None:
+            _DEFAULT_CAST_OPTIONS_ICEBERG = ScanCastOptions(
+                integer_cast="upcast",
+                float_cast=["upcast", "downcast"],
+                datetime_cast=["nanosecond-downcast", "convert-timezone"],
+                missing_struct_fields="insert",
+                extra_struct_fields="ignore",
+                categorical_to_string="allow",
+                _internal_call=True,
+            )
+
+        return _DEFAULT_CAST_OPTIONS_ICEBERG

polars/io/spreadsheet/__init__.py

@@ -0,0 +1,6 @@
+from polars.io.spreadsheet.functions import read_excel, read_ods
+
+__all__ = [
+    "read_excel",
+    "read_ods",
+]

polars/io/spreadsheet/_utils.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, cast
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+
+@contextmanager
+def PortableTemporaryFile(
+    mode: str = "w+b",
+    *,
+    buffering: int = -1,
+    encoding: str | None = None,
+    newline: str | None = None,
+    suffix: str | None = None,
+    prefix: str | None = None,
+    dir: str | Path | None = None,
+    delete: bool = True,
+    errors: str | None = None,
+) -> Iterator[Any]:
+    """
+    Slightly more resilient version of the standard `NamedTemporaryFile`.
+
+    Plays better with Windows when using the 'delete' option.
+    """
+    from tempfile import NamedTemporaryFile
+
+    params = cast(
+        Any,
+        {
+            "mode": mode,
+            "buffering": buffering,
+            "encoding": encoding,
+            "newline": newline,
+            "suffix": suffix,
+            "prefix": prefix,
+            "dir": dir,
+            "delete": False,
+            "errors": errors,
+        },
+    )
+
+    with NamedTemporaryFile(**params) as tmp:
+        try:
+            yield tmp
+        finally:
+            tmp.close()
+            if delete:
+                Path(tmp.name).unlink(missing_ok=True)