Add DataFrame fill_null

datafusion/core/src/dataframe/mod.rs

@@ -51,13 +51,17 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use datafusion_common::config::{CsvOptions, JsonOptions};
 use datafusion_common::{
     exec_err, not_impl_err, plan_err, Column, DFSchema, DataFusionError, ParamValues,
-    SchemaError, UnnestOptions,
+    ScalarValue, SchemaError, UnnestOptions,
 };
-use datafusion_expr::dml::InsertOp;
-use datafusion_expr::{case, is_null, lit, SortExpr};
 use datafusion_expr::{
-    utils::COUNT_STAR_EXPANSION, TableProviderFilterPushDown, UNNAMED_TABLE,
+    case,
+    dml::InsertOp,
+    expr::{Alias, ScalarFunction},
+    is_null, lit,
+    utils::COUNT_STAR_EXPANSION,
+    SortExpr, TableProviderFilterPushDown, UNNAMED_TABLE,
 };
+use datafusion_functions::core::coalesce;
 use datafusion_functions_aggregate::expr_fn::{
     avg, count, max, median, min, stddev, sum,
 };
@@ -1926,6 +1930,71 @@ impl DataFrame {
             plan,
         })
     }
+
+    /// Fill null values in specified columns with a given value
+    /// If no columns are specified, applies to all columns
+    /// Only fills if the value can be cast to the column's type
+    ///
+    /// # Arguments
+    /// * `value` - Value to fill nulls with
+    /// * `columns` - Optional list of column names to fill. If None, fills all columns
+    pub fn fill_null(
+        &self,
+        value: ScalarValue,
+        columns: Option<Vec<String>>,
+    ) -> Result<DataFrame> {
+        let cols = match columns {
+            Some(names) => self.find_columns(&names)?,
+            None => self
+                .logical_plan()
+                .schema()
+                .fields()
+                .iter()
+                .map(|f| f.as_ref().clone())
+                .collect(),
+        };
+
+        // Create projections for each column
+        let projections = self
+            .logical_plan()
+            .schema()
+            .fields()
+            .iter()
+            .map(|field| {
+                if cols.contains(field) {
+                    // Try to cast fill value to column type. If the cast fails, fallback to the original column.
+                    match value.clone().cast_to(field.data_type()) {
+                        Ok(fill_value) => Expr::Alias(Alias {
+                            expr: Box::new(Expr::ScalarFunction(ScalarFunction {
+                                func: coalesce(),
+                                args: vec![col(field.name()), lit(fill_value)],
+                            })),
+                            relation: None,
+                            name: field.name().to_string(),
+                        }),
+                        Err(_) => col(field.name()),
+                    }
+                } else {
+                    col(field.name())
+                }
+            })
+            .collect::<Vec<_>>();
+
+        self.clone().select(projections)
+    }
+
+    // Helper to find columns from names
+    fn find_columns(&self, names: &[String]) -> Result<Vec<Field>> {
+        let schema = self.logical_plan().schema();
+        names
+            .iter()
+            .map(|name| {
+                schema.field_with_name(None, name).cloned().map_err(|_| {
+                    DataFusionError::Plan(format!("Column '{}' not found", name))
+                })
+            })
+            .collect()
+    }
 }
 
 #[derive(Debug)]

datafusion/core/tests/dataframe/mod.rs

@@ -5342,3 +5342,96 @@ async fn test_insert_into_checking() -> Result<()> {
 
     Ok(())
 }
+
+async fn create_null_table() -> Result<DataFrame> {
+    // create a DataFrame with null values
+    //    "+---+----+",
+    //    "| a | b |",
+    //    "+---+---+",
+    //    "| 1 | x |",
+    //    "|   |   |",
+    //    "| 3 | z |",
+    //    "+---+---+",
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("a", DataType::Int32, true),
+        Field::new("b", DataType::Utf8, true),
+    ]));
+    let a_values = Int32Array::from(vec![Some(1), None, Some(3)]);
+    let b_values = StringArray::from(vec![Some("x"), None, Some("z")]);
+    let batch = RecordBatch::try_new(
+        schema.clone(),
+        vec![Arc::new(a_values), Arc::new(b_values)],
+    )?;
+
+    let ctx = SessionContext::new();
+    let table = MemTable::try_new(schema.clone(), vec![vec![batch]])?;
+    ctx.register_table("t_null", Arc::new(table))?;
+    let df = ctx.table("t_null").await?;
+    Ok(df)
+}
+
+#[tokio::test]
+async fn test_fill_null() -> Result<()> {
+    let df = create_null_table().await?;
+
+    // Use fill_null to replace nulls on each column.
+    let df_filled = df
+        .fill_null(ScalarValue::Int32(Some(0)), Some(vec!["a".to_string()]))?
+        .fill_null(
+            ScalarValue::Utf8(Some("default".to_string())),
+            Some(vec!["b".to_string()]),
+        )?;
+
+    let results = df_filled.collect().await?;
+    let expected = [
+        "+---+---------+",
+        "| a | b       |",
+        "+---+---------+",
+        "| 1 | x       |",
+        "| 0 | default |",
+        "| 3 | z       |",
+        "+---+---------+",
+    ];
+    assert_batches_sorted_eq!(expected, &results);
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_fill_null_all_columns() -> Result<()> {
+    let df = create_null_table().await?;
+
+    // Use fill_null to replace nulls on all columns.
+    // Only column "b" will be replaced since ScalarValue::Utf8(Some("default".to_string()))
+    // can be cast to Utf8.
+    let df_filled = df.fill_null(ScalarValue::Utf8(Some("default".to_string())), None)?;
+
+    let results = df_filled.clone().collect().await?;
+
+    let expected = [
+        "+---+---------+",
+        "| a | b       |",
+        "+---+---------+",
+        "| 1 | x       |",
+        "|   | default |",
+        "| 3 | z       |",
+        "+---+---------+",
+    ];
+
+    assert_batches_sorted_eq!(expected, &results);
+
+    // Fill column "a" null values with a value that cannot be cast to Int32.
+    let df_filled = df_filled.fill_null(ScalarValue::Int32(Some(0)), None)?;
+
+    let results = df_filled.collect().await?;
+    let expected = [
+        "+---+---------+",
+        "| a | b       |",
+        "+---+---------+",
+        "| 1 | x       |",
+        "| 0 | default |",
+        "| 3 | z       |",
+        "+---+---------+",
+    ];
+    assert_batches_sorted_eq!(expected, &results);
+    Ok(())
+}