diff --git a/crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/conversion.rs b/crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/conversion.rs index 9474e9cb78..f7941bf41d 100644 --- a/crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/conversion.rs +++ b/crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/conversion.rs @@ -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, df: &DataFrame, span: Span) { +// returns true if there is an index column contained in the dataframe +pub fn add_separator(values: &mut Vec, 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)) } 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 ea121994c8..644c1b11eb 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 @@ -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, ShellError> { let df = &self.df; @@ -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)?; @@ -323,7 +330,6 @@ impl NuDataFrame { pub fn head(&self, rows: Option, span: Span) -> Result, ShellError> { let to_row = rows.unwrap_or(5); let values = self.to_rows(0, to_row, span)?; - Ok(values) } @@ -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) } @@ -368,11 +373,14 @@ impl NuDataFrame { .map(|col| (col.name().to_string(), col.into_iter())) .collect::)>>(); + 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)));