From c0901ef7072d3c3d9c7686c967f62945de8a8049 Mon Sep 17 00:00:00 2001 From: Fernando Herrera Date: Sun, 26 Jun 2022 17:32:18 -0500 Subject: [PATCH] Dataframe with real index (#5892) * remove extra print * dataframe with real index * corrected dataframe tests * clippy error * clippy error --- .../nu-command/src/dataframe/eager/to_nu.rs | 15 ++++++++----- .../values/nu_dataframe/conversion.rs | 6 ++++++ .../src/dataframe/values/nu_dataframe/mod.rs | 8 ++++++- crates/nu-command/src/viewers/table.rs | 21 ++++++++++++++++++- 4 files changed, 43 insertions(+), 7 deletions(-) diff --git a/crates/nu-command/src/dataframe/eager/to_nu.rs b/crates/nu-command/src/dataframe/eager/to_nu.rs index 73880bd70b..b455967861 100644 --- a/crates/nu-command/src/dataframe/eager/to_nu.rs +++ b/crates/nu-command/src/dataframe/eager/to_nu.rs @@ -34,15 +34,20 @@ impl Command for ToNu { } fn examples(&self) -> Vec { - let cols = vec!["a".into(), "b".into()]; + let cols = vec!["index".into(), "a".into(), "b".into()]; let rec_1 = Value::Record { cols: cols.clone(), - vals: vec![Value::test_int(1), Value::test_int(2)], + vals: vec![Value::test_int(0), Value::test_int(1), Value::test_int(2)], span: Span::test_data(), }; let rec_2 = Value::Record { + cols: cols.clone(), + vals: vec![Value::test_int(1), Value::test_int(3), Value::test_int(4)], + span: Span::test_data(), + }; + let rec_3 = Value::Record { cols, - vals: vec![Value::test_int(3), Value::test_int(4)], + vals: vec![Value::test_int(2), Value::test_int(3), Value::test_int(4)], span: Span::test_data(), }; @@ -51,7 +56,7 @@ impl Command for ToNu { description: "Shows head rows from dataframe", example: "[[a b]; [1 2] [3 4]] | into df | into nu", result: Some(Value::List { - vals: vec![rec_1, rec_2.clone()], + vals: vec![rec_1, rec_2], span: Span::test_data(), }), }, @@ -59,7 +64,7 @@ impl Command for ToNu { description: "Shows tail rows from dataframe", example: "[[a b]; [1 2] [5 6] [3 4]] | into df | into nu -t -n 1", result: Some(Value::List { - vals: vec![rec_2], + vals: vec![rec_3], span: Span::test_data(), }), }, diff --git a/crates/nu-command/src/dataframe/values/nu_dataframe/conversion.rs b/crates/nu-command/src/dataframe/values/nu_dataframe/conversion.rs index 7c1c94f9ff..0322d8b952 100644 --- a/crates/nu-command/src/dataframe/values/nu_dataframe/conversion.rs +++ b/crates/nu-command/src/dataframe/values/nu_dataframe/conversion.rs @@ -555,6 +555,12 @@ pub fn add_separator(values: &mut Vec, df: &DataFrame, span: Span) { let mut cols = vec![]; let mut vals = vec![]; + cols.push("index".to_string()); + vals.push(Value::String { + val: "...".into(), + span, + }); + for name in df.get_column_names() { cols.push(name.to_string()); vals.push(Value::String { diff --git a/crates/nu-command/src/dataframe/values/nu_dataframe/mod.rs b/crates/nu-command/src/dataframe/values/nu_dataframe/mod.rs index 60772fe8da..bd80a9da44 100644 --- a/crates/nu-command/src/dataframe/values/nu_dataframe/mod.rs +++ b/crates/nu-command/src/dataframe/values/nu_dataframe/mod.rs @@ -400,10 +400,16 @@ impl NuDataFrame { let values = (0..size) .into_iter() - .map(|_| { + .map(|i| { let mut cols = vec![]; let mut vals = vec![]; + cols.push("index".into()); + vals.push(Value::Int { + val: (i + from_row) as i64, + span, + }); + for (name, col) in &mut iterators { cols.push(name.clone()); diff --git a/crates/nu-command/src/viewers/table.rs b/crates/nu-command/src/viewers/table.rs index 26f640db3b..a3811430ca 100644 --- a/crates/nu-command/src/viewers/table.rs +++ b/crates/nu-command/src/viewers/table.rs @@ -17,6 +17,7 @@ use terminal_size::{Height, Width}; const STREAM_PAGE_SIZE: usize = 1000; const STREAM_TIMEOUT_CHECK_INTERVAL: usize = 100; +const INDEX_COLUMN_NAME: &str = "index"; fn get_width_param(width_param: Option) -> usize { if let Some(col) = width_param { @@ -41,6 +42,10 @@ impl Command for Table { "Render the table." } + fn extra_usage(&self) -> &str { + "If the table contains a column called 'index', this column is used as the table index instead of the usual continuous index" + } + fn search_terms(&self) -> Vec<&str> { vec!["display", "render"] } @@ -374,9 +379,23 @@ fn convert_to_table( // String1 = datatype, String2 = value as string let mut row: Vec<(String, String)> = vec![]; if !disable_index { - row = vec![("string".to_string(), (row_num + row_offset).to_string())]; + let row_val = match &item { + Value::Record { .. } => item + .get_data_by_key(INDEX_COLUMN_NAME) + .map(|value| value.into_string("", config)), + _ => None, + } + .unwrap_or_else(|| (row_num + row_offset).to_string()); + row = vec![("string".to_string(), (row_val).to_string())]; } + // The header with the INDEX is removed from the table headers since + // it is added to the natural table index + headers = headers + .into_iter() + .filter(|header| header != INDEX_COLUMN_NAME) + .collect(); + if headers.is_empty() { row.push(( item.get_type().to_string(),