cj-zhukov opened a new issue, #85:
URL: https://github.com/apache/datafusion-spark/issues/85
_Moved from apache/datafusion#26185, opened before the `datafusion-spark`
crate was moved to this repository. The implementation PR
apache/datafusion#26187 (reviewed there, not merged before the move) is being
ported here._
### 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)
}
```
Spark exposes this as an ordinary expression, `to_json(struct)`. DataFusion
already has the first half (`named_struct`, `struct`) but no serializer.
### Describe the solution you'd like
A Spark-compatible `to_json` scalar function in this crate, 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 struct fields are omitted
(`ignoreNullFields = true`) while NULL map entries are kept, map keys are
written as strings, binary is Base64, non-finite floats are the strings `"NaN"`
/ `"Infinity"` / `"-Infinity"`, a NULL struct produces NULL, and nested
structs, lists and maps 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`).
Not in scope for a first version: the Spark `options` argument
(`ignoreNullFields`, `timestampFormat`) and map/array top-level inputs.
### Describe alternatives you've considered
- Keeping the `collect()`-based helper in user code: works, but executes the
plan and loses laziness, filter pushdown and the original lineage.
- `arrow_cast(struct_col, 'Utf8')`: produces the Arrow display format (`{a:
1, b: x}`), not JSON.
### Additional context
- Spark `to_json`:
https://spark.apache.org/docs/latest/api/sql/index.html#to_json
- Original issue and review history: apache/datafusion#26185,
apache/datafusion#26187
### 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]
