polars_lazy/physical_plan/
exotic.rs

1use polars_core::prelude::*;
2use polars_expr::{ExpressionConversionState, create_physical_expr};
3
4use crate::prelude::*;
5
6#[cfg(feature = "pivot")]
7pub(crate) fn prepare_eval_expr(expr: Expr) -> Expr {
8    expr.map_expr(|e| match e {
9        Expr::Column(_) => Expr::Column(PlSmallStr::EMPTY),
10        Expr::Nth(_) => Expr::Column(PlSmallStr::EMPTY),
11        e => e,
12    })
13}
14
15pub(crate) fn prepare_expression_for_context(
16    name: PlSmallStr,
17    expr: &Expr,
18    dtype: &DataType,
19    ctxt: Context,
20) -> PolarsResult<Arc<dyn PhysicalExpr>> {
21    let mut lp_arena = Arena::with_capacity(8);
22    let mut expr_arena = Arena::with_capacity(10);
23
24    // create a dummy lazyframe and run a very simple optimization run so that
25    // type coercion and simplify expression optimizations run.
26    let column = Series::full_null(name, 0, dtype);
27    let df = column.into_frame();
28    let input_schema = df.schema().clone();
29    let lf = df
30        .lazy()
31        .without_optimizations()
32        .with_simplify_expr(true)
33        .select([expr.clone()]);
34    let optimized = lf.optimize(&mut lp_arena, &mut expr_arena)?;
35    let lp = lp_arena.get(optimized);
36    let aexpr = lp
37        .get_exprs()
38        .pop()
39        .ok_or_else(|| polars_err!(ComputeError: "expected expressions in the context"))?;
40
41    create_physical_expr(
42        &aexpr,
43        ctxt,
44        &expr_arena,
45        &input_schema,
46        &mut ExpressionConversionState::new(true),
47    )
48}