From 1bded8572c751dcab7a2c3940b8b84074f426a94 Mon Sep 17 00:00:00 2001
From: Jack Wright <56345+ayax79@users.noreply.github.com>
Date: Sat, 13 Apr 2024 04:33:29 -0700
Subject: [PATCH] Ensure that two columns named index don't exist when
 converting a Dataframe to a nu Value. (#12501)

# Description
@maxim-uvarov discovered an issue with the current implementation. When
executing [[index a]; [1 1]] | polars into-df, a plugin_failed_to_decode
error occurs. This happens because a Record is created with two columns
named "index" as an index column is added during conversion. This pull
request addresses the problem by not adding an index column if there is
already a column named "index" in the dataframe.

---------

Co-authored-by: Jack Wright <jack.wright@disqo.com>
---
 .../dataframe/values/nu_dataframe/conversion.rs  |  8 ++++++--
 .../src/dataframe/values/nu_dataframe/mod.rs     | 16 ++++++++++++----
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/conversion.rs b/crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/conversion.rs
index 9474e9cb78..f7941bf41d 100644
--- a/crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/conversion.rs
+++ b/crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/conversion.rs
@@ -154,12 +154,16 @@ pub fn create_column(
 
 // Adds a separator to the vector of values using the column names from the
 // dataframe to create the Values Row
-pub fn add_separator(values: &mut Vec<Value>, df: &DataFrame, span: Span) {
+// returns true if there is an index column contained in the dataframe
+pub fn add_separator(values: &mut Vec<Value>, df: &DataFrame, has_index: bool, span: Span) {
     let mut record = Record::new();
 
-    record.push("index", Value::string("...", span));
+    if !has_index {
+        record.push("index", Value::string("...", span));
+    }
 
     for name in df.get_column_names() {
+        // there should only be one index field
         record.push(name, Value::string("...", span))
     }
 
diff --git a/crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/mod.rs b/crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/mod.rs
index ea121994c8..644c1b11eb 100644
--- a/crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/mod.rs
+++ b/crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/mod.rs
@@ -296,6 +296,13 @@ impl NuDataFrame {
         }
     }
 
+    pub fn has_index(&self) -> bool {
+        self.columns(Span::unknown())
+            .unwrap_or_default() // just assume there isn't an index
+            .iter()
+            .any(|col| col.name() == "index")
+    }
+
     // Print is made out a head and if the dataframe is too large, then a tail
     pub fn print(&self, span: Span) -> Result<Vec<Value>, ShellError> {
         let df = &self.df;
@@ -304,7 +311,7 @@ impl NuDataFrame {
         if df.height() > size {
             let sample_size = size / 2;
             let mut values = self.head(Some(sample_size), span)?;
-            conversion::add_separator(&mut values, df, span);
+            conversion::add_separator(&mut values, df, self.has_index(), span);
             let remaining = df.height() - sample_size;
             let tail_size = remaining.min(sample_size);
             let mut tail_values = self.tail(Some(tail_size), span)?;
@@ -323,7 +330,6 @@ impl NuDataFrame {
     pub fn head(&self, rows: Option<usize>, span: Span) -> Result<Vec<Value>, ShellError> {
         let to_row = rows.unwrap_or(5);
         let values = self.to_rows(0, to_row, span)?;
-
         Ok(values)
     }
 
@@ -334,7 +340,6 @@ impl NuDataFrame {
         let from_row = to_row.saturating_sub(size);
 
         let values = self.to_rows(from_row, to_row, span)?;
-
         Ok(values)
     }
 
@@ -368,11 +373,14 @@ impl NuDataFrame {
             .map(|col| (col.name().to_string(), col.into_iter()))
             .collect::<Vec<(String, std::vec::IntoIter<Value>)>>();
 
+        let has_index = self.has_index();
         let values = (0..size)
             .map(|i| {
                 let mut record = Record::new();
 
-                record.push("index", Value::int((i + from_row) as i64, span));
+                if !has_index {
+                    record.push("index", Value::int((i + from_row) as i64, span));
+                }
 
                 for (name, col) in &mut iterators {
                     record.push(name.clone(), col.next().unwrap_or(Value::nothing(span)));