Skip to content

Commit

Permalink
Ensure that two columns named index don't exist when converting a Dat…
Browse files Browse the repository at this point in the history
…aframe to a nu Value. (nushell#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>
  • Loading branch information
ayax79 and ayax79 authored Apr 13, 2024
1 parent f975c99 commit 1bded85
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 6 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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))
}

Expand Down
16 changes: 12 additions & 4 deletions crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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)?;
Expand All @@ -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)
}

Expand All @@ -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)
}

Expand Down Expand Up @@ -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)));
Expand Down

0 comments on commit 1bded85

Please sign in to comment.