Skip to main content

polars_lazy/frame/
mod.rs

1//! Lazy variant of a [DataFrame].
2#[cfg(feature = "python")]
3mod python;
4
5mod cached_arenas;
6mod err;
7#[cfg(not(target_arch = "wasm32"))]
8mod exitable;
9
10use std::num::NonZeroUsize;
11use std::sync::mpsc::{Receiver, sync_channel};
12use std::sync::{Arc, Mutex};
13
14pub use anonymous_scan::*;
15#[cfg(feature = "csv")]
16pub use csv::*;
17#[cfg(not(target_arch = "wasm32"))]
18pub use exitable::*;
19pub use file_list_reader::*;
20#[cfg(feature = "json")]
21pub use ndjson::*;
22#[cfg(feature = "parquet")]
23pub use parquet::*;
24use polars_compute::rolling::QuantileMethod;
25use polars_core::error::feature_gated;
26#[cfg(feature = "pivot")]
27use polars_core::frame::PivotColumnNaming;
28use polars_core::prelude::*;
29use polars_core::query_result::QueryResult;
30use polars_io::RowIndex;
31use polars_mem_engine::scan_predicate::functions::apply_scan_predicate_to_scan_ir;
32use polars_mem_engine::{Executor, create_multiple_physical_plans, create_physical_plan};
33use polars_ops::frame::{JoinBuildSide, JoinCoalesce, MaintainOrderJoin};
34#[cfg(feature = "is_between")]
35use polars_ops::prelude::ClosedInterval;
36pub use polars_plan::frame::{AllowedOptimizations, OptFlags};
37use polars_utils::pl_str::PlSmallStr;
38
39use crate::frame::cached_arenas::CachedArena;
40use crate::prelude::*;
41
42pub trait IntoLazy {
43    fn lazy(self) -> LazyFrame;
44}
45
46impl IntoLazy for DataFrame {
47    /// Convert the `DataFrame` into a `LazyFrame`
48    fn lazy(self) -> LazyFrame {
49        let lp = DslBuilder::from_existing_df(self).build();
50        LazyFrame {
51            logical_plan: lp,
52            opt_state: Default::default(),
53            cached_arena: Default::default(),
54        }
55    }
56}
57
58impl IntoLazy for LazyFrame {
59    fn lazy(self) -> LazyFrame {
60        self
61    }
62}
63
64/// Lazy abstraction over an eager `DataFrame`.
65///
66/// It really is an abstraction over a logical plan. The methods of this struct will incrementally
67/// modify a logical plan until output is requested (via [`collect`](crate::frame::LazyFrame::collect)).
68#[derive(Clone, Default)]
69#[must_use]
70pub struct LazyFrame {
71    pub logical_plan: DslPlan,
72    pub(crate) opt_state: OptFlags,
73    pub(crate) cached_arena: Arc<Mutex<Option<CachedArena>>>,
74}
75
76impl From<DslPlan> for LazyFrame {
77    fn from(plan: DslPlan) -> Self {
78        Self {
79            logical_plan: plan,
80            opt_state: OptFlags::default(),
81            cached_arena: Default::default(),
82        }
83    }
84}
85
86impl LazyFrame {
87    pub(crate) fn from_inner(
88        logical_plan: DslPlan,
89        opt_state: OptFlags,
90        cached_arena: Arc<Mutex<Option<CachedArena>>>,
91    ) -> Self {
92        Self {
93            logical_plan,
94            opt_state,
95            cached_arena,
96        }
97    }
98
99    pub(crate) fn get_plan_builder(self) -> DslBuilder {
100        DslBuilder::from(self.logical_plan)
101    }
102
103    fn get_opt_state(&self) -> OptFlags {
104        self.opt_state
105    }
106
107    pub fn from_logical_plan(logical_plan: DslPlan, opt_state: OptFlags) -> Self {
108        LazyFrame {
109            logical_plan,
110            opt_state,
111            cached_arena: Default::default(),
112        }
113    }
114
115    /// Get current optimizations.
116    pub fn get_current_optimizations(&self) -> OptFlags {
117        self.opt_state
118    }
119
120    /// Set allowed optimizations.
121    pub fn with_optimizations(mut self, opt_state: OptFlags) -> Self {
122        self.opt_state = opt_state;
123        self
124    }
125
126    /// Turn off all optimizations.
127    pub fn without_optimizations(self) -> Self {
128        self.with_optimizations(OptFlags::from_bits_truncate(0) | OptFlags::TYPE_COERCION)
129    }
130
131    /// Toggle projection pushdown optimization.
132    pub fn with_projection_pushdown(mut self, toggle: bool) -> Self {
133        self.opt_state.set(OptFlags::PROJECTION_PUSHDOWN, toggle);
134        self
135    }
136
137    /// Toggle cluster with columns optimization.
138    pub fn with_cluster_with_columns(mut self, toggle: bool) -> Self {
139        self.opt_state.set(OptFlags::CLUSTER_WITH_COLUMNS, toggle);
140        self
141    }
142
143    /// Check if operations are order dependent and unset maintaining_order if
144    /// the order would not be observed.
145    pub fn with_check_order(mut self, toggle: bool) -> Self {
146        self.opt_state.set(OptFlags::CHECK_ORDER_OBSERVE, toggle);
147        self
148    }
149
150    /// Toggle predicate pushdown optimization.
151    pub fn with_predicate_pushdown(mut self, toggle: bool) -> Self {
152        self.opt_state.set(OptFlags::PREDICATE_PUSHDOWN, toggle);
153        self
154    }
155
156    /// Toggle type coercion optimization.
157    pub fn with_type_coercion(mut self, toggle: bool) -> Self {
158        self.opt_state.set(OptFlags::TYPE_COERCION, toggle);
159        self
160    }
161
162    /// Toggle type check optimization.
163    pub fn with_type_check(mut self, toggle: bool) -> Self {
164        self.opt_state.set(OptFlags::TYPE_CHECK, toggle);
165        self
166    }
167
168    /// Toggle expression simplification optimization on or off.
169    pub fn with_simplify_expr(mut self, toggle: bool) -> Self {
170        self.opt_state.set(OptFlags::SIMPLIFY_EXPR, toggle);
171        self
172    }
173
174    /// Toggle common subplan elimination optimization on or off
175    #[cfg(feature = "cse")]
176    pub fn with_comm_subplan_elim(mut self, toggle: bool) -> Self {
177        self.opt_state.set(OptFlags::COMM_SUBPLAN_ELIM, toggle);
178        self
179    }
180
181    /// Toggle common subexpression elimination optimization on or off
182    #[cfg(feature = "cse")]
183    pub fn with_comm_subexpr_elim(mut self, toggle: bool) -> Self {
184        self.opt_state.set(OptFlags::COMM_SUBEXPR_ELIM, toggle);
185        self
186    }
187
188    /// Toggle slice pushdown optimization.
189    pub fn with_slice_pushdown(mut self, toggle: bool) -> Self {
190        self.opt_state.set(OptFlags::SLICE_PUSHDOWN, toggle);
191        self
192    }
193
194    #[cfg(feature = "streaming")]
195    pub fn with_streaming(mut self, toggle: bool) -> Self {
196        self.opt_state.set(OptFlags::STREAMING, toggle);
197        self
198    }
199
200    pub fn with_gpu(mut self, toggle: bool) -> Self {
201        self.opt_state.set(OptFlags::GPU, toggle);
202        self
203    }
204
205    /// Try to estimate the number of rows so that joins can determine which side to keep in memory.
206    pub fn with_row_estimate(mut self, toggle: bool) -> Self {
207        self.opt_state.set(OptFlags::ROW_ESTIMATE, toggle);
208        self
209    }
210
211    /// Run every node eagerly. This turns off multi-node optimizations.
212    pub fn _with_eager(mut self, toggle: bool) -> Self {
213        self.opt_state.set(OptFlags::EAGER, toggle);
214        self
215    }
216
217    /// Return a String describing the naive (un-optimized) logical plan.
218    pub fn describe_plan(&self) -> PolarsResult<String> {
219        Ok(self.clone().to_alp()?.describe())
220    }
221
222    /// Return a String describing the naive (un-optimized) logical plan in tree format.
223    pub fn describe_plan_tree(&self) -> PolarsResult<String> {
224        Ok(self.clone().to_alp()?.describe_tree_format())
225    }
226
227    /// Return a String describing the optimized logical plan.
228    ///
229    /// Returns `Err` if optimizing the logical plan fails.
230    pub fn describe_optimized_plan(&self) -> PolarsResult<String> {
231        Ok(self.clone().to_alp_optimized()?.describe())
232    }
233
234    /// Return a String describing the optimized logical plan in tree format.
235    ///
236    /// Returns `Err` if optimizing the logical plan fails.
237    pub fn describe_optimized_plan_tree(&self) -> PolarsResult<String> {
238        Ok(self.clone().to_alp_optimized()?.describe_tree_format())
239    }
240
241    /// Return a String describing the logical plan.
242    ///
243    /// If `optimized` is `true`, explains the optimized plan. If `optimized` is `false`,
244    /// explains the naive, un-optimized plan.
245    pub fn explain(&self, optimized: bool) -> PolarsResult<String> {
246        if optimized {
247            self.describe_optimized_plan()
248        } else {
249            self.describe_plan()
250        }
251    }
252
253    /// Add a sort operation to the logical plan.
254    ///
255    /// Sorts the LazyFrame by the column name specified using the provided options.
256    ///
257    /// # Example
258    ///
259    /// Sort DataFrame by 'sepal_width' column:
260    /// ```rust
261    /// # use polars_core::prelude::*;
262    /// # use polars_lazy::prelude::*;
263    /// fn sort_by_a(df: DataFrame) -> LazyFrame {
264    ///     df.lazy().sort(["sepal_width"], Default::default())
265    /// }
266    /// ```
267    /// Sort by a single column with specific order:
268    /// ```
269    /// # use polars_core::prelude::*;
270    /// # use polars_lazy::prelude::*;
271    /// fn sort_with_specific_order(df: DataFrame, descending: bool) -> LazyFrame {
272    ///     df.lazy().sort(
273    ///         ["sepal_width"],
274    ///         SortMultipleOptions::new()
275    ///             .with_order_descending(descending)
276    ///     )
277    /// }
278    /// ```
279    /// Sort by multiple columns with specifying order for each column:
280    /// ```
281    /// # use polars_core::prelude::*;
282    /// # use polars_lazy::prelude::*;
283    /// fn sort_by_multiple_columns_with_specific_order(df: DataFrame) -> LazyFrame {
284    ///     df.lazy().sort(
285    ///         ["sepal_width", "sepal_length"],
286    ///         SortMultipleOptions::new()
287    ///             .with_order_descending_multi([false, true])
288    ///     )
289    /// }
290    /// ```
291    /// See [`SortMultipleOptions`] for more options.
292    pub fn sort(self, by: impl IntoVec<PlSmallStr>, sort_options: SortMultipleOptions) -> Self {
293        let opt_state = self.get_opt_state();
294        let lp = self
295            .get_plan_builder()
296            .sort(by.into_vec().into_iter().map(col).collect(), sort_options)
297            .build();
298        Self::from_logical_plan(lp, opt_state)
299    }
300
301    /// Add a sort operation to the logical plan.
302    ///
303    /// Sorts the LazyFrame by the provided list of expressions, which will be turned into
304    /// concrete columns before sorting.
305    ///
306    /// See [`SortMultipleOptions`] for more options.
307    ///
308    /// # Example
309    ///
310    /// ```rust
311    /// use polars_core::prelude::*;
312    /// use polars_lazy::prelude::*;
313    ///
314    /// /// Sort DataFrame by 'sepal_width' column
315    /// fn example(df: DataFrame) -> LazyFrame {
316    ///       df.lazy()
317    ///         .sort_by_exprs(vec![col("sepal_width")], Default::default())
318    /// }
319    /// ```
320    pub fn sort_by_exprs<E: AsRef<[Expr]>>(
321        self,
322        by_exprs: E,
323        sort_options: SortMultipleOptions,
324    ) -> Self {
325        let by_exprs = by_exprs.as_ref().to_vec();
326        if by_exprs.is_empty() {
327            self
328        } else {
329            let opt_state = self.get_opt_state();
330            let lp = self.get_plan_builder().sort(by_exprs, sort_options).build();
331            Self::from_logical_plan(lp, opt_state)
332        }
333    }
334
335    pub fn top_k<E: AsRef<[Expr]>>(
336        self,
337        k: IdxSize,
338        by_exprs: E,
339        sort_options: SortMultipleOptions,
340    ) -> Self {
341        // this will optimize to top-k
342        self.sort_by_exprs(
343            by_exprs,
344            sort_options.with_order_reversed().with_nulls_last(true),
345        )
346        .slice(0, k)
347    }
348
349    pub fn bottom_k<E: AsRef<[Expr]>>(
350        self,
351        k: IdxSize,
352        by_exprs: E,
353        sort_options: SortMultipleOptions,
354    ) -> Self {
355        // this will optimize to bottom-k
356        self.sort_by_exprs(by_exprs, sort_options.with_nulls_last(true))
357            .slice(0, k)
358    }
359
360    /// Reverse the `DataFrame` from top to bottom.
361    ///
362    /// Row `i` becomes row `number_of_rows - i - 1`.
363    ///
364    /// # Example
365    ///
366    /// ```rust
367    /// use polars_core::prelude::*;
368    /// use polars_lazy::prelude::*;
369    ///
370    /// fn example(df: DataFrame) -> LazyFrame {
371    ///       df.lazy()
372    ///         .reverse()
373    /// }
374    /// ```
375    pub fn reverse(self) -> Self {
376        self.select(vec![col(PlSmallStr::from_static("*")).reverse()])
377    }
378
379    /// Rename columns in the DataFrame.
380    ///
381    /// `existing` and `new` are iterables of the same length containing the old and
382    /// corresponding new column names. Renaming happens to all `existing` columns
383    /// simultaneously, not iteratively. If `strict` is true, all columns in `existing`
384    /// must be present in the `LazyFrame` when `rename` is called; otherwise, only
385    /// those columns that are actually found will be renamed (others will be ignored).
386    pub fn rename<I, J, T, S>(self, existing: I, new: J, strict: bool) -> Self
387    where
388        I: IntoIterator<Item = T>,
389        J: IntoIterator<Item = S>,
390        T: AsRef<str>,
391        S: AsRef<str>,
392    {
393        let iter = existing.into_iter();
394        let cap = iter.size_hint().0;
395        let mut existing_vec: Vec<PlSmallStr> = Vec::with_capacity(cap);
396        let mut new_vec: Vec<PlSmallStr> = Vec::with_capacity(cap);
397
398        // TODO! should this error if `existing` and `new` have different lengths?
399        // Currently, the longer of the two is truncated.
400        for (existing, new) in iter.zip(new) {
401            let existing = existing.as_ref();
402            let new = new.as_ref();
403            if new != existing {
404                existing_vec.push(existing.into());
405                new_vec.push(new.into());
406            }
407        }
408
409        self.map_private(DslFunction::Rename {
410            existing: existing_vec.into(),
411            new: new_vec.into(),
412            strict,
413        })
414    }
415
416    /// Removes columns from the DataFrame.
417    /// Note that it's better to only select the columns you need
418    /// and let the projection pushdown optimize away the unneeded columns.
419    ///
420    /// Any given columns that are not in the schema will give a [`PolarsError::ColumnNotFound`]
421    /// error while materializing the [`LazyFrame`].
422    pub fn drop(self, columns: Selector) -> Self {
423        let opt_state = self.get_opt_state();
424        let lp = self.get_plan_builder().drop(columns).build();
425        Self::from_logical_plan(lp, opt_state)
426    }
427
428    /// Shift the values by a given period and fill the parts that will be empty due to this operation
429    /// with `Nones`.
430    ///
431    /// See the method on [Series](polars_core::series::SeriesTrait::shift) for more info on the `shift` operation.
432    pub fn shift<E: Into<Expr>>(self, n: E) -> Self {
433        self.select(vec![col(PlSmallStr::from_static("*")).shift(n.into())])
434    }
435
436    /// Shift the values by a given period and fill the parts that will be empty due to this operation
437    /// with the result of the `fill_value` expression.
438    ///
439    /// See the method on [Series](polars_core::series::SeriesTrait::shift) for more info on the `shift` operation.
440    pub fn shift_and_fill<E: Into<Expr>, IE: Into<Expr>>(self, n: E, fill_value: IE) -> Self {
441        self.select(vec![
442            col(PlSmallStr::from_static("*")).shift_and_fill(n.into(), fill_value.into()),
443        ])
444    }
445
446    /// Fill None values in the DataFrame with an expression.
447    pub fn fill_null<E: Into<Expr>>(self, fill_value: E) -> LazyFrame {
448        let opt_state = self.get_opt_state();
449        let lp = self.get_plan_builder().fill_null(fill_value.into()).build();
450        Self::from_logical_plan(lp, opt_state)
451    }
452
453    /// Fill NaN values in the DataFrame with an expression.
454    pub fn fill_nan<E: Into<Expr>>(self, fill_value: E) -> LazyFrame {
455        let opt_state = self.get_opt_state();
456        let lp = self.get_plan_builder().fill_nan(fill_value.into()).build();
457        Self::from_logical_plan(lp, opt_state)
458    }
459
460    /// Caches the result into a new LazyFrame.
461    ///
462    /// This should be used to prevent computations running multiple times.
463    pub fn cache(self) -> Self {
464        let opt_state = self.get_opt_state();
465        let lp = self.get_plan_builder().cache().build();
466        Self::from_logical_plan(lp, opt_state)
467    }
468
469    /// Cast named frame columns, resulting in a new LazyFrame with updated dtypes
470    pub fn cast(self, dtypes: PlHashMap<&str, DataType>, strict: bool) -> Self {
471        let cast_cols: Vec<Expr> = dtypes
472            .into_iter()
473            .map(|(name, dt)| {
474                let name = PlSmallStr::from_str(name);
475
476                if strict {
477                    col(name).strict_cast(dt)
478                } else {
479                    col(name).cast(dt)
480                }
481            })
482            .collect();
483
484        if cast_cols.is_empty() {
485            self
486        } else {
487            self.with_columns(cast_cols)
488        }
489    }
490
491    /// Cast all frame columns to the given dtype, resulting in a new LazyFrame
492    pub fn cast_all(self, dtype: impl Into<DataTypeExpr>, strict: bool) -> Self {
493        self.with_columns(vec![if strict {
494            col(PlSmallStr::from_static("*")).strict_cast(dtype)
495        } else {
496            col(PlSmallStr::from_static("*")).cast(dtype)
497        }])
498    }
499
500    pub fn optimize(
501        self,
502        lp_arena: &mut Arena<IR>,
503        expr_arena: &mut Arena<AExpr>,
504    ) -> PolarsResult<Node> {
505        self.optimize_with_scratch(lp_arena, expr_arena, &mut vec![])
506    }
507
508    pub fn to_alp_optimized(mut self) -> PolarsResult<IRPlan> {
509        let (mut lp_arena, mut expr_arena) = self.get_arenas();
510        let node = self.optimize_with_scratch(&mut lp_arena, &mut expr_arena, &mut vec![])?;
511
512        Ok(IRPlan::new(node, lp_arena, expr_arena))
513    }
514
515    pub fn to_alp(mut self) -> PolarsResult<IRPlan> {
516        let (mut lp_arena, mut expr_arena) = self.get_arenas();
517        let node = to_alp(
518            self.logical_plan,
519            &mut expr_arena,
520            &mut lp_arena,
521            &mut self.opt_state,
522        )?;
523        let plan = IRPlan::new(node, lp_arena, expr_arena);
524        Ok(plan)
525    }
526
527    pub(crate) fn optimize_with_scratch(
528        self,
529        lp_arena: &mut Arena<IR>,
530        expr_arena: &mut Arena<AExpr>,
531        scratch: &mut Vec<Node>,
532    ) -> PolarsResult<Node> {
533        let lp_top = optimize(
534            self.logical_plan,
535            self.opt_state,
536            lp_arena,
537            expr_arena,
538            scratch,
539            apply_scan_predicate_to_scan_ir,
540        )?;
541
542        Ok(lp_top)
543    }
544
545    fn prepare_collect_post_opt<P>(
546        mut self,
547        check_sink: bool,
548        query_start: Option<std::time::Instant>,
549        post_opt: P,
550    ) -> PolarsResult<(ExecutionState, Box<dyn Executor>, bool)>
551    where
552        P: FnOnce(
553            Node,
554            &mut Arena<IR>,
555            &mut Arena<AExpr>,
556            Option<std::time::Duration>,
557        ) -> PolarsResult<()>,
558    {
559        let (mut lp_arena, mut expr_arena) = self.get_arenas();
560
561        let mut scratch = vec![];
562        let lp_top = self.optimize_with_scratch(&mut lp_arena, &mut expr_arena, &mut scratch)?;
563
564        post_opt(
565            lp_top,
566            &mut lp_arena,
567            &mut expr_arena,
568            // Post optimization callback gets the time since the
569            // query was started as its "base" timepoint.
570            query_start.map(|s| s.elapsed()),
571        )?;
572
573        // sink should be replaced
574        let no_file_sink = if check_sink {
575            !matches!(
576                lp_arena.get(lp_top),
577                IR::Sink {
578                    payload: SinkTypeIR::File { .. },
579                    ..
580                }
581            )
582        } else {
583            true
584        };
585        let physical_plan = create_physical_plan(
586            lp_top,
587            &mut lp_arena,
588            &mut expr_arena,
589            BUILD_STREAMING_EXECUTOR,
590        )?;
591
592        let state = ExecutionState::new();
593        Ok((state, physical_plan, no_file_sink))
594    }
595
596    // post_opt: A function that is called after optimization. This can be used to modify the IR jit.
597    pub fn _collect_post_opt<P>(self, post_opt: P) -> PolarsResult<DataFrame>
598    where
599        P: FnOnce(
600            Node,
601            &mut Arena<IR>,
602            &mut Arena<AExpr>,
603            Option<std::time::Duration>,
604        ) -> PolarsResult<()>,
605    {
606        let (mut state, mut physical_plan, _) =
607            self.prepare_collect_post_opt(false, None, post_opt)?;
608        physical_plan.execute(&mut state)
609    }
610
611    #[allow(unused_mut)]
612    fn prepare_collect(
613        self,
614        check_sink: bool,
615        query_start: Option<std::time::Instant>,
616    ) -> PolarsResult<(ExecutionState, Box<dyn Executor>, bool)> {
617        self.prepare_collect_post_opt(check_sink, query_start, |_, _, _, _| Ok(()))
618    }
619
620    /// Execute all the lazy operations and collect them into a [`DataFrame`] using a specified
621    /// `engine`.
622    ///
623    /// The query is optimized prior to execution.
624    pub fn collect_with_engine(mut self, engine: Engine) -> PolarsResult<QueryResult> {
625        let engine = match engine {
626            Engine::Streaming => Engine::Streaming,
627            _ if std::env::var("POLARS_FORCE_STREAMING").as_deref() == Ok("1") => Engine::Streaming,
628            Engine::Auto => Engine::InMemory,
629            v => v,
630        };
631
632        if engine != Engine::Streaming
633            && std::env::var("POLARS_AUTO_STREAMING").as_deref() == Ok("1")
634        {
635            feature_gated!("streaming", {
636                if let Some(r) = self.clone()._collect_with_streaming_suppress_todo_panic() {
637                    return r;
638                }
639            })
640        }
641        match engine {
642            Engine::Streaming => {
643                feature_gated!("streaming", self = self.with_streaming(true))
644            },
645            Engine::Gpu => self = self.with_gpu(true),
646            _ => (),
647        }
648
649        let mut ir_plan = self.to_alp_optimized()?;
650
651        ir_plan.ensure_root_node_is_sink();
652
653        match engine {
654            Engine::Streaming => feature_gated!("streaming", {
655                polars_stream::run_query(
656                    ir_plan.lp_top,
657                    &mut ir_plan.lp_arena,
658                    &mut ir_plan.expr_arena,
659                )
660            }),
661            Engine::InMemory | Engine::Gpu => {
662                if let IR::SinkMultiple { inputs } = ir_plan.root() {
663                    polars_ensure!(
664                        engine != Engine::Gpu,
665                        InvalidOperation:
666                        "collect_all is not supported for the gpu engine"
667                    );
668
669                    return create_multiple_physical_plans(
670                        inputs.clone().as_slice(),
671                        &mut ir_plan.lp_arena,
672                        &mut ir_plan.expr_arena,
673                        BUILD_STREAMING_EXECUTOR,
674                    )?
675                    .execute()
676                    .map(QueryResult::Multiple);
677                }
678
679                let mut physical_plan = create_physical_plan(
680                    ir_plan.lp_top,
681                    &mut ir_plan.lp_arena,
682                    &mut ir_plan.expr_arena,
683                    BUILD_STREAMING_EXECUTOR,
684                )?;
685                let mut state = ExecutionState::new();
686                physical_plan.execute(&mut state).map(QueryResult::Single)
687            },
688            Engine::Auto => unreachable!(),
689        }
690    }
691
692    pub fn explain_all(plans: Vec<DslPlan>, opt_state: OptFlags) -> PolarsResult<String> {
693        let sink_multiple = LazyFrame {
694            logical_plan: DslPlan::SinkMultiple { inputs: plans },
695            opt_state,
696            cached_arena: Default::default(),
697        };
698        sink_multiple.explain(true)
699    }
700
701    pub fn collect_all_with_engine(
702        plans: Vec<DslPlan>,
703        engine: Engine,
704        opt_state: OptFlags,
705    ) -> PolarsResult<Vec<DataFrame>> {
706        if plans.is_empty() {
707            return Ok(Vec::new());
708        }
709
710        LazyFrame {
711            logical_plan: DslPlan::SinkMultiple { inputs: plans },
712            opt_state,
713            cached_arena: Default::default(),
714        }
715        .collect_with_engine(engine)
716        .map(|r| r.unwrap_multiple())
717    }
718
719    /// Execute all the lazy operations and collect them into a [`DataFrame`].
720    ///
721    /// The query is optimized prior to execution.
722    ///
723    /// # Example
724    ///
725    /// ```rust
726    /// use polars_core::prelude::*;
727    /// use polars_lazy::prelude::*;
728    ///
729    /// fn example(df: DataFrame) -> PolarsResult<DataFrame> {
730    ///     df.lazy()
731    ///       .group_by([col("foo")])
732    ///       .agg([col("bar").sum(), col("ham").mean().alias("avg_ham")])
733    ///       .collect()
734    /// }
735    /// ```
736    pub fn collect(self) -> PolarsResult<DataFrame> {
737        self.collect_with_engine(Engine::Auto).map(|r| match r {
738            QueryResult::Single(df) => df,
739            // TODO: Should return query results
740            QueryResult::Multiple(_) => DataFrame::empty(),
741        })
742    }
743
744    /// Collect the query in batches.
745    ///
746    /// If lazy is true the query will not start until the first poll (or until
747    /// start is called on CollectBatches).
748    #[cfg(feature = "async")]
749    pub fn collect_batches(
750        self,
751        engine: Engine,
752        maintain_order: bool,
753        chunk_size: Option<NonZeroUsize>,
754        lazy: bool,
755    ) -> PolarsResult<CollectBatches> {
756        let (send, recv) = sync_channel(1);
757        let runner_send = send.clone();
758        let ldf = self.sink_batches(
759            PlanCallback::new(move |df| {
760                // Stop if receiver has closed.
761                let send_result = send.send(Ok(df));
762                Ok(send_result.is_err())
763            }),
764            maintain_order,
765            chunk_size,
766        )?;
767        let runner = move || {
768            // We use spawn_blocking here as it has a high blocking thread pool limit.
769            polars_core::runtime::ASYNC.spawn_blocking(move || {
770                if let Err(e) = ldf.collect_with_engine(engine) {
771                    runner_send.send(Err(e)).ok();
772                }
773            });
774        };
775
776        let mut collect_batches = CollectBatches {
777            recv,
778            runner: Some(Box::new(runner)),
779        };
780        if !lazy {
781            collect_batches.start();
782        }
783        Ok(collect_batches)
784    }
785
786    // post_opt: A function that is called after optimization. This can be used to modify the IR jit.
787    // This version does profiling of the node execution.
788    pub fn _profile_post_opt<P>(self, post_opt: P) -> PolarsResult<(DataFrame, DataFrame)>
789    where
790        P: FnOnce(
791            Node,
792            &mut Arena<IR>,
793            &mut Arena<AExpr>,
794            Option<std::time::Duration>,
795        ) -> PolarsResult<()>,
796    {
797        let query_start = std::time::Instant::now();
798        let (mut state, mut physical_plan, _) =
799            self.prepare_collect_post_opt(false, Some(query_start), post_opt)?;
800        state.time_nodes(query_start);
801        let out = physical_plan.execute(&mut state)?;
802        let timer_df = state.finish_timer()?;
803        Ok((out, timer_df))
804    }
805
806    /// Profile a LazyFrame.
807    ///
808    /// This will run the query and return a tuple
809    /// containing the materialized DataFrame and a DataFrame that contains profiling information
810    /// of each node that is executed.
811    ///
812    /// The units of the timings are microseconds.
813    pub fn profile(self) -> PolarsResult<(DataFrame, DataFrame)> {
814        self._profile_post_opt(|_, _, _, _| Ok(()))
815    }
816
817    pub fn sink_batches(
818        mut self,
819        function: PlanCallback<DataFrame, bool>,
820        maintain_order: bool,
821        chunk_size: Option<NonZeroUsize>,
822    ) -> PolarsResult<Self> {
823        use polars_plan::prelude::sink::CallbackSinkType;
824
825        polars_ensure!(
826            !matches!(self.logical_plan, DslPlan::Sink { .. }),
827            InvalidOperation: "cannot create a sink on top of another sink"
828        );
829
830        self.logical_plan = DslPlan::Sink {
831            input: Arc::new(self.logical_plan),
832            payload: SinkType::Callback(CallbackSinkType {
833                function,
834                maintain_order,
835                chunk_size,
836            }),
837        };
838
839        Ok(self)
840    }
841
842    /// Collect with the streaming engine. Returns `None` if the streaming engine panics with a todo!.
843    #[cfg(feature = "streaming")]
844    fn _collect_with_streaming_suppress_todo_panic(
845        mut self,
846    ) -> Option<PolarsResult<polars_core::query_result::QueryResult>> {
847        self.opt_state |= OptFlags::STREAMING;
848        let mut ir_plan = match self.to_alp_optimized() {
849            Ok(v) => v,
850            Err(e) => return Some(Err(e)),
851        };
852
853        ir_plan.ensure_root_node_is_sink();
854
855        let f = || {
856            polars_stream::run_query(
857                ir_plan.lp_top,
858                &mut ir_plan.lp_arena,
859                &mut ir_plan.expr_arena,
860            )
861        };
862
863        match std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)) {
864            Ok(v) => Some(v),
865            Err(e) => {
866                // Fallback to normal engine if error is due to not being implemented
867                // and auto_streaming is set, otherwise propagate error.
868                if e.downcast_ref::<&str>()
869                    .is_some_and(|s| s.starts_with("not yet implemented"))
870                {
871                    if polars_core::config::verbose() {
872                        eprintln!(
873                            "caught unimplemented error in new streaming engine, falling back to normal engine"
874                        );
875                    }
876                    None
877                } else {
878                    std::panic::resume_unwind(e)
879                }
880            },
881        }
882    }
883
884    pub fn sink(
885        mut self,
886        sink_type: SinkDestination,
887        file_format: FileWriteFormat,
888        unified_sink_args: UnifiedSinkArgs,
889    ) -> PolarsResult<Self> {
890        polars_ensure!(
891            !matches!(self.logical_plan, DslPlan::Sink { .. }),
892            InvalidOperation: "cannot create a sink on top of another sink"
893        );
894
895        self.logical_plan = DslPlan::Sink {
896            input: Arc::new(self.logical_plan),
897            payload: match sink_type {
898                SinkDestination::File { target } => SinkType::File(FileSinkOptions {
899                    target,
900                    file_format,
901                    unified_sink_args,
902                }),
903                SinkDestination::Partitioned {
904                    base_path,
905                    file_path_provider,
906                    partition_strategy,
907                    max_rows_per_file,
908                    approximate_bytes_per_file,
909                } => SinkType::Partitioned(PartitionedSinkOptions {
910                    base_path,
911                    file_path_provider,
912                    partition_strategy,
913                    file_format,
914                    unified_sink_args,
915                    max_rows_per_file,
916                    approximate_bytes_per_file,
917                }),
918            },
919        };
920        Ok(self)
921    }
922
923    /// Filter frame rows that match a predicate expression.
924    ///
925    /// The expression must yield boolean values (note that rows where the
926    /// predicate resolves to `null` are *not* included in the resulting frame).
927    ///
928    /// # Example
929    ///
930    /// ```rust
931    /// use polars_core::prelude::*;
932    /// use polars_lazy::prelude::*;
933    ///
934    /// fn example(df: DataFrame) -> LazyFrame {
935    ///       df.lazy()
936    ///         .filter(col("sepal_width").is_not_null())
937    ///         .select([col("sepal_width"), col("sepal_length")])
938    /// }
939    /// ```
940    pub fn filter(self, predicate: Expr) -> Self {
941        let opt_state = self.get_opt_state();
942        let lp = self.get_plan_builder().filter(predicate).build();
943        Self::from_logical_plan(lp, opt_state)
944    }
945
946    /// Remove frame rows that match a predicate expression.
947    ///
948    /// The expression must yield boolean values (note that rows where the
949    /// predicate resolves to `null` are *not* removed from the resulting frame).
950    ///
951    /// # Example
952    ///
953    /// ```rust
954    /// use polars_core::prelude::*;
955    /// use polars_lazy::prelude::*;
956    ///
957    /// fn example(df: DataFrame) -> LazyFrame {
958    ///       df.lazy()
959    ///         .remove(col("sepal_width").is_null())
960    ///         .select([col("sepal_width"), col("sepal_length")])
961    /// }
962    /// ```
963    pub fn remove(self, predicate: Expr) -> Self {
964        self.filter(predicate.neq_missing(lit(true)))
965    }
966
967    /// Select (and optionally rename, with [`alias`](crate::dsl::Expr::alias)) columns from the query.
968    ///
969    /// Columns can be selected with [`col`];
970    /// If you want to select all columns use `col(PlSmallStr::from_static("*"))`.
971    ///
972    /// # Example
973    ///
974    /// ```rust
975    /// use polars_core::prelude::*;
976    /// use polars_lazy::prelude::*;
977    ///
978    /// /// This function selects column "foo" and column "bar".
979    /// /// Column "bar" is renamed to "ham".
980    /// fn example(df: DataFrame) -> LazyFrame {
981    ///       df.lazy()
982    ///         .select([col("foo"),
983    ///                   col("bar").alias("ham")])
984    /// }
985    ///
986    /// /// This function selects all columns except "foo"
987    /// fn exclude_a_column(df: DataFrame) -> LazyFrame {
988    ///       df.lazy()
989    ///         .select([all().exclude_cols(["foo"]).as_expr()])
990    /// }
991    /// ```
992    pub fn select<E: AsRef<[Expr]>>(self, exprs: E) -> Self {
993        let exprs = exprs.as_ref().to_vec();
994        self.select_impl(
995            exprs,
996            ProjectionOptions {
997                run_parallel: true,
998                duplicate_check: true,
999                should_broadcast: true,
1000            },
1001        )
1002    }
1003
1004    pub fn select_seq<E: AsRef<[Expr]>>(self, exprs: E) -> Self {
1005        let exprs = exprs.as_ref().to_vec();
1006        self.select_impl(
1007            exprs,
1008            ProjectionOptions {
1009                run_parallel: false,
1010                duplicate_check: true,
1011                should_broadcast: true,
1012            },
1013        )
1014    }
1015
1016    fn select_impl(self, exprs: Vec<Expr>, options: ProjectionOptions) -> Self {
1017        let opt_state = self.get_opt_state();
1018        let lp = self.get_plan_builder().project(exprs, options).build();
1019        Self::from_logical_plan(lp, opt_state)
1020    }
1021
1022    /// Performs a "group-by" on a `LazyFrame`, producing a [`LazyGroupBy`], which can subsequently be aggregated.
1023    ///
1024    /// Takes a list of expressions to group on.
1025    ///
1026    /// # Example
1027    ///
1028    /// ```rust
1029    /// use polars_core::prelude::*;
1030    /// use polars_lazy::prelude::*;
1031    ///
1032    /// fn example(df: DataFrame) -> LazyFrame {
1033    ///       df.lazy()
1034    ///        .group_by([col("date")])
1035    ///        .agg([
1036    ///            col("rain").min().alias("min_rain"),
1037    ///            col("rain").sum().alias("sum_rain"),
1038    ///            col("rain").quantile(lit(0.5), QuantileMethod::Nearest).alias("median_rain"),
1039    ///        ])
1040    /// }
1041    /// ```
1042    pub fn group_by<E: AsRef<[IE]>, IE: Into<Expr> + Clone>(self, by: E) -> LazyGroupBy {
1043        let keys = by
1044            .as_ref()
1045            .iter()
1046            .map(|e| e.clone().into())
1047            .collect::<Vec<_>>();
1048        let opt_state = self.get_opt_state();
1049
1050        #[cfg(feature = "dynamic_group_by")]
1051        {
1052            LazyGroupBy {
1053                logical_plan: self.logical_plan,
1054                opt_state,
1055                keys,
1056                predicates: vec![],
1057                maintain_order: false,
1058                dynamic_options: None,
1059                rolling_options: None,
1060            }
1061        }
1062
1063        #[cfg(not(feature = "dynamic_group_by"))]
1064        {
1065            LazyGroupBy {
1066                logical_plan: self.logical_plan,
1067                opt_state,
1068                keys,
1069                predicates: vec![],
1070                maintain_order: false,
1071            }
1072        }
1073    }
1074
1075    /// Create rolling groups based on a time column.
1076    ///
1077    /// Also works for index values of type UInt32, UInt64, Int32, or Int64.
1078    ///
1079    /// Different from a [`group_by_dynamic`][`Self::group_by_dynamic`], the windows are now determined by the
1080    /// individual values and are not of constant intervals. For constant intervals use
1081    /// *group_by_dynamic*
1082    #[cfg(feature = "dynamic_group_by")]
1083    pub fn rolling<E: AsRef<[Expr]>>(
1084        mut self,
1085        index_column: Expr,
1086        group_by: E,
1087        mut options: RollingGroupOptions,
1088    ) -> LazyGroupBy {
1089        if let Expr::Column(name) = index_column {
1090            options.index_column = name;
1091        } else {
1092            let output_field = index_column
1093                .to_field(&self.collect_schema().unwrap())
1094                .unwrap();
1095            return self.with_column(index_column).rolling(
1096                Expr::Column(output_field.name().clone()),
1097                group_by,
1098                options,
1099            );
1100        }
1101        let opt_state = self.get_opt_state();
1102        LazyGroupBy {
1103            logical_plan: self.logical_plan,
1104            opt_state,
1105            predicates: vec![],
1106            keys: group_by.as_ref().to_vec(),
1107            maintain_order: true,
1108            dynamic_options: None,
1109            rolling_options: Some(options),
1110        }
1111    }
1112
1113    /// Group based on a time value (or index value of type Int32, Int64).
1114    ///
1115    /// Time windows are calculated and rows are assigned to windows. Different from a
1116    /// normal group_by is that a row can be member of multiple groups. The time/index
1117    /// window could be seen as a rolling window, with a window size determined by
1118    /// dates/times/values instead of slots in the DataFrame.
1119    ///
1120    /// A window is defined by:
1121    ///
1122    /// - every: interval of the window
1123    /// - period: length of the window
1124    /// - offset: offset of the window
1125    ///
1126    /// The `group_by` argument should be empty `[]` if you don't want to combine this
1127    /// with a ordinary group_by on these keys.
1128    #[cfg(feature = "dynamic_group_by")]
1129    pub fn group_by_dynamic<E: AsRef<[Expr]>>(
1130        mut self,
1131        index_column: Expr,
1132        group_by: E,
1133        mut options: DynamicGroupOptions,
1134    ) -> LazyGroupBy {
1135        if let Expr::Column(name) = index_column {
1136            options.index_column = name;
1137        } else {
1138            let output_field = index_column
1139                .to_field(&self.collect_schema().unwrap())
1140                .unwrap();
1141            return self.with_column(index_column).group_by_dynamic(
1142                Expr::Column(output_field.name().clone()),
1143                group_by,
1144                options,
1145            );
1146        }
1147        let opt_state = self.get_opt_state();
1148        LazyGroupBy {
1149            logical_plan: self.logical_plan,
1150            opt_state,
1151            predicates: vec![],
1152            keys: group_by.as_ref().to_vec(),
1153            maintain_order: true,
1154            dynamic_options: Some(options),
1155            rolling_options: None,
1156        }
1157    }
1158
1159    /// Similar to [`group_by`][`Self::group_by`], but order of the DataFrame is maintained.
1160    pub fn group_by_stable<E: AsRef<[IE]>, IE: Into<Expr> + Clone>(self, by: E) -> LazyGroupBy {
1161        let keys = by
1162            .as_ref()
1163            .iter()
1164            .map(|e| e.clone().into())
1165            .collect::<Vec<_>>();
1166        let opt_state = self.get_opt_state();
1167
1168        #[cfg(feature = "dynamic_group_by")]
1169        {
1170            LazyGroupBy {
1171                logical_plan: self.logical_plan,
1172                opt_state,
1173                keys,
1174                predicates: vec![],
1175                maintain_order: true,
1176                dynamic_options: None,
1177                rolling_options: None,
1178            }
1179        }
1180
1181        #[cfg(not(feature = "dynamic_group_by"))]
1182        {
1183            LazyGroupBy {
1184                logical_plan: self.logical_plan,
1185                opt_state,
1186                keys,
1187                predicates: vec![],
1188                maintain_order: true,
1189            }
1190        }
1191    }
1192
1193    /// Left anti join this query with another lazy query.
1194    ///
1195    /// Matches on the values of the expressions `left_on` and `right_on`. For more
1196    /// flexible join logic, see [`join`](LazyFrame::join) or
1197    /// [`join_builder`](LazyFrame::join_builder).
1198    ///
1199    /// # Example
1200    ///
1201    /// ```rust
1202    /// use polars_core::prelude::*;
1203    /// use polars_lazy::prelude::*;
1204    /// fn anti_join_dataframes(ldf: LazyFrame, other: LazyFrame) -> LazyFrame {
1205    ///         ldf
1206    ///         .anti_join(other, col("foo"), col("bar").cast(DataType::String))
1207    /// }
1208    /// ```
1209    #[cfg(feature = "semi_anti_join")]
1210    pub fn anti_join<E: Into<Expr>>(self, other: LazyFrame, left_on: E, right_on: E) -> LazyFrame {
1211        self.join(
1212            other,
1213            [left_on.into()],
1214            [right_on.into()],
1215            JoinArgs::new(JoinType::Anti),
1216        )
1217    }
1218
1219    /// Creates the Cartesian product from both frames, preserving the order of the left keys.
1220    #[cfg(feature = "cross_join")]
1221    pub fn cross_join(self, other: LazyFrame, suffix: Option<PlSmallStr>) -> LazyFrame {
1222        self.join(
1223            other,
1224            vec![],
1225            vec![],
1226            JoinArgs::new(JoinType::Cross).with_suffix(suffix),
1227        )
1228    }
1229
1230    /// Left outer join this query with another lazy query.
1231    ///
1232    /// Matches on the values of the expressions `left_on` and `right_on`. For more
1233    /// flexible join logic, see [`join`](LazyFrame::join) or
1234    /// [`join_builder`](LazyFrame::join_builder).
1235    ///
1236    /// # Example
1237    ///
1238    /// ```rust
1239    /// use polars_core::prelude::*;
1240    /// use polars_lazy::prelude::*;
1241    /// fn left_join_dataframes(ldf: LazyFrame, other: LazyFrame) -> LazyFrame {
1242    ///         ldf
1243    ///         .left_join(other, col("foo"), col("bar"))
1244    /// }
1245    /// ```
1246    pub fn left_join<E: Into<Expr>>(self, other: LazyFrame, left_on: E, right_on: E) -> LazyFrame {
1247        self.join(
1248            other,
1249            [left_on.into()],
1250            [right_on.into()],
1251            JoinArgs::new(JoinType::Left),
1252        )
1253    }
1254
1255    /// Inner join this query with another lazy query.
1256    ///
1257    /// Matches on the values of the expressions `left_on` and `right_on`. For more
1258    /// flexible join logic, see [`join`](LazyFrame::join) or
1259    /// [`join_builder`](LazyFrame::join_builder).
1260    ///
1261    /// # Example
1262    ///
1263    /// ```rust
1264    /// use polars_core::prelude::*;
1265    /// use polars_lazy::prelude::*;
1266    /// fn inner_join_dataframes(ldf: LazyFrame, other: LazyFrame) -> LazyFrame {
1267    ///         ldf
1268    ///         .inner_join(other, col("foo"), col("bar").cast(DataType::String))
1269    /// }
1270    /// ```
1271    pub fn inner_join<E: Into<Expr>>(self, other: LazyFrame, left_on: E, right_on: E) -> LazyFrame {
1272        self.join(
1273            other,
1274            [left_on.into()],
1275            [right_on.into()],
1276            JoinArgs::new(JoinType::Inner),
1277        )
1278    }
1279
1280    /// Full outer join this query with another lazy query.
1281    ///
1282    /// Matches on the values of the expressions `left_on` and `right_on`. For more
1283    /// flexible join logic, see [`join`](LazyFrame::join) or
1284    /// [`join_builder`](LazyFrame::join_builder).
1285    ///
1286    /// # Example
1287    ///
1288    /// ```rust
1289    /// use polars_core::prelude::*;
1290    /// use polars_lazy::prelude::*;
1291    /// fn full_join_dataframes(ldf: LazyFrame, other: LazyFrame) -> LazyFrame {
1292    ///         ldf
1293    ///         .full_join(other, col("foo"), col("bar"))
1294    /// }
1295    /// ```
1296    pub fn full_join<E: Into<Expr>>(self, other: LazyFrame, left_on: E, right_on: E) -> LazyFrame {
1297        self.join(
1298            other,
1299            [left_on.into()],
1300            [right_on.into()],
1301            JoinArgs::new(JoinType::Full),
1302        )
1303    }
1304
1305    /// Left semi join this query with another lazy query.
1306    ///
1307    /// Matches on the values of the expressions `left_on` and `right_on`. For more
1308    /// flexible join logic, see [`join`](LazyFrame::join) or
1309    /// [`join_builder`](LazyFrame::join_builder).
1310    ///
1311    /// # Example
1312    ///
1313    /// ```rust
1314    /// use polars_core::prelude::*;
1315    /// use polars_lazy::prelude::*;
1316    /// fn semi_join_dataframes(ldf: LazyFrame, other: LazyFrame) -> LazyFrame {
1317    ///         ldf
1318    ///         .semi_join(other, col("foo"), col("bar").cast(DataType::String))
1319    /// }
1320    /// ```
1321    #[cfg(feature = "semi_anti_join")]
1322    pub fn semi_join<E: Into<Expr>>(self, other: LazyFrame, left_on: E, right_on: E) -> LazyFrame {
1323        self.join(
1324            other,
1325            [left_on.into()],
1326            [right_on.into()],
1327            JoinArgs::new(JoinType::Semi),
1328        )
1329    }
1330
1331    /// Generic function to join two LazyFrames.
1332    ///
1333    /// `join` can join on multiple columns, given as two list of expressions, and with a
1334    /// [`JoinType`] specified by `how`. Non-joined column names in the right DataFrame
1335    /// that already exist in this DataFrame are suffixed with `"_right"`. For control
1336    /// over how columns are renamed and parallelization options, use
1337    /// [`join_builder`](LazyFrame::join_builder).
1338    ///
1339    /// Any provided `args.slice` parameter is not considered, but set by the internal optimizer.
1340    ///
1341    /// # Example
1342    ///
1343    /// ```rust
1344    /// use polars_core::prelude::*;
1345    /// use polars_lazy::prelude::*;
1346    ///
1347    /// fn example(ldf: LazyFrame, other: LazyFrame) -> LazyFrame {
1348    ///         ldf
1349    ///         .join(other, [col("foo"), col("bar")], [col("foo"), col("bar")], JoinArgs::new(JoinType::Inner))
1350    /// }
1351    /// ```
1352    pub fn join<E: AsRef<[Expr]>>(
1353        self,
1354        other: LazyFrame,
1355        left_on: E,
1356        right_on: E,
1357        args: JoinArgs,
1358    ) -> LazyFrame {
1359        let left_on = left_on.as_ref().to_vec();
1360        let right_on = right_on.as_ref().to_vec();
1361
1362        self._join_impl(other, left_on, right_on, args)
1363    }
1364
1365    fn _join_impl(
1366        self,
1367        other: LazyFrame,
1368        left_on: Vec<Expr>,
1369        right_on: Vec<Expr>,
1370        args: JoinArgs,
1371    ) -> LazyFrame {
1372        let JoinArgs {
1373            how,
1374            validation,
1375            suffix,
1376            slice,
1377            nulls_equal,
1378            coalesce,
1379            maintain_order,
1380            build_side,
1381        } = args;
1382
1383        if slice.is_some() {
1384            panic!("impl error: slice is not handled")
1385        }
1386
1387        let mut builder = self
1388            .join_builder()
1389            .with(other)
1390            .left_on(left_on)
1391            .right_on(right_on)
1392            .how(how)
1393            .validate(validation)
1394            .join_nulls(nulls_equal)
1395            .coalesce(coalesce)
1396            .maintain_order(maintain_order)
1397            .build_side(build_side);
1398
1399        if let Some(suffix) = suffix {
1400            builder = builder.suffix(suffix);
1401        }
1402
1403        // Note: args.slice is set by the optimizer
1404        builder.finish()
1405    }
1406
1407    /// Consume `self` and return a [`JoinBuilder`] to customize a join on this LazyFrame.
1408    ///
1409    /// After the `JoinBuilder` has been created and set up, calling
1410    /// [`finish()`](JoinBuilder::finish) on it will give back the `LazyFrame`
1411    /// representing the `join` operation.
1412    pub fn join_builder(self) -> JoinBuilder {
1413        JoinBuilder::new(self)
1414    }
1415
1416    /// Gathers rows from this DataFrame based on the indices in idxs.
1417    ///
1418    /// idxs must only have a single column of indices.
1419    pub fn gather(self, idxs: LazyFrame, null_on_oob: bool) -> LazyFrame {
1420        let opt_state = self.get_opt_state();
1421        let lp = self
1422            .get_plan_builder()
1423            .gather(idxs.logical_plan, null_on_oob)
1424            .build();
1425        Self::from_logical_plan(lp, opt_state)
1426    }
1427
1428    /// Add or replace a column, given as an expression, to a DataFrame.
1429    ///
1430    /// # Example
1431    ///
1432    /// ```rust
1433    /// use polars_core::prelude::*;
1434    /// use polars_lazy::prelude::*;
1435    /// fn add_column(df: DataFrame) -> LazyFrame {
1436    ///     df.lazy()
1437    ///         .with_column(
1438    ///             when(col("sepal_length").lt(lit(5.0)))
1439    ///             .then(lit(10))
1440    ///             .otherwise(lit(1))
1441    ///             .alias("new_column_name"),
1442    ///         )
1443    /// }
1444    /// ```
1445    pub fn with_column(self, expr: Expr) -> LazyFrame {
1446        let opt_state = self.get_opt_state();
1447        let lp = self
1448            .get_plan_builder()
1449            .with_columns(
1450                vec![expr],
1451                ProjectionOptions {
1452                    run_parallel: false,
1453                    duplicate_check: true,
1454                    should_broadcast: true,
1455                },
1456            )
1457            .build();
1458        Self::from_logical_plan(lp, opt_state)
1459    }
1460
1461    /// Add or replace multiple columns, given as expressions, to a DataFrame.
1462    ///
1463    /// # Example
1464    ///
1465    /// ```rust
1466    /// use polars_core::prelude::*;
1467    /// use polars_lazy::prelude::*;
1468    /// fn add_columns(df: DataFrame) -> LazyFrame {
1469    ///     df.lazy()
1470    ///         .with_columns(
1471    ///             vec![lit(10).alias("foo"), lit(100).alias("bar")]
1472    ///          )
1473    /// }
1474    /// ```
1475    pub fn with_columns<E: AsRef<[Expr]>>(self, exprs: E) -> LazyFrame {
1476        let exprs = exprs.as_ref().to_vec();
1477        self.with_columns_impl(
1478            exprs,
1479            ProjectionOptions {
1480                run_parallel: true,
1481                duplicate_check: true,
1482                should_broadcast: true,
1483            },
1484        )
1485    }
1486
1487    /// Add or replace multiple columns to a DataFrame, but evaluate them sequentially.
1488    pub fn with_columns_seq<E: AsRef<[Expr]>>(self, exprs: E) -> LazyFrame {
1489        let exprs = exprs.as_ref().to_vec();
1490        self.with_columns_impl(
1491            exprs,
1492            ProjectionOptions {
1493                run_parallel: false,
1494                duplicate_check: true,
1495                should_broadcast: true,
1496            },
1497        )
1498    }
1499
1500    /// Match or evolve to a certain schema.
1501    pub fn match_to_schema(
1502        self,
1503        schema: SchemaRef,
1504        per_column: Arc<[MatchToSchemaPerColumn]>,
1505        extra_columns: ExtraColumnsPolicy,
1506    ) -> LazyFrame {
1507        let opt_state = self.get_opt_state();
1508        let lp = self
1509            .get_plan_builder()
1510            .match_to_schema(schema, per_column, extra_columns)
1511            .build();
1512        Self::from_logical_plan(lp, opt_state)
1513    }
1514
1515    pub fn pipe_with_schema(
1516        self,
1517        callback: PlanCallback<(Vec<DslPlan>, Vec<SchemaRef>), DslPlan>,
1518    ) -> Self {
1519        let opt_state = self.get_opt_state();
1520        let lp = self
1521            .get_plan_builder()
1522            .pipe_with_schema(vec![], callback)
1523            .build();
1524        Self::from_logical_plan(lp, opt_state)
1525    }
1526
1527    pub fn pipe_with_schemas(
1528        self,
1529        others: Vec<LazyFrame>,
1530        callback: PlanCallback<(Vec<DslPlan>, Vec<SchemaRef>), DslPlan>,
1531    ) -> Self {
1532        let opt_state = self.get_opt_state();
1533        let lp = self
1534            .get_plan_builder()
1535            .pipe_with_schema(
1536                others.into_iter().map(|lf| lf.logical_plan).collect(),
1537                callback,
1538            )
1539            .build();
1540        Self::from_logical_plan(lp, opt_state)
1541    }
1542
1543    fn with_columns_impl(self, exprs: Vec<Expr>, options: ProjectionOptions) -> LazyFrame {
1544        let opt_state = self.get_opt_state();
1545        let lp = self.get_plan_builder().with_columns(exprs, options).build();
1546        Self::from_logical_plan(lp, opt_state)
1547    }
1548
1549    pub fn with_context<C: AsRef<[LazyFrame]>>(self, contexts: C) -> LazyFrame {
1550        let contexts = contexts
1551            .as_ref()
1552            .iter()
1553            .map(|lf| lf.logical_plan.clone())
1554            .collect();
1555        let opt_state = self.get_opt_state();
1556        let lp = self.get_plan_builder().with_context(contexts).build();
1557        Self::from_logical_plan(lp, opt_state)
1558    }
1559
1560    /// Aggregate all the columns as their maximum values.
1561    ///
1562    /// Aggregated columns will have the same names as the original columns.
1563    pub fn max(self) -> Self {
1564        self.map_private(DslFunction::Stats(StatsFunction::Max))
1565    }
1566
1567    /// Aggregate all the columns as their minimum values.
1568    ///
1569    /// Aggregated columns will have the same names as the original columns.
1570    pub fn min(self) -> Self {
1571        self.map_private(DslFunction::Stats(StatsFunction::Min))
1572    }
1573
1574    /// Aggregate all the columns as their sum values.
1575    ///
1576    /// Aggregated columns will have the same names as the original columns.
1577    ///
1578    /// - Boolean columns will sum to a `u32` containing the number of `true`s.
1579    /// - For integer columns, the ordinary checks for overflow are performed:
1580    ///   if running in `debug` mode, overflows will panic, whereas in `release` mode overflows will
1581    ///   silently wrap.
1582    /// - String columns will sum to None.
1583    pub fn sum(self) -> Self {
1584        self.map_private(DslFunction::Stats(StatsFunction::Sum))
1585    }
1586
1587    /// Aggregate all the columns as their mean values.
1588    ///
1589    /// - Boolean and integer columns are converted to `f64` before computing the mean.
1590    /// - String columns will have a mean of None.
1591    pub fn mean(self) -> Self {
1592        self.map_private(DslFunction::Stats(StatsFunction::Mean))
1593    }
1594
1595    /// Aggregate all the columns as their median values.
1596    ///
1597    /// - Boolean and integer results are converted to `f64`. However, they are still
1598    ///   susceptible to overflow before this conversion occurs.
1599    /// - String columns will sum to None.
1600    pub fn median(self) -> Self {
1601        self.map_private(DslFunction::Stats(StatsFunction::Median))
1602    }
1603
1604    /// Aggregate all the columns as their quantile values.
1605    pub fn quantile(self, quantile: Expr, method: QuantileMethod) -> Self {
1606        self.map_private(DslFunction::Stats(StatsFunction::Quantile {
1607            quantile,
1608            method,
1609        }))
1610    }
1611
1612    /// Aggregate all the columns as their standard deviation values.
1613    ///
1614    /// `ddof` is the "Delta Degrees of Freedom"; `N - ddof` will be the denominator when
1615    /// computing the variance, where `N` is the number of rows.
1616    /// > In standard statistical practice, `ddof=1` provides an unbiased estimator of the
1617    /// > variance of a hypothetical infinite population. `ddof=0` provides a maximum
1618    /// > likelihood estimate of the variance for normally distributed variables. The
1619    /// > standard deviation computed in this function is the square root of the estimated
1620    /// > variance, so even with `ddof=1`, it will not be an unbiased estimate of the
1621    /// > standard deviation per se.
1622    ///
1623    /// Source: [Numpy](https://numpy.org/doc/stable/reference/generated/numpy.std.html#)
1624    pub fn std(self, ddof: u8) -> Self {
1625        self.map_private(DslFunction::Stats(StatsFunction::Std { ddof }))
1626    }
1627
1628    /// Aggregate all the columns as their variance values.
1629    ///
1630    /// `ddof` is the "Delta Degrees of Freedom"; `N - ddof` will be the denominator when
1631    /// computing the variance, where `N` is the number of rows.
1632    /// > In standard statistical practice, `ddof=1` provides an unbiased estimator of the
1633    /// > variance of a hypothetical infinite population. `ddof=0` provides a maximum
1634    /// > likelihood estimate of the variance for normally distributed variables.
1635    ///
1636    /// Source: [Numpy](https://numpy.org/doc/stable/reference/generated/numpy.var.html#)
1637    pub fn var(self, ddof: u8) -> Self {
1638        self.map_private(DslFunction::Stats(StatsFunction::Var { ddof }))
1639    }
1640
1641    /// Apply explode operation. [See eager explode](polars_core::frame::DataFrame::explode).
1642    pub fn explode(self, columns: Selector, options: ExplodeOptions) -> LazyFrame {
1643        self.explode_impl(columns, options, false)
1644    }
1645
1646    /// Apply explode operation. [See eager explode](polars_core::frame::DataFrame::explode).
1647    fn explode_impl(
1648        self,
1649        columns: Selector,
1650        options: ExplodeOptions,
1651        allow_empty: bool,
1652    ) -> LazyFrame {
1653        let opt_state = self.get_opt_state();
1654        let lp = self
1655            .get_plan_builder()
1656            .explode(columns, options, allow_empty)
1657            .build();
1658        Self::from_logical_plan(lp, opt_state)
1659    }
1660
1661    /// Aggregate all the columns as the sum of their null value count.
1662    pub fn null_count(self) -> LazyFrame {
1663        self.select(vec![col(PlSmallStr::from_static("*")).null_count()])
1664    }
1665
1666    /// Drop non-unique rows and maintain the order of kept rows.
1667    ///
1668    /// `subset` is an optional `Vec` of column names to consider for uniqueness; if
1669    /// `None`, all columns are considered.
1670    pub fn unique_stable(
1671        self,
1672        subset: Option<Selector>,
1673        keep_strategy: UniqueKeepStrategy,
1674    ) -> LazyFrame {
1675        let subset = subset.map(|s| vec![Expr::Selector(s)]);
1676        self.unique_stable_generic(subset, keep_strategy)
1677    }
1678
1679    pub fn unique_stable_generic(
1680        self,
1681        subset: Option<Vec<Expr>>,
1682        keep_strategy: UniqueKeepStrategy,
1683    ) -> LazyFrame {
1684        let opt_state = self.get_opt_state();
1685        let options = DistinctOptionsDSL {
1686            subset,
1687            maintain_order: true,
1688            keep_strategy,
1689        };
1690        let lp = self.get_plan_builder().distinct(options).build();
1691        Self::from_logical_plan(lp, opt_state)
1692    }
1693
1694    /// Drop non-unique rows without maintaining the order of kept rows.
1695    ///
1696    /// The order of the kept rows may change; to maintain the original row order, use
1697    /// [`unique_stable`](LazyFrame::unique_stable).
1698    ///
1699    /// `subset` is an optional `Vec` of column names to consider for uniqueness; if None,
1700    /// all columns are considered.
1701    pub fn unique(self, subset: Option<Selector>, keep_strategy: UniqueKeepStrategy) -> LazyFrame {
1702        let subset = subset.map(|s| vec![Expr::Selector(s)]);
1703        self.unique_generic(subset, keep_strategy)
1704    }
1705
1706    pub fn unique_generic(
1707        self,
1708        subset: Option<Vec<Expr>>,
1709        keep_strategy: UniqueKeepStrategy,
1710    ) -> LazyFrame {
1711        let opt_state = self.get_opt_state();
1712        let options = DistinctOptionsDSL {
1713            subset,
1714            maintain_order: false,
1715            keep_strategy,
1716        };
1717        let lp = self.get_plan_builder().distinct(options).build();
1718        Self::from_logical_plan(lp, opt_state)
1719    }
1720
1721    /// Drop rows containing one or more NaN values.
1722    ///
1723    /// `subset` is an optional `Vec` of column names to consider for NaNs; if None, all
1724    /// floating point columns are considered.
1725    pub fn drop_nans(self, subset: Option<Selector>) -> LazyFrame {
1726        let opt_state = self.get_opt_state();
1727        let lp = self.get_plan_builder().drop_nans(subset).build();
1728        Self::from_logical_plan(lp, opt_state)
1729    }
1730
1731    /// Drop rows containing one or more None values.
1732    ///
1733    /// `subset` is an optional `Vec` of column names to consider for nulls; if None, all
1734    /// columns are considered.
1735    pub fn drop_nulls(self, subset: Option<Selector>) -> LazyFrame {
1736        let opt_state = self.get_opt_state();
1737        let lp = self.get_plan_builder().drop_nulls(subset).build();
1738        Self::from_logical_plan(lp, opt_state)
1739    }
1740
1741    /// Slice the DataFrame using an offset (starting row) and a length.
1742    ///
1743    /// If `offset` is negative, it is counted from the end of the DataFrame. For
1744    /// instance, `lf.slice(-5, 3)` gets three rows, starting at the row fifth from the
1745    /// end.
1746    ///
1747    /// If `offset` and `len` are such that the slice extends beyond the end of the
1748    /// DataFrame, the portion between `offset` and the end will be returned. In this
1749    /// case, the number of rows in the returned DataFrame will be less than `len`.
1750    pub fn slice(self, offset: i64, len: IdxSize) -> LazyFrame {
1751        let opt_state = self.get_opt_state();
1752        let lp = self.get_plan_builder().slice(offset, len).build();
1753        Self::from_logical_plan(lp, opt_state)
1754    }
1755
1756    /// Remove all the rows of the LazyFrame.
1757    pub fn clear(self) -> LazyFrame {
1758        self.slice(0, 0)
1759    }
1760
1761    /// Get the first row.
1762    ///
1763    /// Equivalent to `self.slice(0, 1)`.
1764    pub fn first(self) -> LazyFrame {
1765        self.slice(0, 1)
1766    }
1767
1768    /// Get the last row.
1769    ///
1770    /// Equivalent to `self.slice(-1, 1)`.
1771    pub fn last(self) -> LazyFrame {
1772        self.slice(-1, 1)
1773    }
1774
1775    /// Get the last `n` rows.
1776    ///
1777    /// Equivalent to `self.slice(-(n as i64), n)`.
1778    pub fn tail(self, n: IdxSize) -> LazyFrame {
1779        let neg_tail = -(n as i64);
1780        self.slice(neg_tail, n)
1781    }
1782
1783    #[cfg(feature = "pivot")]
1784    #[expect(clippy::too_many_arguments)]
1785    pub fn pivot(
1786        self,
1787        on: Selector,
1788        on_columns: Arc<DataFrame>,
1789        index: Selector,
1790        values: Selector,
1791        agg: Expr,
1792        maintain_order: bool,
1793        separator: PlSmallStr,
1794        column_naming: PivotColumnNaming,
1795    ) -> LazyFrame {
1796        let opt_state = self.get_opt_state();
1797        let lp = self
1798            .get_plan_builder()
1799            .pivot(
1800                on,
1801                on_columns,
1802                index,
1803                values,
1804                agg,
1805                maintain_order,
1806                separator,
1807                column_naming,
1808            )
1809            .build();
1810        Self::from_logical_plan(lp, opt_state)
1811    }
1812
1813    /// Unpivot the DataFrame from wide to long format.
1814    ///
1815    /// See [`UnpivotArgsIR`] for information on how to unpivot a DataFrame.
1816    #[cfg(feature = "pivot")]
1817    pub fn unpivot(self, args: UnpivotArgsDSL) -> LazyFrame {
1818        let opt_state = self.get_opt_state();
1819        let lp = self.get_plan_builder().unpivot(args).build();
1820        Self::from_logical_plan(lp, opt_state)
1821    }
1822
1823    /// Limit the DataFrame to the first `n` rows.
1824    pub fn limit(self, n: IdxSize) -> LazyFrame {
1825        self.slice(0, n)
1826    }
1827
1828    /// Apply a function/closure once the logical plan get executed.
1829    ///
1830    /// The function has access to the whole materialized DataFrame at the time it is
1831    /// called.
1832    ///
1833    /// To apply specific functions to specific columns, use [`Expr::map`] in conjunction
1834    /// with `LazyFrame::with_column` or `with_columns`.
1835    ///
1836    /// ## Warning
1837    /// This can blow up in your face if the schema is changed due to the operation. The
1838    /// optimizer relies on a correct schema.
1839    ///
1840    /// You can toggle certain optimizations off.
1841    pub fn map<F>(
1842        self,
1843        function: F,
1844        optimizations: AllowedOptimizations,
1845        schema: Option<Arc<dyn UdfSchema>>,
1846        name: Option<&'static str>,
1847    ) -> LazyFrame
1848    where
1849        F: 'static + Fn(DataFrame) -> PolarsResult<DataFrame> + Send + Sync,
1850    {
1851        let opt_state = self.get_opt_state();
1852        let lp = self
1853            .get_plan_builder()
1854            .map(
1855                function,
1856                optimizations,
1857                schema,
1858                PlSmallStr::from_static(name.unwrap_or("ANONYMOUS UDF")),
1859            )
1860            .build();
1861        Self::from_logical_plan(lp, opt_state)
1862    }
1863
1864    #[cfg(feature = "python")]
1865    pub fn map_python(
1866        self,
1867        function: polars_utils::python_function::PythonFunction,
1868        optimizations: AllowedOptimizations,
1869        schema: Option<SchemaRef>,
1870        validate_output: bool,
1871    ) -> LazyFrame {
1872        let opt_state = self.get_opt_state();
1873        let lp = self
1874            .get_plan_builder()
1875            .map_python(function, optimizations, schema, validate_output)
1876            .build();
1877        Self::from_logical_plan(lp, opt_state)
1878    }
1879
1880    pub(crate) fn map_private(self, function: DslFunction) -> LazyFrame {
1881        let opt_state = self.get_opt_state();
1882        let lp = self.get_plan_builder().map_private(function).build();
1883        Self::from_logical_plan(lp, opt_state)
1884    }
1885
1886    /// Add a new column at index 0 that counts the rows.
1887    ///
1888    /// `name` is the name of the new column. `offset` is where to start counting from; if
1889    /// `None`, it is set to `0`.
1890    ///
1891    /// # Warning
1892    /// This can have a negative effect on query performance. This may for instance block
1893    /// predicate pushdown optimization.
1894    pub fn with_row_index<S>(self, name: S, offset: Option<IdxSize>) -> LazyFrame
1895    where
1896        S: Into<PlSmallStr>,
1897    {
1898        let name = name.into();
1899
1900        match &self.logical_plan {
1901            v @ DslPlan::Scan {
1902                scan_type,
1903                unified_scan_args,
1904                ..
1905            } if unified_scan_args.row_index.is_none()
1906                && !matches!(
1907                    &**scan_type,
1908                    FileScanDsl::Anonymous { .. } | FileScanDsl::ExpandedPaths { .. }
1909                ) =>
1910            {
1911                let DslPlan::Scan {
1912                    sources,
1913                    mut unified_scan_args,
1914                    scan_type,
1915                    cached_ir: _,
1916                } = v.clone()
1917                else {
1918                    unreachable!()
1919                };
1920
1921                unified_scan_args.row_index = Some(RowIndex {
1922                    name,
1923                    offset: offset.unwrap_or(0),
1924                });
1925
1926                DslPlan::Scan {
1927                    sources,
1928                    unified_scan_args,
1929                    scan_type,
1930                    cached_ir: Default::default(),
1931                }
1932                .into()
1933            },
1934            _ => self.map_private(DslFunction::RowIndex { name, offset }),
1935        }
1936    }
1937
1938    /// Return the number of non-null elements for each column.
1939    pub fn count(self) -> LazyFrame {
1940        self.select(vec![col(PlSmallStr::from_static("*")).count()])
1941    }
1942
1943    /// Unnest the given `Struct` columns: the fields of the `Struct` type will be
1944    /// inserted as columns.
1945    #[cfg(feature = "dtype-struct")]
1946    pub fn unnest(self, cols: Selector, separator: Option<PlSmallStr>) -> Self {
1947        self.map_private(DslFunction::Unnest {
1948            columns: cols,
1949            separator,
1950        })
1951    }
1952
1953    #[cfg(feature = "merge_sorted")]
1954    pub fn merge_sorted<S>(
1955        self,
1956        other: LazyFrame,
1957        key: S,
1958        maintain_order: bool,
1959    ) -> PolarsResult<LazyFrame>
1960    where
1961        S: Into<PlSmallStr>,
1962    {
1963        let key = key.into();
1964
1965        let lp = DslPlan::MergeSorted {
1966            input_left: Arc::new(self.logical_plan),
1967            input_right: Arc::new(other.logical_plan),
1968            key,
1969            maintain_order,
1970        };
1971        Ok(LazyFrame::from_logical_plan(lp, self.opt_state))
1972    }
1973
1974    pub fn hint(self, hint: HintIR) -> PolarsResult<LazyFrame> {
1975        let lp = DslPlan::MapFunction {
1976            input: Arc::new(self.logical_plan),
1977            function: DslFunction::Hint(hint),
1978        };
1979        Ok(LazyFrame::from_logical_plan(lp, self.opt_state))
1980    }
1981}
1982
1983/// Utility struct for lazy group_by operation.
1984#[derive(Clone)]
1985pub struct LazyGroupBy {
1986    pub logical_plan: DslPlan,
1987    opt_state: OptFlags,
1988    keys: Vec<Expr>,
1989    predicates: Vec<Expr>,
1990    maintain_order: bool,
1991    #[cfg(feature = "dynamic_group_by")]
1992    dynamic_options: Option<DynamicGroupOptions>,
1993    #[cfg(feature = "dynamic_group_by")]
1994    rolling_options: Option<RollingGroupOptions>,
1995}
1996
1997impl From<LazyGroupBy> for LazyFrame {
1998    fn from(lgb: LazyGroupBy) -> Self {
1999        Self {
2000            logical_plan: lgb.logical_plan,
2001            opt_state: lgb.opt_state,
2002            cached_arena: Default::default(),
2003        }
2004    }
2005}
2006
2007impl LazyGroupBy {
2008    /// Filter groups with a predicate after aggregation.
2009    ///
2010    /// Similarly to the [LazyGroupBy::agg] method, the predicate must run an aggregation as it
2011    /// is evaluated on the groups.
2012    /// This method can be chained in which case all predicates must evaluate to `true` for a
2013    /// group to be kept.
2014    ///
2015    /// # Example
2016    ///
2017    /// ```rust
2018    /// use polars_core::prelude::*;
2019    /// use polars_lazy::prelude::*;
2020    ///
2021    /// fn example(df: DataFrame) -> LazyFrame {
2022    ///       df.lazy()
2023    ///        .group_by_stable([col("date")])
2024    ///        .having(col("rain").sum().gt(lit(10)))
2025    ///        .agg([col("rain").min().alias("min_rain")])
2026    /// }
2027    /// ```
2028    pub fn having(mut self, predicate: Expr) -> Self {
2029        self.predicates.push(predicate);
2030        self
2031    }
2032
2033    /// Group by and aggregate.
2034    ///
2035    /// Select a column with [col] and choose an aggregation.
2036    /// If you want to aggregate all columns use `col(PlSmallStr::from_static("*"))`.
2037    ///
2038    /// # Example
2039    ///
2040    /// ```rust
2041    /// use polars_core::prelude::*;
2042    /// use polars_lazy::prelude::*;
2043    ///
2044    /// fn example(df: DataFrame) -> LazyFrame {
2045    ///       df.lazy()
2046    ///        .group_by_stable([col("date")])
2047    ///        .agg([
2048    ///            col("rain").min().alias("min_rain"),
2049    ///            col("rain").sum().alias("sum_rain"),
2050    ///            col("rain").quantile(lit(0.5), QuantileMethod::Nearest).alias("median_rain"),
2051    ///        ])
2052    /// }
2053    /// ```
2054    pub fn agg<E: AsRef<[Expr]>>(self, aggs: E) -> LazyFrame {
2055        #[cfg(feature = "dynamic_group_by")]
2056        let lp = DslBuilder::from(self.logical_plan)
2057            .group_by(
2058                self.keys,
2059                self.predicates,
2060                aggs,
2061                None,
2062                self.maintain_order,
2063                self.dynamic_options,
2064                self.rolling_options,
2065            )
2066            .build();
2067
2068        #[cfg(not(feature = "dynamic_group_by"))]
2069        let lp = DslBuilder::from(self.logical_plan)
2070            .group_by(self.keys, self.predicates, aggs, None, self.maintain_order)
2071            .build();
2072        LazyFrame::from_logical_plan(lp, self.opt_state)
2073    }
2074
2075    /// Return first n rows of each group
2076    pub fn head(self, n: Option<usize>) -> LazyFrame {
2077        let keys = self
2078            .keys
2079            .iter()
2080            .filter_map(|expr| expr_output_name(expr).ok())
2081            .collect::<Vec<_>>();
2082
2083        self.agg([all().as_expr().head(n)]).explode_impl(
2084            all() - by_name(keys.iter().cloned(), false, false),
2085            ExplodeOptions {
2086                empty_as_null: true,
2087                keep_nulls: true,
2088            },
2089            true,
2090        )
2091    }
2092
2093    /// Return last n rows of each group
2094    pub fn tail(self, n: Option<usize>) -> LazyFrame {
2095        let keys = self
2096            .keys
2097            .iter()
2098            .filter_map(|expr| expr_output_name(expr).ok())
2099            .collect::<Vec<_>>();
2100
2101        self.agg([all().as_expr().tail(n)]).explode_impl(
2102            all() - by_name(keys.iter().cloned(), false, false),
2103            ExplodeOptions {
2104                empty_as_null: true,
2105                keep_nulls: true,
2106            },
2107            true,
2108        )
2109    }
2110
2111    /// Apply a function over the groups as a new DataFrame.
2112    ///
2113    /// **It is not recommended that you use this as materializing the DataFrame is very
2114    /// expensive.**
2115    pub fn apply(self, f: PlanCallback<DataFrame, DataFrame>, schema: SchemaRef) -> LazyFrame {
2116        if !self.predicates.is_empty() {
2117            panic!("not yet implemented: `apply` cannot be used with `having` predicates");
2118        }
2119
2120        #[cfg(feature = "dynamic_group_by")]
2121        let options = GroupbyOptions {
2122            dynamic: self.dynamic_options,
2123            rolling: self.rolling_options,
2124            slice: None,
2125        };
2126
2127        #[cfg(not(feature = "dynamic_group_by"))]
2128        let options = GroupbyOptions { slice: None };
2129
2130        let lp = DslPlan::GroupBy {
2131            input: Arc::new(self.logical_plan),
2132            keys: self.keys,
2133            predicates: vec![],
2134            aggs: vec![],
2135            apply: Some((f, schema)),
2136            maintain_order: self.maintain_order,
2137            options: Arc::new(options),
2138        };
2139        LazyFrame::from_logical_plan(lp, self.opt_state)
2140    }
2141}
2142
2143#[must_use]
2144pub struct JoinBuilder {
2145    lf: LazyFrame,
2146    how: JoinType,
2147    other: Option<LazyFrame>,
2148    left_on: Vec<Expr>,
2149    right_on: Vec<Expr>,
2150    allow_parallel: bool,
2151    force_parallel: bool,
2152    suffix: Option<PlSmallStr>,
2153    validation: JoinValidation,
2154    nulls_equal: bool,
2155    coalesce: JoinCoalesce,
2156    maintain_order: MaintainOrderJoin,
2157    build_side: Option<JoinBuildSide>,
2158}
2159impl JoinBuilder {
2160    /// Create the `JoinBuilder` with the provided `LazyFrame` as the left table.
2161    pub fn new(lf: LazyFrame) -> Self {
2162        Self {
2163            lf,
2164            other: None,
2165            how: JoinType::Inner,
2166            left_on: vec![],
2167            right_on: vec![],
2168            allow_parallel: true,
2169            force_parallel: false,
2170            suffix: None,
2171            validation: Default::default(),
2172            nulls_equal: false,
2173            coalesce: Default::default(),
2174            maintain_order: Default::default(),
2175            build_side: None,
2176        }
2177    }
2178
2179    /// The right table in the join.
2180    pub fn with(mut self, other: LazyFrame) -> Self {
2181        self.other = Some(other);
2182        self
2183    }
2184
2185    /// Select the join type.
2186    pub fn how(mut self, how: JoinType) -> Self {
2187        self.how = how;
2188        self
2189    }
2190
2191    pub fn validate(mut self, validation: JoinValidation) -> Self {
2192        self.validation = validation;
2193        self
2194    }
2195
2196    /// The expressions you want to join both tables on.
2197    ///
2198    /// The passed expressions must be valid in both `LazyFrame`s in the join.
2199    pub fn on<E: AsRef<[Expr]>>(mut self, on: E) -> Self {
2200        let on = on.as_ref().to_vec();
2201        self.left_on.clone_from(&on);
2202        self.right_on = on;
2203        self
2204    }
2205
2206    /// The expressions you want to join the left table on.
2207    ///
2208    /// The passed expressions must be valid in the left table.
2209    pub fn left_on<E: AsRef<[Expr]>>(mut self, on: E) -> Self {
2210        self.left_on = on.as_ref().to_vec();
2211        self
2212    }
2213
2214    /// The expressions you want to join the right table on.
2215    ///
2216    /// The passed expressions must be valid in the right table.
2217    pub fn right_on<E: AsRef<[Expr]>>(mut self, on: E) -> Self {
2218        self.right_on = on.as_ref().to_vec();
2219        self
2220    }
2221
2222    /// Allow parallel table evaluation.
2223    pub fn allow_parallel(mut self, allow: bool) -> Self {
2224        self.allow_parallel = allow;
2225        self
2226    }
2227
2228    /// Force parallel table evaluation.
2229    pub fn force_parallel(mut self, force: bool) -> Self {
2230        self.force_parallel = force;
2231        self
2232    }
2233
2234    /// Join on null values. By default null values will never produce matches.
2235    pub fn join_nulls(mut self, nulls_equal: bool) -> Self {
2236        self.nulls_equal = nulls_equal;
2237        self
2238    }
2239
2240    /// Suffix to add duplicate column names in join.
2241    /// Defaults to `"_right"` if this method is never called.
2242    pub fn suffix<S>(mut self, suffix: S) -> Self
2243    where
2244        S: Into<PlSmallStr>,
2245    {
2246        self.suffix = Some(suffix.into());
2247        self
2248    }
2249
2250    /// Whether to coalesce join columns.
2251    pub fn coalesce(mut self, coalesce: JoinCoalesce) -> Self {
2252        self.coalesce = coalesce;
2253        self
2254    }
2255
2256    /// Whether to preserve the row order.
2257    pub fn maintain_order(mut self, maintain_order: MaintainOrderJoin) -> Self {
2258        self.maintain_order = maintain_order;
2259        self
2260    }
2261
2262    /// Whether to prefer a specific build side.
2263    pub fn build_side(mut self, build_side: Option<JoinBuildSide>) -> Self {
2264        self.build_side = build_side;
2265        self
2266    }
2267
2268    /// Finish builder
2269    pub fn finish(self) -> LazyFrame {
2270        let opt_state = self.lf.opt_state;
2271        let other = self.other.expect("'with' not set in join builder");
2272
2273        let args = JoinArgs {
2274            how: self.how,
2275            validation: self.validation,
2276            suffix: self.suffix,
2277            slice: None,
2278            nulls_equal: self.nulls_equal,
2279            coalesce: self.coalesce,
2280            maintain_order: self.maintain_order,
2281            build_side: self.build_side,
2282        };
2283
2284        let lp = self
2285            .lf
2286            .get_plan_builder()
2287            .join(
2288                other.logical_plan,
2289                self.left_on,
2290                self.right_on,
2291                JoinOptions {
2292                    allow_parallel: self.allow_parallel,
2293                    force_parallel: self.force_parallel,
2294                    args,
2295                }
2296                .into(),
2297            )
2298            .build();
2299        LazyFrame::from_logical_plan(lp, opt_state)
2300    }
2301
2302    // Finish with join predicates
2303    pub fn join_where(self, predicates: Vec<Expr>) -> LazyFrame {
2304        let opt_state = self.lf.opt_state;
2305        let other = self.other.expect("with not set");
2306
2307        // Decompose `And` conjunctions into their component expressions
2308        fn decompose_and(predicate: Expr, expanded_predicates: &mut Vec<Expr>) {
2309            if let Expr::BinaryExpr {
2310                op: Operator::And,
2311                left,
2312                right,
2313            } = predicate
2314            {
2315                decompose_and((*left).clone(), expanded_predicates);
2316                decompose_and((*right).clone(), expanded_predicates);
2317            } else {
2318                expanded_predicates.push(predicate);
2319            }
2320        }
2321        let mut expanded_predicates = Vec::with_capacity(predicates.len() * 2);
2322        for predicate in predicates {
2323            decompose_and(predicate, &mut expanded_predicates);
2324        }
2325        let predicates: Vec<Expr> = expanded_predicates;
2326
2327        // Decompose `is_between` predicates to allow for cleaner expression of range joins
2328        #[cfg(feature = "is_between")]
2329        let predicates: Vec<Expr> = {
2330            let mut expanded_predicates = Vec::with_capacity(predicates.len() * 2);
2331            for predicate in predicates {
2332                if let Expr::Function {
2333                    function: FunctionExpr::Boolean(BooleanFunction::IsBetween { closed }),
2334                    input,
2335                    ..
2336                } = &predicate
2337                {
2338                    if let [expr, lower, upper] = input.as_slice() {
2339                        match closed {
2340                            ClosedInterval::Both => {
2341                                expanded_predicates.push(expr.clone().gt_eq(lower.clone()));
2342                                expanded_predicates.push(expr.clone().lt_eq(upper.clone()));
2343                            },
2344                            ClosedInterval::Right => {
2345                                expanded_predicates.push(expr.clone().gt(lower.clone()));
2346                                expanded_predicates.push(expr.clone().lt_eq(upper.clone()));
2347                            },
2348                            ClosedInterval::Left => {
2349                                expanded_predicates.push(expr.clone().gt_eq(lower.clone()));
2350                                expanded_predicates.push(expr.clone().lt(upper.clone()));
2351                            },
2352                            ClosedInterval::None => {
2353                                expanded_predicates.push(expr.clone().gt(lower.clone()));
2354                                expanded_predicates.push(expr.clone().lt(upper.clone()));
2355                            },
2356                        }
2357                        continue;
2358                    }
2359                }
2360                expanded_predicates.push(predicate);
2361            }
2362            expanded_predicates
2363        };
2364
2365        let args = JoinArgs {
2366            how: self.how,
2367            validation: self.validation,
2368            suffix: self.suffix,
2369            slice: None,
2370            nulls_equal: self.nulls_equal,
2371            coalesce: self.coalesce,
2372            maintain_order: self.maintain_order,
2373            build_side: self.build_side,
2374        };
2375        let options = JoinOptions {
2376            allow_parallel: self.allow_parallel,
2377            force_parallel: self.force_parallel,
2378            args,
2379        };
2380
2381        let lp = DslPlan::Join {
2382            input_left: Arc::new(self.lf.logical_plan),
2383            input_right: Arc::new(other.logical_plan),
2384            left_on: Default::default(),
2385            right_on: Default::default(),
2386            predicates,
2387            options: Arc::from(options),
2388        };
2389
2390        LazyFrame::from_logical_plan(lp, opt_state)
2391    }
2392}
2393
2394pub const BUILD_STREAMING_EXECUTOR: Option<polars_mem_engine::StreamingExecutorBuilder> = {
2395    #[cfg(not(feature = "streaming"))]
2396    {
2397        None
2398    }
2399    #[cfg(feature = "streaming")]
2400    {
2401        Some(polars_stream::build_streaming_query_executor)
2402    }
2403};
2404
2405pub struct CollectBatches {
2406    recv: Receiver<PolarsResult<DataFrame>>,
2407    runner: Option<Box<dyn FnOnce() + Send + 'static>>,
2408}
2409
2410impl CollectBatches {
2411    /// Start running the query, if not already.
2412    pub fn start(&mut self) {
2413        if let Some(runner) = self.runner.take() {
2414            runner()
2415        }
2416    }
2417}
2418
2419impl Iterator for CollectBatches {
2420    type Item = PolarsResult<DataFrame>;
2421
2422    fn next(&mut self) -> Option<Self::Item> {
2423        self.start();
2424        self.recv.recv().ok()
2425    }
2426}