06-polars Lazy API

class LazyFrame:
    """
    Representation of a Lazy computation graph/query against a DataFrame.

    This allows for whole-query optimisation in addition to parallelism, and
    is the preferred (and highest-performance) mode of operation for polars.

    def __init__(
        self,
        data: FrameInitTypes | None = None,
        schema: SchemaDefinition | None = None,
        *,
        schema_overrides: SchemaDict | None = None,
        strict: bool = True,
        orient: Orientation | None = None,
        infer_schema_length: int | None = N_INFER_DEFAULT,
        nan_to_null: bool = False,
        height: int | None = None,
    ) -> None:
        from polars.dataframe import DataFrame

        self._ldf = (
            DataFrame(
                data=data,
                schema=schema,
                schema_overrides=schema_overrides,
                strict=strict,
                orient=orient,
                infer_schema_length=infer_schema_length,
                nan_to_null=nan_to_null,
                height=height,
            )
            .lazy()
            ._ldf
        )

/// Lazy abstraction over an eager `DataFrame`.
///
/// It really is an abstraction over a logical plan. The methods of this struct will incrementally
/// modify a logical plan until output is requested (via [`collect`](crate::frame::LazyFrame::collect)).
#[derive(Clone, Default)]
#[must_use]
pub struct LazyFrame {
    pub logical_plan: DslPlan,
    pub(crate) opt_state: OptFlags,
    pub(crate) cached_arena: Arc<Mutex<Option<CachedArena>>>,
}

pub trait IntoLazy {
    fn lazy(self) -> LazyFrame;
}

impl IntoLazy for DataFrame {
    /// Convert the `DataFrame` into a `LazyFrame`
    fn lazy(self) -> LazyFrame {
        let lp = DslBuilder::from_existing_df(self).build();
        LazyFrame {
            logical_plan: lp,
            opt_state: Default::default(),
            cached_arena: Default::default(),
        }
    }
}

pub enum DslPlan {
    #[cfg(feature = "python")]
    PythonScan {
        options: crate::dsl::python_dsl::PythonOptionsDsl,
    },
    /// Filter on a boolean mask
    Filter {
        input: Arc<DslPlan>,
        predicate: Expr,
    },
    /// Cache the input at this point in the LP
    Cache {
        input: Arc<DslPlan>,
        id: UniqueId,
    },
    Scan {

    pub fn collect_with_engine(mut self, engine: Engine) -> PolarsResult<QueryResult> {
        let engine = match engine {
            Engine::Streaming => Engine::Streaming,
            _ if std::env::var("POLARS_FORCE_STREAMING").as_deref() == Ok("1") => Engine::Streaming,
            Engine::Auto => Engine::InMemory,
            v => v,
        };

        if engine != Engine::Streaming
            && std::env::var("POLARS_AUTO_STREAMING").as_deref() == Ok("1")
        {
            feature_gated!("streaming", {
                if let Some(r) = self.clone()._collect_with_streaming_suppress_todo_panic() {
                    return r;
                }
            })
        }
        match engine {
            Engine::Streaming => {
                feature_gated!("streaming", self = self.with_streaming(true))
            },
            Engine::Gpu => self = self.with_gpu(true),
            _ => (),
        }

    pub(crate) fn optimize_with_scratch(
        self,
        ir_arena: &mut Arena<IR>,
        expr_arena: &mut Arena<AExpr>,
        scratch: &mut Vec<Node>,
    ) -> PolarsResult<Node> {
        let mut opt_flags = self.opt_state;
        // Unset CSE
        // This can be turned on again during ir-conversion.
        #[allow(clippy::eq_op)]
        #[cfg(feature = "cse")]
        if opt_flags.contains(OptFlags::EAGER) {
            opt_flags &= !(OptFlags::COMM_SUBEXPR_ELIM | OptFlags::COMM_SUBEXPR_ELIM);
        }
        let root = to_alp(self.logical_plan, expr_arena, ir_arena, &mut opt_flags)?;

        let lp_top = optimize(
            root,
            opt_flags,
            ir_arena,
            expr_arena,
            scratch,
            apply_scan_predicate_to_scan_ir,
        )?;

        Ok(lp_top)
    }

bitflags! {
#[derive(Copy, Clone, Debug)]

        /// Only read columns that are used later in the query.
        const PROJECTION_PUSHDOWN = 1;
        /// Apply predicates/filters as early as possible.
        const PREDICATE_PUSHDOWN = 1 << 2;
        /// Cluster sequential `with_columns` calls to independent calls.
        const CLUSTER_WITH_COLUMNS = 1 << 3;
        /// Run many type coercion optimization rules until fixed point.
        const TYPE_COERCION = 1 << 4;
        /// Run many expression optimization rules until fixed point.
        const SIMPLIFY_EXPR = 1 << 5;
        /// Do type checking of the IR.
        const TYPE_CHECK = 1 << 6;
        /// Pushdown slices/limits.
        const SLICE_PUSHDOWN = 1 << 7;
        /// Run common-subplan-elimination. This elides duplicate plans and caches their
        /// outputs.
        const COMM_SUBPLAN_ELIM = 1 << 8;
        /// Run common-subexpression-elimination. This elides duplicate expressions and caches their
        /// outputs.
        const COMM_SUBEXPR_ELIM = 1 << 9;
        /// Is the query going to run on the GPU engine.
        const GPU = 1 << 10;
        /// Run on the streaming engine.
        const STREAMING = 1 << 11;
        /// Run every node eagerly. This turns off multi-node optimizations.
        const EAGER = 1 << 12;
        /// Try to estimate the number of rows so that joins can determine which side to keep in memory.
        const ROW_ESTIMATE = 1 << 13;
        /// Replace simple projections with a faster inlined projection that skips the expression engine.
        const FAST_PROJECTION = 1 << 14;
        /// Check if operations are order dependent and unset maintaining_order if
        /// the order would not be observed.
        const CHECK_ORDER_OBSERVE = 1 << 15;
        /// Collapse consecutive sort nodes and pull them up through selecting nodes.
        const SORT_COLLAPSE = 1 << 16;
        /// Pre-partition hive partitioned joins or group-by's
        /// Only works if PREDICATE_PUSHDOWN is set
        const PARTITION_HIVE = 1 << 17;
    }

pub fn optimize(
    mut root: Node,
    opt_flags: OptFlags,
    ir_arena: &mut Arena<IR>,
    expr_arena: &mut Arena<AExpr>,
    scratch: &mut Vec<Node>,
    apply_scan_predicate_to_scan_ir: fn(
        Node,
        &mut Arena<IR>,
        &mut Arena<AExpr>,
    ) -> PolarsResult<()>,
) -> PolarsResult<Node> {

    // Note: This can be in opt_flags in the future if needed.
    let pushdown_maintain_errors = pushdown_maintain_errors();

    let mut repeat_slice_pd_after_filter_pd = false;

    // Should be run before projection pushdown.
    // This allows columns only needed for filters to be dropped early.
    if opt_flags.predicate_pushdown() {
        let mut predicate_pushdown_opt = PredicatePushDown::new(
            pushdown_maintain_errors,
            opt_flags.streaming(),
            opt_flags.partition_hive(),
        );
        let ir = ir_arena.take(root);
        let ir = predicate_pushdown_opt.optimize(ir, ir_arena, expr_arena)?;
        ir_arena.replace(root, ir);
    }

    if opt_flags.fast_projection() {
        rules.push(Box::new(SimpleProjectionAndCollapse::new(
            opt_flags.eager(),
        )));
    }

    root = opt.optimize_loop(&mut rules, expr_arena, ir_arena, root)?;

    /// Return a String describing the naive (un-optimized) logical plan.
    pub fn describe_plan(&self) -> PolarsResult<String> {
        Ok(self.clone().to_alp()?.describe())
    }

    /// Return a String describing the naive (un-optimized) logical plan in tree format.
    pub fn describe_plan_tree(&self) -> PolarsResult<String> {
        Ok(self.clone().to_alp()?.describe_tree_format())
    }

    /// Return a String describing the optimized logical plan.
    ///
    /// Returns `Err` if optimizing the logical plan fails.
    pub fn describe_optimized_plan(&self) -> PolarsResult<String> {
        Ok(self.clone().to_alp_optimized()?.describe())
    }

    /// Return a String describing the optimized logical plan in tree format.
    ///
    /// Returns `Err` if optimizing the logical plan fails.
    pub fn describe_optimized_plan_tree(&self) -> PolarsResult<String> {
        Ok(self.clone().to_alp_optimized()?.describe_tree_format())
    }

    /// Return a String describing the logical plan.
    ///
    /// If `optimized` is `true`, explains the optimized plan. If `optimized` is `false`,
    /// explains the naive, un-optimized plan.
    pub fn explain(&self, optimized: bool) -> PolarsResult<String> {
        if optimized {
            self.describe_optimized_plan()
        } else {
            self.describe_plan()
        }
    }

        engine = _select_engine(engine)

        if optimized:
            optimizations = optimizations.__copy__()
            optimizations._pyoptflags.streaming = engine == "streaming"
            ldf = self._ldf.with_optimizations(optimizations._pyoptflags)
            if format == "tree":
                return ldf.describe_optimized_plan_tree()
            else:
                return ldf.describe_optimized_plan()

        if format == "tree":
            return self._ldf.describe_plan_tree()
        else:
            return self._ldf.describe_plan()

        let mut ir_plan = self.to_alp_optimized()?;

        ir_plan.ensure_root_node_is_sink();

        match engine {
            Engine::Streaming => feature_gated!("streaming", {
                polars_stream::run_query(
                    ir_plan.lp_top,
                    &mut ir_plan.lp_arena,
                    &mut ir_plan.expr_arena,
                )
            }),
            Engine::InMemory | Engine::Gpu => {

                let mut physical_plan = create_physical_plan(
                    ir_plan.lp_top,
                    &mut ir_plan.lp_arena,
                    &mut ir_plan.expr_arena,
                    BUILD_STREAMING_EXECUTOR,

            if k not in (  # except "private" kwargs

        if background:
            issue_unstable_warning("background mode is considered unstable.")
            return InProcessQuery(ldf.collect_concurrently())