cj-zhukov opened a new issue, #26185:
URL: https://github.com/apache/datafusion/issues/26185
### Is your feature request related to a problem or challenge?
There is no way to turn a struct, or a set of columns, into a JSON string
column without leaving the lazy plan. Today this needs `collect()`,
concatenating the batches, serializing with `arrow_json::ArrayWriter`, and
rebuilding the `DataFrame` with `read_batch`, which materializes the whole
table and discards the original plan:
```rust
/// Create json like string column new_col from cols
///
/// # Examples
/// ```
/// # use color_eyre::Result;
/// use datafusion::prelude::*;
/// # use datafusion_example::utils::dataframe::df_cols_to_json;
/// # use datafusion_example::utils::datasets::ExampleDataset;
/// # use datafusion::assert_batches_eq;
/// # #[tokio::main]
/// # async fn main() -> Result<()> {
/// let ctx = SessionContext::new();
/// let cars = ExampleDataset::Cars;
/// let df = cars.dataframe(&ctx).await?;
/// let res = df_cols_to_json(&ctx, df, &["car", "speed"],
"new_col").await?.sort(vec![col("time").sort(true, true)])?;
/// assert_batches_eq!(
/// &[
/// "+---------------------+------------------------------+",
/// "| time | new_col |",
/// "+---------------------+------------------------------+",
/// "| 1996-04-12T12:05:03 | {\"car\":\"red\",\"speed\":20.0} |",
/// "| 1996-04-12T12:05:03 | {\"car\":\"green\",\"speed\":10.0} |",
/// "| 1996-04-12T12:05:04 | {\"car\":\"red\",\"speed\":20.3} |",
/// "| 1996-04-12T12:05:04 | {\"car\":\"green\",\"speed\":10.3} |",
/// "| 1996-04-12T12:05:05 | {\"car\":\"red\",\"speed\":21.4} |",
/// "| 1996-04-12T12:05:05 | {\"car\":\"green\",\"speed\":10.4} |",
/// "| 1996-04-12T12:05:06 | {\"car\":\"red\",\"speed\":21.5} |",
/// "| 1996-04-12T12:05:06 | {\"car\":\"green\",\"speed\":10.5} |",
/// "| 1996-04-12T12:05:07 | {\"car\":\"red\",\"speed\":19.0} |",
/// "| 1996-04-12T12:05:07 | {\"car\":\"green\",\"speed\":11.0} |",
/// "| 1996-04-12T12:05:08 | {\"car\":\"red\",\"speed\":18.0} |",
/// "| 1996-04-12T12:05:08 | {\"car\":\"green\",\"speed\":12.0} |",
/// "| 1996-04-12T12:05:09 | {\"car\":\"red\",\"speed\":17.0} |",
/// "| 1996-04-12T12:05:09 | {\"car\":\"green\",\"speed\":14.0} |",
/// "| 1996-04-12T12:05:10 | {\"car\":\"red\",\"speed\":7.0} |",
/// "| 1996-04-12T12:05:10 | {\"car\":\"green\",\"speed\":15.0} |",
/// "| 1996-04-12T12:05:11 | {\"car\":\"red\",\"speed\":7.1} |",
/// "| 1996-04-12T12:05:11 | {\"car\":\"green\",\"speed\":15.1} |",
/// "| 1996-04-12T12:05:12 | {\"car\":\"red\",\"speed\":7.2} |",
/// "| 1996-04-12T12:05:12 | {\"car\":\"green\",\"speed\":15.2} |",
/// "| 1996-04-12T12:05:13 | {\"car\":\"red\",\"speed\":3.0} |",
/// "| 1996-04-12T12:05:13 | {\"car\":\"green\",\"speed\":8.0} |",
/// "| 1996-04-12T12:05:14 | {\"car\":\"red\",\"speed\":1.0} |",
/// "| 1996-04-12T12:05:14 | {\"car\":\"green\",\"speed\":2.0} |",
/// "| 1996-04-12T12:05:15 | {\"car\":\"red\",\"speed\":0.0} |",
/// "+---------------------+------------------------------+",
/// ],
/// &res.collect().await?
/// );
/// # Ok(())
/// # }
/// ```
pub async fn df_cols_to_json(
ctx: &SessionContext,
df: DataFrame,
cols: &[&str],
new_col: &str,
) -> Result<DataFrame, UtilsError> {
let schema_ref = Arc::new(df.schema().as_arrow().clone());
let arrays = concat_arrays(df).await?;
let batch = RecordBatch::try_new(schema_ref.clone(), arrays)?;
let selected_arrays: Vec<ArrayRef> = cols
.iter()
.map(|col| {
batch
.column_by_name(col)
.ok_or_else(|| DataFusionError::Plan(format!("column {col}
not found")))
.cloned()
})
.collect::<Result<_, _>>()?;
let selected_fields: Vec<Field> = cols
.iter()
.map(|col| schema_ref.field_with_name(col).cloned())
.collect::<Result<_, _>>()?;
let selected_schema = Arc::new(Schema::new(selected_fields));
let selected_batch = RecordBatch::try_new(selected_schema.clone(),
selected_arrays)?;
let mut writer = arrow_json::ArrayWriter::new(Vec::new());
writer.write(&selected_batch)?;
writer.finish()?;
let json_data = writer.into_inner();
let json_rows: Vec<Map<String, Value>> =
serde_json::from_reader(json_data.as_slice())?;
let str_rows: Vec<String> = json_rows
.iter()
.map(serde_json::to_string)
.collect::<Result<_, _>>()?;
let mut all_columns = batch.columns().to_vec();
all_columns.push(Arc::new(StringArray::from(str_rows)) as ArrayRef);
let mut new_fields: Vec<Field> = schema_ref
.fields()
.iter()
.map(|f| f.as_ref().clone())
.collect();
new_fields.push(Field::new(new_col, DataType::Utf8, true));
let new_schema = Arc::new(Schema::new(new_fields));
let new_batch = RecordBatch::try_new(new_schema, all_columns)?;
let res = ctx.read_batch(new_batch)?.drop_columns(cols)?;
Ok(res)
}
```
Other engines expose this as an ordinary expression.
Related #21301
### Describe the solution you'd like
A Spark-compatible `to_json` scalar function in `datafusion-spark`, plus a
small Expr builder so it can be used with `with_column` either on an existing
struct column or on a list of column names:
```rust
use datafusion::prelude::*;
use datafusion_spark::expr_fn::to_json;
let df = dataframe!(
"id" => [10, 20, 30],
"foo" => [1, 2, 3],
"bar" => [None, Some("bak"), Some("baz")],
)?;
// pack columns into a struct and serialize it
let df = df.with_column("json", to_json(&["foo", "bar"]))?;
// or serialize an existing struct column
// let df = df.with_column("json", to_json(col("event")))?;
df.show().await?;
+----+-----+-----+-----------------------+
| id | foo | bar | json |
+----+-----+-----+-----------------------+
| 10 | 1 | | {"foo":1} |
| 20 | 2 | bak | {"foo":2,"bar":"bak"} |
| 30 | 3 | baz | {"foo":3,"bar":"baz"} |
+----+-----+-----+-----------------------+
```
The same function from SQL:
```sql
SELECT id, to_json(named_struct('foo', foo, 'bar', bar)) AS json FROM t;
```
Semantics follow Spark’s defaults: NULL fields are omitted
(`ignoreNullFields = true`), a NULL struct produces NULL, nested structs/lists
are serialized recursively. `to_json(&[...])` is plain sugar for
`to_json(named_struct(...))`, so `with_column` needs no changes and the result
stays a lazy projection (serialization runs per batch with the arrow-json row
encoder, no collect).
### Describe alternatives you've considered
Keeping the collect()-based helper in user code: works, but executes the
plan and loses laziness, filter push down and the original lineage - see
example above.
### Additional context
Spark `to_json`:
https://spark.apache.org/docs/latest/api/sql/index.html#to_json
Polars `struct.json_encode`:
https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.struct.json_encode.html
### Willingness to contribute
I plan to contribute to this feature independently
--
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]
