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`
This commit is contained in:
Jack Wright 2024-06-09 19:45:25 -07:00 committed by GitHub
parent 650ae537c3
commit 021b8633cb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 30 additions and 16 deletions

View File

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

View File

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