DriverTrac/venv/lib/python3.12/site-packages/polars/_utils/cloud.py

44 lines
1.2 KiB
Python

from __future__ import annotations
from typing import TYPE_CHECKING
import polars._plr as plr
from polars.lazyframe.opt_flags import DEFAULT_QUERY_OPT_FLAGS
if TYPE_CHECKING:
from polars import LazyFrame, QueryOptFlags
def prepare_cloud_plan(
lf: LazyFrame,
*,
allow_local_scans: bool,
optimizations: QueryOptFlags = DEFAULT_QUERY_OPT_FLAGS,
) -> bytes:
"""
Prepare the given LazyFrame for execution on Polars Cloud.
Parameters
----------
lf
The LazyFrame to prepare.
allow_local_scans
Whether or not to allow local scans in the plan.
optimizations
Optimizations to enable or disable in the query optimizer.
Raises
------
InvalidOperationError
If the given LazyFrame is not eligible to be run on Polars Cloud.
The following conditions will disqualify a LazyFrame from being eligible:
- Contains a user-defined function
- Scans or sinks to a local filesystem
ComputeError
If the given LazyFrame cannot be serialized.
"""
optimizations = optimizations.__copy__()
pylf = lf._ldf.with_optimizations(optimizations._pyoptflags)
return plr.prepare_cloud_plan(pylf, allow_local_scans=allow_local_scans)