diff --git a/src/dataframe.rs b/src/dataframe.rs index 46fba137c..508ad6d59 100644 --- a/src/dataframe.rs +++ b/src/dataframe.rs @@ -19,6 +19,7 @@ use std::collections::HashMap; use std::ffi::CString; use std::sync::Arc; +use crate::datafusion_expr::Expr; use arrow::array::{new_null_array, RecordBatch, RecordBatchIterator, RecordBatchReader}; use arrow::compute::can_cast_types; use arrow::error::ArrowError; @@ -473,7 +474,19 @@ impl PyDataFrame { } fn with_column(&self, name: &str, expr: PyExpr) -> PyDataFusionResult { - let df = self.df.as_ref().clone().with_column(name, expr.into())?; + let expr: Expr = expr.into(); + let aliased = expr.alias(name); + + let df_schema = self.df.as_ref().schema().clone(); + let mut proj_exprs: Vec = df_schema + .fields() + .iter() + .filter(|field| field.name() != name) + .map(|field| col(field.name())) + .collect(); + proj_exprs.push(aliased); + + let df = self.df.as_ref().clone().select(proj_exprs)?; Ok(Self::new(df)) }