kosiew commented on code in PR #22702:
URL: https://github.com/apache/datafusion/pull/22702#discussion_r3340603906
##########
datafusion/core/tests/dataframe/mod.rs:
##########
@@ -6539,6 +6539,150 @@ async fn test_fill_null_all_columns() -> Result<()> {
Ok(())
}
+async fn create_nan_table() -> Result<DataFrame> {
+ // create a DataFrame with a NaN value in a float column "a" and a
+ // non-float column "b" that must stay untouched by fill_nan.
+ // "+-----+---+",
+ // "| a | b |",
+ // "+-----+---+",
+ // "| 1.0 | 1 |",
+ // "| NaN | 2 |",
+ // "| 3.0 | 3 |",
+ // "+-----+---+",
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Float64, true),
+ Field::new("b", DataType::Int32, true),
+ ]));
+ let a_values = Float64Array::from(vec![Some(1.0), Some(f64::NAN),
Some(3.0)]);
+ let b_values = Int32Array::from(vec![Some(1), Some(2), Some(3)]);
+ 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_nan", Arc::new(table))?;
+ let df = ctx.table("t_nan").await?;
+ Ok(df)
+}
+
+#[tokio::test]
+async fn test_fill_nan() -> Result<()> {
+ let df = create_nan_table().await?;
+
+ // Fill NaNs in the float column "a" with 0.0.
+ let df_filled =
+ df.fill_nan(ScalarValue::Float64(Some(0.0)), vec!["a".to_string()])?;
Review Comment:
Nice test coverage around the happy path and cast failure case. One thing
that seems worth locking in is the documented behavior that replacements are
applied when the provided value can be cast to the target float type. Right now
the positive tests use an already-Float64 value, while the negative test covers
an uncastable value. Would it make sense to add a positive cast case as well,
for example `df.fill_nan(ScalarValue::Int32(Some(0)), vec!["a".to_string()])?`?
That would help ensure cross-type casting continues to work as intended.
##########
datafusion/core/src/dataframe/mod.rs:
##########
@@ -2527,6 +2528,78 @@ impl DataFrame {
.collect()
}
+ /// Fill NaN values in specified columns with a given value
+ /// If no columns are specified (empty vector), applies to all columns
+ /// Only fills if the value can be cast to the column's type
+ ///
+ /// # Arguments
+ /// * `value` - Value to fill NaNs with
+ /// * `columns` - List of column names to fill. If empty, fills all
columns.
+ ///
+ /// # Example
+ /// ```
+ /// # use datafusion::prelude::*;
+ /// # use datafusion::error::Result;
+ /// # use datafusion_common::ScalarValue;
+ /// # #[tokio::main]
+ /// # async fn main() -> Result<()> {
+ /// let ctx = SessionContext::new();
+ /// let df = ctx
+ /// .read_csv("tests/data/example.csv", CsvReadOptions::new())
+ /// .await?;
+ /// // Fill NaN in only columns "a" and "c":
+ /// let df = df.fill_nan(ScalarValue::from(0.0), vec!["a".to_owned(),
"c".to_owned()])?;
+ /// // Fill NaN across all columns:
+ /// let df = df.fill_nan(ScalarValue::from(0.0), vec![])?;
+ /// # Ok(())
+ /// # }
+ /// ```
+ #[expect(clippy::needless_pass_by_value)]
+ pub fn fill_nan(
+ &self,
+ value: ScalarValue,
+ columns: Vec<String>,
+ ) -> Result<DataFrame> {
+ let cols = if columns.is_empty() {
Review Comment:
I noticed `fill_nan` and `fill_null` share quite a bit of structure around
column selection, cast handling, alias construction, and the fallback behavior
when casting is not possible. It might be worth extracting a small private
helper that takes the scalar function and applicability check. That could keep
the no-op-on-cast-failure behavior defined in one place and make it a little
easier to keep the two APIs in sync over time.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
