cj-zhukov opened a new issue, #25298:
URL: https://github.com/apache/datafusion/issues/25298
### 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`:
```rust
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.
```rust
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
https://github.com/apache/datafusion/issues/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`.
--
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]
