From b8e6dde4ff82dd988bced7c61cfd768137d1eb41 Mon Sep 17 00:00:00 2001
From: Julian Amarilla <julian.amarilla@compass.com>
Date: Fri, 25 Apr 2025 02:01:35 -0300
Subject: [PATCH] make polars commands propagate metadata

---
 .../dataframe/command/aggregation/agg_groups.rs    | 14 ++++++++++++++
 .../src/dataframe/command/aggregation/aggregate.rs | 14 ++++++++++++++
 .../src/dataframe/command/aggregation/count.rs     | 14 ++++++++++++++
 .../src/dataframe/command/aggregation/implode.rs   | 14 ++++++++++++++
 .../src/dataframe/command/core/summary.rs          | 14 ++++++++++++++
 .../src/dataframe/command/data/first.rs            | 14 ++++++++++++++
 .../src/dataframe/command/data/join.rs             | 14 ++++++++++++++
 .../src/dataframe/command/data/len.rs              | 14 ++++++++++++++
 .../src/dataframe/command/string/str_slice.rs      | 14 ++++++++++++++
 9 files changed, 126 insertions(+)

diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/agg_groups.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/agg_groups.rs
index a8d58bd6f0..47bacfa694 100644
--- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/agg_groups.rs
+++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/agg_groups.rs
@@ -60,6 +60,20 @@ impl PluginCommand for ExprAggGroups {
         engine: &EngineInterface,
         call: &EvaluatedCall,
         input: PipelineData,
+    ) -> Result<PipelineData, LabeledError> {
+        let metadata = input.metadata();
+        self.run_inner(plugin, engine, call, input)
+            .map(|pd| pd.set_metadata(metadata))
+    }
+}
+
+impl ExprAggGroups {
+    fn run_inner(
+        &self,
+        plugin: &PolarsPlugin,
+        engine: &EngineInterface,
+        call: &EvaluatedCall,
+        input: PipelineData,
     ) -> Result<PipelineData, LabeledError> {
         let value = input.into_value(call.head)?;
         match PolarsPluginObject::try_from_value(plugin, &value)? {
diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/aggregate.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/aggregate.rs
index 64ff8952c5..84bf8ddcfa 100644
--- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/aggregate.rs
+++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/aggregate.rs
@@ -126,6 +126,20 @@ impl PluginCommand for LazyAggregate {
         engine: &EngineInterface,
         call: &EvaluatedCall,
         input: PipelineData,
+    ) -> Result<PipelineData, LabeledError> {
+        let metadata = input.metadata();
+        self.run_inner(plugin, engine, call, input)
+            .map(|pd| pd.set_metadata(metadata))
+    }
+}
+
+impl LazyAggregate {
+    fn run_inner(
+        &self,
+        plugin: &PolarsPlugin,
+        engine: &EngineInterface,
+        call: &EvaluatedCall,
+        input: PipelineData,
     ) -> Result<PipelineData, LabeledError> {
         let vals: Vec<Value> = call.rest(0)?;
         let value = Value::list(vals, call.head);
diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/count.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/count.rs
index 3b1b32952b..39c3b9f325 100644
--- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/count.rs
+++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/count.rs
@@ -55,6 +55,20 @@ impl PluginCommand for ExprCount {
         engine: &EngineInterface,
         call: &EvaluatedCall,
         input: PipelineData,
+    ) -> Result<PipelineData, LabeledError> {
+        let metadata = input.metadata();
+        self.run_inner(plugin, engine, call, input)
+            .map(|pd| pd.set_metadata(metadata))
+    }
+}
+
+impl ExprCount {
+    fn run_inner(
+        &self,
+        plugin: &PolarsPlugin,
+        engine: &EngineInterface,
+        call: &EvaluatedCall,
+        input: PipelineData,
     ) -> Result<PipelineData, LabeledError> {
         let value = input.into_value(call.head)?;
         match PolarsPluginObject::try_from_value(plugin, &value)? {
diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/implode.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/implode.rs
index c46dd6f7a8..fb2f45947f 100644
--- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/implode.rs
+++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/implode.rs
@@ -52,6 +52,20 @@ impl PluginCommand for ExprImplode {
         engine: &EngineInterface,
         call: &EvaluatedCall,
         input: PipelineData,
+    ) -> Result<PipelineData, LabeledError> {
+        let metadata = input.metadata();
+        self.run_inner(plugin, engine, call, input)
+            .map(|pd| pd.set_metadata(metadata))
+    }
+}
+
+impl ExprImplode {
+    fn run_inner(
+        &self,
+        plugin: &PolarsPlugin,
+        engine: &EngineInterface,
+        call: &EvaluatedCall,
+        input: PipelineData,
     ) -> Result<PipelineData, LabeledError> {
         let value = input.into_value(call.head)?;
         match PolarsPluginObject::try_from_value(plugin, &value)? {
diff --git a/crates/nu_plugin_polars/src/dataframe/command/core/summary.rs b/crates/nu_plugin_polars/src/dataframe/command/core/summary.rs
index b00a8bb201..342b76d97e 100644
--- a/crates/nu_plugin_polars/src/dataframe/command/core/summary.rs
+++ b/crates/nu_plugin_polars/src/dataframe/command/core/summary.rs
@@ -111,6 +111,20 @@ impl PluginCommand for Summary {
         engine: &EngineInterface,
         call: &EvaluatedCall,
         input: PipelineData,
+    ) -> Result<PipelineData, LabeledError> {
+        let metadata = input.metadata();
+        self.run_inner(plugin, engine, call, input)
+            .map(|pd| pd.set_metadata(metadata))
+    }
+}
+
+impl Summary {
+    fn run_inner(
+        &self,
+        plugin: &PolarsPlugin,
+        engine: &EngineInterface,
+        call: &EvaluatedCall,
+        input: PipelineData,
     ) -> Result<PipelineData, LabeledError> {
         command(plugin, engine, call, input).map_err(LabeledError::from)
     }
diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/first.rs b/crates/nu_plugin_polars/src/dataframe/command/data/first.rs
index 552b173e5c..eca853a5ae 100644
--- a/crates/nu_plugin_polars/src/dataframe/command/data/first.rs
+++ b/crates/nu_plugin_polars/src/dataframe/command/data/first.rs
@@ -96,6 +96,20 @@ impl PluginCommand for FirstDF {
         engine: &EngineInterface,
         call: &EvaluatedCall,
         input: PipelineData,
+    ) -> Result<PipelineData, LabeledError> {
+        let metadata = input.metadata();
+        self.run_inner(plugin, engine, call, input)
+            .map(|pd| pd.set_metadata(metadata))
+    }
+}
+
+impl FirstDF {
+    fn run_inner(
+        &self,
+        plugin: &PolarsPlugin,
+        engine: &EngineInterface,
+        call: &EvaluatedCall,
+        input: PipelineData,
     ) -> Result<PipelineData, LabeledError> {
         let value = input.into_value(call.head)?;
         match PolarsPluginObject::try_from_value(plugin, &value)? {
diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/join.rs b/crates/nu_plugin_polars/src/dataframe/command/data/join.rs
index bd140b5c2a..dcf5f6c434 100644
--- a/crates/nu_plugin_polars/src/dataframe/command/data/join.rs
+++ b/crates/nu_plugin_polars/src/dataframe/command/data/join.rs
@@ -240,6 +240,20 @@ impl PluginCommand for LazyJoin {
         engine: &EngineInterface,
         call: &EvaluatedCall,
         input: PipelineData,
+    ) -> Result<PipelineData, LabeledError> {
+        let metadata = input.metadata();
+        self.run_inner(plugin, engine, call, input)
+            .map(|pd| pd.set_metadata(metadata))
+    }
+}
+
+impl LazyJoin {
+    fn run_inner(
+        &self,
+        plugin: &PolarsPlugin,
+        engine: &EngineInterface,
+        call: &EvaluatedCall,
+        input: PipelineData,
     ) -> Result<PipelineData, LabeledError> {
         let left = call.has_flag("left")?;
         let full = call.has_flag("full")?;
diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/len.rs b/crates/nu_plugin_polars/src/dataframe/command/data/len.rs
index f537a6baa1..77f88ec6dc 100644
--- a/crates/nu_plugin_polars/src/dataframe/command/data/len.rs
+++ b/crates/nu_plugin_polars/src/dataframe/command/data/len.rs
@@ -55,6 +55,20 @@ impl PluginCommand for ExprLen {
         plugin: &Self::Plugin,
         engine: &EngineInterface,
         call: &EvaluatedCall,
+        input: PipelineData,
+    ) -> Result<PipelineData, LabeledError> {
+        let metadata = input.metadata();
+        self.run_inner(plugin, engine, call, input)
+            .map(|pd| pd.set_metadata(metadata))
+    }
+}
+
+impl ExprLen {
+    fn run_inner(
+        &self,
+        plugin: &PolarsPlugin,
+        engine: &EngineInterface,
+        call: &EvaluatedCall,
         _input: PipelineData,
     ) -> Result<PipelineData, LabeledError> {
         let res: NuExpression = polars::prelude::len().into();
diff --git a/crates/nu_plugin_polars/src/dataframe/command/string/str_slice.rs b/crates/nu_plugin_polars/src/dataframe/command/string/str_slice.rs
index 02d35e6984..94a040354d 100644
--- a/crates/nu_plugin_polars/src/dataframe/command/string/str_slice.rs
+++ b/crates/nu_plugin_polars/src/dataframe/command/string/str_slice.rs
@@ -116,6 +116,20 @@ impl PluginCommand for StrSlice {
         engine: &EngineInterface,
         call: &EvaluatedCall,
         input: PipelineData,
+    ) -> Result<PipelineData, LabeledError> {
+        let metadata = input.metadata();
+        self.run_inner(plugin, engine, call, input)
+            .map(|pd| pd.set_metadata(metadata))
+    }
+}
+
+impl StrSlice {
+    fn run_inner(
+        &self,
+        plugin: &PolarsPlugin,
+        engine: &EngineInterface,
+        call: &EvaluatedCall,
+        input: PipelineData,
     ) -> Result<PipelineData, LabeledError> {
         let value = input.into_value(call.head)?;
         match PolarsPluginObject::try_from_value(plugin, &value)? {