Method with_column

datafusion/core/src/dataframe/mod.rs:2180–2232

Add or replace a column in the DataFrame. # Example ``` # use datafusion::prelude::*; # use datafusion::error::Result; # #[tokio::main] # async fn main() -> Result<()> { let ctx = SessionContext::new(); let df = ctx .read_csv("tests/data/example.csv", CsvReadOptions::new()) .await?; let df = df.with_column("ab_sum", col("a") + col("b"))?; # Ok(()) # } ```

(self, name: &str, expr: Expr)

    /// # }
    /// ```
    pub fn with_column(self, name: &str, expr: Expr) -> Result<DataFrame> {
        let window_func_exprs = find_window_exprs([&expr]);

        let original_names: HashSet<String> = self
            .plan
            .schema()
            .iter()
            .map(|(_, f)| f.name().clone())
            .collect();

        // Maybe build window plan
        let plan = if window_func_exprs.is_empty() {
            self.plan
        } else {
            LogicalPlanBuilder::window_plan(self.plan, window_func_exprs)?
        };

        let new_column = expr.alias(name);
        let mut col_exists = false;

        let mut fields: Vec<(Expr, bool)> = plan
            .schema()
            .iter()
            .filter_map(|(qualifier, field)| {
                // Skip new fields introduced by window_plan
                if !original_names.contains(field.name()) {
                    return None;
                }

                if field.name() == name {
                    col_exists = true;
                    Some((new_column.clone(), true))
                } else {
                    let e = col(Column::from((qualifier, field)));
                    Some((e, self.projection_requires_validation))
                }
            })
            .collect();

        if !col_exists {
            fields.push((new_column, true));
        }

        let project_plan = LogicalPlanBuilder::from(plan)
            .project_with_validation(fields)?
            .build()?;

        Ok(DataFrame {
            session_state: self.session_state,
            plan: project_plan,
            projection_requires_validation: false,
        })
    }

    /// Rename one column by applying a new projection. This is a no-op if the column to be
    /// renamed does not exist.