05-polars Expressions and contexts

class Expr:
    """Expressions that can be used in various contexts."""

    # NOTE: This `= None` is needed to generate the docs with sphinx_accessor.
    _pyexpr: PyExpr = None  # type: ignore[assignment]

    @property
    def bin(self) -> ExprBinaryNameSpace:
        """
        Create an object namespace of all binary related methods.

        See the individual method pages for full details
        """
        return ExprBinaryNameSpace(self)

    def __add__(self, other: IntoExpr) -> Expr:
        other_pyexpr = parse_into_expression(other, str_as_lit=True)
        return wrap_expr(self._pyexpr + other_pyexpr)

    def __eq__(self, other: IntoExpr) -> Expr:  # type: ignore[override]
        warn_null_comparison(other)
        other_pyexpr = parse_into_expression(other, str_as_lit=True)
        return wrap_expr(self._pyexpr.eq(other_pyexpr))

    def __floordiv__(self, other: IntoExpr) -> Expr:
        other_pyexpr = parse_into_expression(other)
        return wrap_expr(self._pyexpr // other_pyexpr)

def parse_into_expression(
    input: IntoExpr,
    *,
    str_as_lit: bool = False,
    list_as_series: bool = False,
    structify: bool = False,
    dtype: PolarsDataType | None = None,
    require_selector: bool = False,
) -> PyExpr:

    if isinstance(input, pl.Expr):
        expr = input
        if structify:
            expr = _structify_expression(expr)
    elif isinstance(input, str) and not str_as_lit:
        expr = F.col(input)
    else:
        if require_selector:
            msg = f"cannot turn {qualified_type_name(input)!r} into selector"
            raise TypeError(msg)
        elif isinstance(input, list) and list_as_series:
            expr = F.lit(pl.Series(input), dtype=dtype)
        else:
            expr = F.lit(input, dtype=dtype)

#[pyclass(from_py_object)] // Not marked as frozen for pickling, but that's the only &mut self method.
#[repr(transparent)]
#[derive(Clone)]
pub struct PyExpr {
    pub inner: Expr,
}

impl From<Expr> for PyExpr {
    fn from(expr: Expr) -> Self {
        PyExpr { inner: expr }
    }
}

#[pyfunction]
pub fn coalesce(exprs: Vec<PyExpr>) -> PyExpr {
    let exprs = exprs.to_exprs();
    dsl::coalesce(&exprs).into()
}

#[pyfunction]
pub fn col(name: &str) -> PyExpr {
    dsl::col(name).into()
}

#[pyfunction]
pub fn element() -> PyExpr {
    dsl::element().into()
}

/// Expressions that can be used in various contexts.
///
/// Queries consist of multiple expressions.
/// When using the polars lazy API, don't construct an `Expr` directly; instead, create one using
/// the functions in the `polars_lazy::dsl` module. See that module's docs for more info.
#[derive(Clone, PartialEq)]
#[must_use]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
pub enum Expr {
    /// Values in a `eval` context.
    ///
    /// Equivalent of `pl.element()`.
    Element,
    Alias(Arc<Expr>, PlSmallStr),
    Column(PlSmallStr),
    Selector(Selector),
    Literal(LiteralValue),
    DataTypeFunction(DataTypeFunction),
    BinaryExpr {

    Cast {
        expr: Arc<Expr>,
        dtype: DataTypeExpr,
        options: CastOptions,
    },
    Sort {
        expr: Arc<Expr>,
        options: SortOptions,
    },
    Gather {
        expr: Arc<Expr>,
        idx: Arc<Expr>,
        returns_scalar: bool,
        null_on_oob: bool,
    },
    SortBy {
        expr: Arc<Expr>,
        by: Vec<Expr>,
        sort_options: SortMultipleOptions,
    },
    Agg(AggExpr),
    /// A ternary operation
    /// if true then "foo" else "bar"
    Ternary {
        predicate: Arc<Expr>,
        truthy: Arc<Expr>,
        falsy: Arc<Expr>,
    },
    Function {
        /// function arguments
        input: Vec<Expr>,
        /// function to apply
        function: FunctionExpr,
    },
    Explode {

    // Polars' `select` operation. This may access full materialized data.
    Select {
        input: Node,
        expr: Vec<ExprIR>,
        schema: SchemaRef,
        options: ProjectionOptions,
    },

    },
    GroupBy {
        input: Node,
        keys: Vec<ExprIR>,
        aggs: Vec<ExprIR>,
        schema: SchemaRef,
        maintain_order: bool,
        options: Arc<GroupbyOptions>,
        apply: Option<PlanCallback<DataFrame, DataFrame>>,
    },

    HStack {
        input: Node,
        exprs: Vec<ExprIR>,
        schema: SchemaRef,
        options: ProjectionOptions,
    },
    Distinct {

// Executor are the executors of the physical plan and produce DataFrames. They
// combine physical expressions, which produce Series.

/// Executors will evaluate physical expressions and collect them in a DataFrame.
///
/// Executors have other executors as input. By having a tree of executors we can execute the
/// physical plan until the last executor is evaluated.
pub trait Executor: Send + Sync {
    fn execute(&mut self, cache: &mut ExecutionState) -> PolarsResult<DataFrame>;

    fn is_cache_prefiller(&self) -> bool {
        false
    }
}

pub trait PhysicalExpr: Send + Sync {
    fn as_expression(&self) -> Option<&Expr> {
        None
    }

    fn as_column(&self) -> Option<PlSmallStr> {
        None
    }

    /// Take a DataFrame and evaluate the expression.
    ///
    /// Note: implementers should implement evaluate_impl instead, as this wraps
    /// that call with an error context.
    fn evaluate(&self, df: &DataFrame, state: &ExecutionState) -> PolarsResult<Column> {
        self.evaluate_impl(df, state).map_err(|e| {
            if let Some(expr) = self.as_expression() {
                e.with_expr_context(expr.to_string().into())
            } else {
                e
            }
        })
    }

    fn evaluate_impl(&self, df: &DataFrame, _state: &ExecutionState) -> PolarsResult<Column>;

fn run_exprs_par(
    df: &DataFrame,
    exprs: &[Arc<dyn PhysicalExpr>],
    state: &ExecutionState,
) -> PolarsResult<Vec<Column>> {
    RAYON.install(|| {
        exprs
            .par_iter()
            .map(|expr| expr.evaluate(df, state))
            .collect()
    })
}

fn run_exprs_seq(
    df: &DataFrame,
    exprs: &[Arc<dyn PhysicalExpr>],
    state: &ExecutionState,
) -> PolarsResult<Vec<Column>> {
    exprs.iter().map(|expr| expr.evaluate(df, state)).collect()
}

pub(super) fn evaluate_physical_expressions(
    df: &mut DataFrame,
    exprs: &[Arc<dyn PhysicalExpr>],
    state: &ExecutionState,
    has_windows: bool,
    run_parallel: bool,
) -> PolarsResult<Vec<Column>> {
    let expr_runner = if has_windows {
        execute_projection_cached_window_fns
    } else if run_parallel && exprs.len() > 1 {
        run_exprs_par
    } else {
        run_exprs_seq
    };

    let selected_columns = expr_runner(df, exprs, state)?;

    if has_windows {
        state.clear_window_expr_cache();
    }

    Ok(selected_columns)
}

    fn execute_hor(
        &mut self,
        df: DataFrame,
        state: &mut ExecutionState,
    ) -> PolarsResult<DataFrame> {
        if self.has_window {
            state.insert_has_window_function_flag()
        }
        let c = self.predicate.evaluate(&df, state)?;
        if self.has_window {
            state.clear_window_expr_cache()
        }

        // @scalar-opt
        // @partition-opt
        df.filter(column_to_mask(&c, df.height())?.as_ref())
    }

pub(super) fn evaluate_aggs(
    df: &DataFrame,
    aggs: &[Arc<dyn PhysicalExpr>],
    groups: &GroupPositions,
    state: &ExecutionState,
) -> PolarsResult<Vec<Column>> {
    RAYON.install(|| {
        aggs.par_iter()
            .map(|expr| {
                let agg = expr.evaluate_on_groups(df, groups, state)?.finalize();
                polars_ensure!(agg.len() == groups.len(), agg_len = agg.len(), groups.len());
                Ok(agg)
            })
            .collect::<PolarsResult<Vec<_>>>()
    })
}

//! this contains code used for rewriting projections, expanding wildcards, regex selection etc.

use super::*;
use crate::constants::{
    POLARS_ELEMENT, POLARS_STRUCTFIELDS, get_pl_element_name, get_pl_structfields_name,
};

pub fn prepare_projection(
    exprs: Vec<Expr>,
    schema: &Schema,
    opt_flags: &mut OptFlags,
) -> PolarsResult<(Vec<Expr>, Schema)> {
    let exprs = rewrite_projections(exprs, &PlIndexSet::new(), schema, opt_flags)?;
    let schema = expressions_to_schema(&exprs, schema, |duplicate_name: &str| {
        format!("projections contained duplicate output name '{duplicate_name}'")
    })?;
    Ok((exprs, schema))
}

pub fn expand_expression(
    expr: &Expr,
    ignored_selector_columns: &PlIndexSet<PlSmallStr>,
    schema: &Schema,
    out: &mut Vec<Expr>,
    opt_flags: &mut OptFlags,
) -> PolarsResult<()> {
    if expr.into_iter().all(|e| !needs_expansion(e)) {
        out.push(expr.clone());
        return Ok(());
    }

    expand_expression_rec(expr, ignored_selector_columns, schema, out, opt_flags)?;
    Ok(())
}

}