refactor: simplify with_column method by using with_column directly

Author: kosiew

python/datafusion/context.py

@@ -731,11 +731,12 @@ def from_polars(self, data: pl.DataFrame, name: str | None = None) -> DataFrame:
         """
         return DataFrame(self.ctx.from_polars(data, name))
 
-    # https://github.com/apache/datafusion-python/pull/1016#discussion_r1983239116
-    # is the discussion on how we arrived at adding register_view
     def register_view(self, name: str, df: DataFrame) -> None:
         """Register a :py:class: `~datafusion.detaframe.DataFrame` as a view.
 
+        The DataFrame is converted into a view before registration so it can be
+        referenced from SQL without materializing results immediately.
+
         Args:
             name (str): The name to register the view under.
             df (DataFrame): The DataFrame to be converted into a view and registered.

src/dataframe.rs

@@ -19,7 +19,6 @@ 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;
@@ -474,19 +473,7 @@ impl PyDataFrame {
     }
 
     fn with_column(&self, name: &str, expr: PyExpr) -> PyDataFusionResult<Self> {
-        let expr: Expr = expr.into();
-        let aliased = expr.alias(name);
-
-        let df_schema = self.df.as_ref().schema().clone();
-        let mut proj_exprs: Vec<Expr> = 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)?;
+        let df = self.df.as_ref().clone().with_column(name, expr.into())?;
         Ok(Self::new(df))
     }