From 021b8633cb8f776ca1a30db2597d107aa92fbf6f Mon Sep 17 00:00:00 2001 From: Jack Wright <56345+ayax79@users.noreply.github.com> Date: Sun, 9 Jun 2024 19:45:25 -0700 Subject: [PATCH] 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` --- .../src/dataframe/eager/to_nu.rs | 14 ++++---- .../src/dataframe/values/nu_dataframe/mod.rs | 32 +++++++++++++------ 2 files changed, 30 insertions(+), 16 deletions(-) diff --git a/crates/nu_plugin_polars/src/dataframe/eager/to_nu.rs b/crates/nu_plugin_polars/src/dataframe/eager/to_nu.rs index 8e3cdffa24..72aff7a728 100644 --- a/crates/nu_plugin_polars/src/dataframe/eager/to_nu.rs +++ b/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 { let rows: Option = 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)? } }; 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 1c322579dd..89a8c862e7 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 @@ -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, ShellError> { + pub fn print(&self, include_index: bool, span: Span) -> Result, 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, span: Span) -> Result, ShellError> { + pub fn head( + &self, + rows: Option, + include_index: bool, + span: Span, + ) -> Result, 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, span: Span) -> Result, ShellError> { + pub fn tail( + &self, + rows: Option, + include_index: bool, + span: Span, + ) -> Result, 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, 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 { - let vals = self.print(span)?; + let vals = self.print(true, span)?; Ok(Value::list(vals, span)) }