Added polars struct-encode-json, providing the ability to encode structs as json (#15678)

# Description
This PR introduces `polars struct-encode-json`. This exposes the ability
to encode struct columns as json strings. This is useful when converting
things to formats like CSV that do not support complex types.

```nushell
> ❯ : [[id person]; [1 {name: "Bob", age: 36}] [2 {name: "Betty", age: 63}]]
                    | polars into-df -s {id: i64, person: {name: str, age: u8}}
                    | polars select id (polars col person | polars struct-json-encode | polars as encoded) 
                    | polars collect
╭───┬────┬───────────────────────────╮
│ # │ id │          encoded          │
├───┼────┼───────────────────────────┤
│ 0 │  1 │ {"age":36,"name":"Bob"}   │
│ 1 │  2 │ {"age":63,"name":"Betty"} │
╰───┴────┴───────────────────────────╯
```

# User-Facing Changes
* Added `polars struct-encode-json`, providing the ability to encode
structs as json
This commit is contained in:
Jack Wright 2025-05-06 13:58:51 -07:00 committed by GitHub
parent ce308ee461
commit ff8831318d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 109 additions and 2 deletions

View File

@ -242,8 +242,46 @@ impl PluginTest {
// Check for equality with the result
if !self.value_eq(expectation, &value)? {
// If they're not equal, print a diff of the debug format
let expectation_formatted = format!("{:#?}", expectation);
let value_formatted = format!("{:#?}", value);
let (expectation_formatted, value_formatted) =
match (expectation, &value) {
(
Value::Custom { val: ex_val, .. },
Value::Custom { val: v_val, .. },
) => {
// We have to serialize both custom values before handing them to the plugin
let expectation_serialized =
PluginCustomValue::serialize_from_custom_value(
ex_val.as_ref(),
expectation.span(),
)?
.with_source(self.source.clone());
let value_serialized =
PluginCustomValue::serialize_from_custom_value(
v_val.as_ref(),
expectation.span(),
)?
.with_source(self.source.clone());
let persistent =
self.source.persistent(None)?.get_plugin(None)?;
let expectation_base = persistent
.custom_value_to_base_value(
expectation_serialized
.into_spanned(expectation.span()),
)?;
let value_base = persistent.custom_value_to_base_value(
value_serialized.into_spanned(value.span()),
)?;
(
format!("{:#?}", expectation_base),
format!("{:#?}", value_base),
)
}
_ => (format!("{:#?}", expectation), format!("{:#?}", value)),
};
let diff = diff_by_line(&expectation_formatted, &value_formatted);
failed_header();
eprintln!("{} {}", bold.paint("Result:"), diff);

View File

@ -34,6 +34,7 @@ mod slice;
mod sort_by_expr;
pub mod sql_context;
pub mod sql_expr;
mod struct_json_encode;
mod take;
mod unnest;
mod unpivot;
@ -114,6 +115,7 @@ pub(crate) fn data_commands() -> Vec<Box<dyn PluginCommand<Plugin = PolarsPlugin
Box::new(LazySortBy),
Box::new(LazyFilter),
Box::new(Shift),
Box::new(struct_json_encode::StructJsonEncode),
Box::new(qcut::QCutSeries),
Box::new(Unique),
Box::new(unnest::UnnestDF),

View File

@ -0,0 +1,67 @@
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
use nu_protocol::{Category, Example, LabeledError, PipelineData, Signature, Span, Type};
use polars::df;
use crate::{
values::{CustomValueSupport, NuDataFrame, NuExpression},
PolarsPlugin,
};
#[derive(Clone)]
pub struct StructJsonEncode;
impl PluginCommand for StructJsonEncode {
type Plugin = PolarsPlugin;
fn name(&self) -> &str {
"polars struct-json-encode"
}
fn description(&self) -> &str {
"Convert this struct to a string column with json values."
}
fn signature(&self) -> Signature {
Signature::build(self.name())
.category(Category::Custom("dataframe".into()))
.input_output_type(Type::custom("expression"), Type::custom("expression"))
}
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Encode a struct as JSON",
example: r#"[[id person]; [1 {name: "Bob", age: 36}] [2 {name: "Betty", age: 63}]]
| polars into-df -s {id: i32, person: {name: str, age: u8}}
| polars select id (polars col person | polars struct-json-encode | polars as encoded)
| polars sort-by id
| polars collect"#,
result: Some(
NuDataFrame::from(
df!(
"id" => [1i32, 2],
"encoded" => [
r#"{"name":"Bob","age":36}"#,
r#"{"name":"Betty","age":63}"#,
],
)
.expect("Should be able to create a simple dataframe"),
)
.into_value(Span::test_data()),
),
}]
}
fn run(
&self,
plugin: &Self::Plugin,
engine: &EngineInterface,
call: &EvaluatedCall,
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
NuExpression::try_from_pipeline(plugin, input, call.head)
.map(|expr| expr.into_polars().struct_().json_encode())
.map(NuExpression::from)
.and_then(|expr| expr.to_pipeline_data(plugin, engine, call.head))
.map_err(LabeledError::from)
}
}