mirror of
https://github.com/nushell/nushell.git
synced 2024-11-25 09:53:43 +01:00
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:
parent
650ae537c3
commit
021b8633cb
@ -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)?
|
||||
}
|
||||
};
|
||||
|
||||
|
@ -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))
|
||||
}
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user