From ff8831318dab1e3960be26c92936664894e70f5a Mon Sep 17 00:00:00 2001
From: Jack Wright <56345+ayax79@users.noreply.github.com>
Date: Tue, 6 May 2025 13:58:51 -0700
Subject: [PATCH] Added `polars struct-encode-json`, providing the ability to
 encode structs as json (#15678)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

# 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
---
 .../nu-plugin-test-support/src/plugin_test.rs | 42 +++++++++++-
 .../src/dataframe/command/data/mod.rs         |  2 +
 .../command/data/struct_json_encode.rs        | 67 +++++++++++++++++++
 3 files changed, 109 insertions(+), 2 deletions(-)
 create mode 100644 crates/nu_plugin_polars/src/dataframe/command/data/struct_json_encode.rs

diff --git a/crates/nu-plugin-test-support/src/plugin_test.rs b/crates/nu-plugin-test-support/src/plugin_test.rs
index 1d66674824..6fcdf47619 100644
--- a/crates/nu-plugin-test-support/src/plugin_test.rs
+++ b/crates/nu-plugin-test-support/src/plugin_test.rs
@@ -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);
diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/mod.rs b/crates/nu_plugin_polars/src/dataframe/command/data/mod.rs
index f4089a03c7..ed2e4a2393 100644
--- a/crates/nu_plugin_polars/src/dataframe/command/data/mod.rs
+++ b/crates/nu_plugin_polars/src/dataframe/command/data/mod.rs
@@ -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),
diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/struct_json_encode.rs b/crates/nu_plugin_polars/src/dataframe/command/data/struct_json_encode.rs
new file mode 100644
index 0000000000..a6d60f5a11
--- /dev/null
+++ b/crates/nu_plugin_polars/src/dataframe/command/data/struct_json_encode.rs
@@ -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)
+    }
+}