Allow the addition of an index column to be optional (#13097)

Per discussion on discord dataframes channel with @maxim-uvarov and pyz.

When converting a dataframe to an nushell value via `polars into-nu`,
the index column should not be added by default and should only be added
when specifying `--index`

crates/nu_plugin_polars/src/dataframe/eager/to_nu.rs

@@ -35,6 +35,7 @@ impl PluginCommand for ToNu {
                 Some('n'),
             )
             .switch("tail", "shows tail rows", Some('t'))
+            .switch("index", "add an index column", Some('i'))
             .input_output_types(vec![
                 (Type::Custom("expression".into()), Type::Any),
                 (Type::Custom("dataframe".into()), Type::table()),
@@ -62,18 +63,18 @@ impl PluginCommand for ToNu {
         vec![
             Example {
                 description: "Shows head rows from dataframe",
-                example: "[[a b]; [1 2] [3 4]] | polars into-df | polars into-nu",
+                example: "[[a b]; [1 2] [3 4]] | polars into-df | polars into-nu --index",
                 result: Some(Value::list(vec![rec_1, rec_2], Span::test_data())),
             },
             Example {
                 description: "Shows tail rows from dataframe",
                 example:
-                    "[[a b]; [1 2] [5 6] [3 4]] | polars into-df | polars into-nu --tail --rows 1",
+                    "[[a b]; [1 2] [5 6] [3 4]] | polars into-df | polars into-nu --tail --rows 1 --index",
                 result: Some(Value::list(vec![rec_3], Span::test_data())),
             },
             Example {
                 description: "Convert a col expression into a nushell value",
-                example: "polars col a | polars into-nu",
+                example: "polars col a | polars into-nu --index",
                 result: Some(Value::test_record(record! {
                     "expr" =>  Value::test_string("column"),
                     "value" => Value::test_string("a"),
@@ -106,17 +107,18 @@ fn dataframe_command(
 ) -> Result<PipelineData, ShellError> {
     let rows: Option<usize> = call.get_flag("rows")?;
     let tail: bool = call.has_flag("tail")?;
+    let index: bool = call.has_flag("index")?;
 
     let df = NuDataFrame::try_from_value_coerce(plugin, &input, call.head)?;
 
     let values = if tail {
-        df.tail(rows, call.head)?
+        df.tail(rows, index, call.head)?
     } else {
         // if rows is specified, return those rows, otherwise return everything
         if rows.is_some() {
-            df.head(rows, call.head)?
+            df.head(rows, index, call.head)?
         } else {
-            df.head(Some(df.height()), call.head)?
+            df.head(Some(df.height()), index, call.head)?
         }
     };
 

crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/mod.rs

@@ -326,22 +326,22 @@ impl NuDataFrame {
     }
 
     // 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> {
+    pub fn print(&self, include_index: bool, span: Span) -> Result<Vec<Value>, ShellError> {
         let df = &self.df;
         let size: usize = 20;
 
         if df.height() > size {
             let sample_size = size / 2;
-            let mut values = self.head(Some(sample_size), span)?;
+            let mut values = self.head(Some(sample_size), include_index, 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)?;
+            let mut tail_values = self.tail(Some(tail_size), include_index, span)?;
             values.append(&mut tail_values);
 
             Ok(values)
         } else {
-            Ok(self.head(Some(size), span)?)
+            Ok(self.head(Some(size), include_index, span)?)
         }
     }
 
@@ -349,26 +349,38 @@ impl NuDataFrame {
         self.df.height()
     }
 
-    pub fn head(&self, rows: Option<usize>, span: Span) -> Result<Vec<Value>, ShellError> {
+    pub fn head(
+        &self,
+        rows: Option<usize>,
+        include_index: bool,
+        span: Span,
+    ) -> Result<Vec<Value>, ShellError> {
         let to_row = rows.unwrap_or(5);
-        let values = self.to_rows(0, to_row, span)?;
+        let values = self.to_rows(0, to_row, include_index, span)?;
         Ok(values)
     }
 
-    pub fn tail(&self, rows: Option<usize>, span: Span) -> Result<Vec<Value>, ShellError> {
+    pub fn tail(
+        &self,
+        rows: Option<usize>,
+        include_index: bool,
+        span: Span,
+    ) -> Result<Vec<Value>, ShellError> {
         let df = &self.df;
         let to_row = df.height();
         let size = rows.unwrap_or(DEFAULT_ROWS);
         let from_row = to_row.saturating_sub(size);
 
-        let values = self.to_rows(from_row, to_row, span)?;
+        let values = self.to_rows(from_row, to_row, include_index, span)?;
         Ok(values)
     }
 
+    /// Converts the dataframe to a nushell list of values
     pub fn to_rows(
         &self,
         from_row: usize,
         to_row: usize,
+        include_index: bool,
         span: Span,
     ) -> Result<Vec<Value>, ShellError> {
         let df = &self.df;
@@ -400,7 +412,7 @@ impl NuDataFrame {
             .map(|i| {
                 let mut record = Record::new();
 
-                if !has_index {
+                if !has_index && include_index {
                     record.push("index", Value::int((i + from_row) as i64, span));
                 }
 
@@ -602,7 +614,7 @@ impl CustomValueSupport for NuDataFrame {
     }
 
     fn base_value(self, span: Span) -> Result<Value, ShellError> {
-        let vals = self.print(span)?;
+        let vals = self.print(true, span)?;
         Ok(Value::list(vals, span))
     }