如何使用Rust DataFusion实现空值填充与新增计算列

核心API对应说明 #

• 空值填充：直接使用SQL标准表达式coalesce即可，无需查找专门命名为fill_na的方法
• 新增计算列：DataFrame提供with_column方法，和Spark同名API作用完全一致；如果使用低版本无该方法，可通过select枚举全部原有列+新列表达式实现相同效果

具体实现逻辑 #

• b列null值填充
  将b列null替换为空字符串""的表达式为coalesce(col("b"), lit(""))，coalesce会依次判断传入的参数，返回第一个非null的值。如果偏好Case When写法，等价实现如下，两种写法运行效果完全一致：

  when(col("b").is_null(), lit("")).otherwise(col("b"))
  

• 新增a+b组合列
  注意：测试数据中a是整数类型、b是字符串类型，直接执行col("a") + col("b")会触发类型不匹配报错，需要先把a列转换为字符串类型再做运算。 调用with_column时传入新列名和计算表达式即可生成新列，返回的新DataFrame会自动保留所有原有列。

完整可运行代码 #

use datafusion::prelude::*;
use datafusion::error::Result;

#[tokio::main]
async fn main() -> Result<()> {
    let file_path = "datalayers/landing/test_file.json";
    
    let ctx = SessionContext::new();
    let df = ctx.read_json(file_path, NdJsonReadOptions::default()).await?;

    // 填充b列空值为空字符串
    let df_filled = df.with_column(
        "b",
        coalesce(col("b"), lit(""))
    )?;

    // 新增a、b列组合的计算列
    let df_result = df_filled.with_column(
        "ab_comb",
        col("a").cast(arrow::datatypes::DataType::Utf8) + col("b")
    )?;

    df_result.show().await?;
    Ok(())
}

低版本兼容写法 #

如果使用的DataFusion版本没有提供with_column方法，直接用select方法枚举所有原有列+新列表达式即可，效果完全相同：

let df_result = df_filled.select(vec![
    col("a"),
    col("b"),
    col("c"),
    (col("a").cast(arrow::datatypes::DataType::Utf8) + col("b")).alias("ab_comb")
])?;

代码运行后输出结果如下：

+---+----+---+---------+
| a | b  | c | ab_comb |
+---+----+---+---------+
| 1 | hi | 3 | 1hi     |
| 5 |    | 7 | 5       |
+---+----+---+---------+

内容的提问来源于stack exchange，提问作者DataPsycho