Add a column of values to a DataFrame

[cj-zhukov]

Is your feature request related to a problem or challenge?

I want to append an Arrow column (ArrayRef) to an existing DataFrame.

DataFrame::with_column only accepts an Expr. That works for derived columns such as col("a") + col("b"), but not for attaching values I already have (for example ["foo", "bar", "baz"] next to existing data).

The workaround is a helper that executes the frame, concatenates its columns, pushes the new array, and builds a new DataFrame with read_batch:

pub async fn add_column_to_df(
    ctx: &SessionContext,
    df: DataFrame,
    data: ArrayRef,
    col_name: &str,
) -> Result<DataFrame> {
    let schema = df.schema().as_arrow().clone();
    let mut arrays = concat_arrays(df).await?;
    let row_count = arrays
        .first()
        .ok_or_else(|| DataFusionError::Execution("Empty DataFrame".into()))?
        .len();

    if data.len() != row_count {
        return Err(DataFusionError::Execution(format!(
            "Column '{col_name}' has length {}, expected {row_count}",
            data.len()
        ))
        .into());
    }

    let new_col_type = data.data_type().clone();
    arrays.push(data);

    make_new_df(ctx, arrays, &Arc::new(schema), col_name, &new_col_type)
}

fn make_new_df(
    ctx: &SessionContext,
    arrays: Vec<ArrayRef>,
    old_schema: &SchemaRef,
    col_name: &str,
    new_col_type: &DataType,
) -> Result<DataFrame> {
    let mut new_fields: Vec<Field> = old_schema
        .fields()
        .iter()
        .map(|f| f.as_ref().clone())
        .collect();
    new_fields.push(Field::new(col_name, new_col_type.clone(), true));
    let new_schema = Arc::new(Schema::new(new_fields));
    let batch = RecordBatch::try_new(new_schema, arrays)?;
    let df = ctx.read_batch(batch)?;
    Ok(df)
}

pub async fn concat_arrays(df: DataFrame) -> Result<Vec<ArrayRef>> {
    let schema = df.schema().clone();
    let batches = df.collect().await?;
    let batches = batches.iter().collect::<Vec<_>>();
    let field_num = schema.fields().len();
    let mut arrays = Vec::with_capacity(field_num);
    for i in 0..field_num {
        let array = batches
            .iter()
            .map(|batch| batch.column(i).as_ref())
            .collect::<Vec<_>>();
        let array = concat(&array)?;
        arrays.push(array);
    }
    Ok(arrays)
}

But this requires a collect(). After that you can no longer keep planning (filter, select, and so on) on the original lazy plan.

Describe the solution you'd like

A DataFrame API that appends a column from an ArrayRef without executing the current plan.

The new column should be stored on the plan and applied when the new DataFrame is collected.

let df = DataFrame::from_columns([("id", id), ("data", data)])?;

// does not collect
let df = df.with_array_columns([("new_col", new_col)])?;

let df = df.filter(col("id").gt(lit(1)))?;
df.collect().await?;

The important part is: append column data without collecting at the call site.

Describe alternatives you've considered

• Collect and rebuild (the helper above). It works, but it executes the DataFrame only to add a column.
• Join a second DataFrame that holds the new column (on a key, or on row_number()). This stays lazy and uses public APIs, but it is needs a key or a positional index.
• with_column(Expr), a list literal, or unnest. These add a computed or broadcast value, not row i of an ArrayRef.

Additional context

There was a discussion about a Polars-style feature in #9672. That issue was closed because the only approach considered at the time was to collect the DataFrame and append the column to the resulting RecordBatch.

[cj-zhukov]

take