From 33a7bc405fa9b4f446bbbc0f0f0d4e488a3278c8 Mon Sep 17 00:00:00 2001 From: Ian Manske Date: Thu, 9 Nov 2023 12:51:46 +0000 Subject: [PATCH] Refactor `drop columns` to fix issues (#10903) # Description This PR refactors `drop columns` and fixes issues #10902 and #6846. Tables with "holes" are now handled consistently, although still somewhat awkwardly. That is, the columns in the first row are used to determine which columns to drop, meaning that the columns displayed all the way to the right by `table` may not be the columns actually being dropped. For example, `[{a: 1}, {b: 2}] | drop column` will drop column `a` instead of `b`. Before, this would give a list of empty records. # User-Facing Changes `drop columns` can now take records as input. --- crates/nu-command/src/filters/drop/column.rs | 221 +++++++++---------- crates/nu-command/tests/commands/drop.rs | 42 ++++ 2 files changed, 152 insertions(+), 111 deletions(-) diff --git a/crates/nu-command/src/filters/drop/column.rs b/crates/nu-command/src/filters/drop/column.rs index 9cc530223..f1496900f 100644 --- a/crates/nu-command/src/filters/drop/column.rs +++ b/crates/nu-command/src/filters/drop/column.rs @@ -1,11 +1,13 @@ use nu_engine::CallExt; -use nu_protocol::ast::{Call, CellPath}; +use nu_protocol::ast::Call; use nu_protocol::engine::{Command, EngineState, Stack}; use nu_protocol::{ - record, Category, Example, FromValue, IntoInterruptiblePipelineData, IntoPipelineData, - PipelineData, Record, ShellError, Signature, Span, SyntaxShape, Type, Value, + record, Category, Example, IntoInterruptiblePipelineData, IntoPipelineData, PipelineData, + ShellError, Signature, Span, Spanned, SyntaxShape, Type, Value, }; +use std::collections::HashSet; + #[derive(Clone)] pub struct DropColumn; @@ -16,7 +18,10 @@ impl Command for DropColumn { fn signature(&self) -> Signature { Signature::build(self.name()) - .input_output_types(vec![(Type::Table(vec![]), Type::Table(vec![]))]) + .input_output_types(vec![ + (Type::Table(vec![]), Type::Table(vec![])), + (Type::Record(vec![]), Type::Record(vec![])), + ]) .optional( "columns", SyntaxShape::Int, @@ -41,148 +46,142 @@ impl Command for DropColumn { input: PipelineData, ) -> Result { // the number of columns to drop - let columns: Option = call.opt(engine_state, stack, 0)?; - let span = call.head; + let columns: Option> = call.opt(engine_state, stack, 0)?; - let columns_to_drop = if let Some(quantity) = columns { - quantity + let columns = if let Some(columns) = columns { + if columns.item < 0 { + return Err(ShellError::NeedsPositiveValue(columns.span)); + } else { + columns.item as usize + } } else { 1 }; - // Make columns to drop is positive - if columns_to_drop < 0 { - if let Some(expr) = call.positional_nth(0) { - return Err(ShellError::NeedsPositiveValue(expr.span)); - } - } - - dropcol(engine_state, span, input, columns_to_drop) + drop_cols(engine_state, input, call.head, columns) } fn examples(&self) -> Vec { - vec![Example { - description: "Remove the last column of a table", - example: "[[lib, extension]; [nu-lib, rs] [nu-core, rb]] | drop column", - result: Some(Value::list( - vec![ - Value::test_record(record!("lib" =>Value::test_string("nu-lib"))), - Value::test_record(record!("lib" =>Value::test_string("nu-core"))), - ], - Span::test_data(), - )), - }] + vec![ + Example { + description: "Remove the last column of a table", + example: "[[lib, extension]; [nu-lib, rs] [nu-core, rb]] | drop column", + result: Some(Value::test_list(vec![ + Value::test_record(record! { "lib" => Value::test_string("nu-lib") }), + Value::test_record(record! { "lib" => Value::test_string("nu-core") }), + ])), + }, + Example { + description: "Remove the last column of a record", + example: "{lib: nu-lib, extension: rs} | drop column", + result: Some(Value::test_record( + record! { "lib" => Value::test_string("nu-lib") }, + )), + }, + ] } } -fn dropcol( +fn drop_cols( engine_state: &EngineState, - span: Span, input: PipelineData, - columns: i64, // the number of columns to drop + head: Span, + columns: usize, ) -> Result { - let mut keep_columns = vec![]; - + // For simplicity and performance, we use the first row's columns + // as the columns for the whole table, and assume that later rows/records + // have these same columns. However, this can give weird results like: + // `[{a: 1}, {b: 2}] | drop column` + // This will drop the column "a" instead of "b" even though column "b" + // is displayed farther to the right. match input { - PipelineData::ListStream(stream, ..) => { - let mut output = vec![]; + PipelineData::ListStream(mut stream, ..) => { + if let Some(mut first) = stream.next() { + let drop_cols = drop_cols_set(&mut first, head, columns)?; - let v: Vec<_> = stream.into_iter().collect(); - let input_cols = get_input_cols(v.clone()); - let kc = get_keep_columns(input_cols, columns); - keep_columns = get_cellpath_columns(kc, span); - - for input_val in v { - let mut record = Record::new(); - - for path in &keep_columns { - let fetcher = input_val.clone().follow_cell_path(&path.members, false)?; - record.push(path.into_string(), fetcher); - } - output.push(Value::record(record, span)) - } - - Ok(output - .into_iter() - .into_pipeline_data(engine_state.ctrlc.clone())) - } - PipelineData::Value(v, ..) => { - let val_span = v.span(); - if let Value::List { - vals: input_vals, .. - } = v - { - let mut output = vec![]; - let input_cols = get_input_cols(input_vals.clone()); - let kc = get_keep_columns(input_cols, columns); - keep_columns = get_cellpath_columns(kc, val_span); - - for input_val in input_vals { - let mut record = Record::new(); - - for path in &keep_columns { - let fetcher = input_val.clone().follow_cell_path(&path.members, false)?; - record.push(path.into_string(), fetcher); - } - output.push(Value::record(record, val_span)) - } - - Ok(output - .into_iter() + Ok(std::iter::once(first) + .chain(stream.map(move |mut v| { + match drop_record_cols(&mut v, head, &drop_cols) { + Ok(()) => v, + Err(e) => Value::error(e, head), + } + })) .into_pipeline_data(engine_state.ctrlc.clone())) } else { - let mut record = Record::new(); - - for cell_path in &keep_columns { - let result = v.clone().follow_cell_path(&cell_path.members, false)?; - record.push(cell_path.into_string(), result); - } - - Ok(Value::record(record, span).into_pipeline_data()) + Ok(PipelineData::Empty) } } - x => Ok(x), + PipelineData::Value(v, ..) => { + let span = v.span(); + match v { + Value::List { mut vals, .. } => { + if let Some((first, rest)) = vals.split_first_mut() { + let drop_cols = drop_cols_set(first, head, columns)?; + for val in rest { + drop_record_cols(val, head, &drop_cols)? + } + } + Ok(Value::list(vals, span).into_pipeline_data()) + } + Value::Record { + val: mut record, .. + } => { + let len = record.len().saturating_sub(columns); + record.truncate(len); + Ok(Value::record(record, span).into_pipeline_data()) + } + // Propagate errors + Value::Error { error, .. } => Err(*error), + val => Err(unsupported_value_error(&val, head)), + } + } + PipelineData::Empty => Ok(PipelineData::Empty), + PipelineData::ExternalStream { span, .. } => Err(ShellError::OnlySupportsThisInputType { + exp_input_type: "table or record".into(), + wrong_type: "raw data".into(), + dst_span: head, + src_span: span, + }), } } -fn get_input_cols(input: Vec) -> Vec { - let rec = input.first(); - match rec { - Some(Value::Record { val, .. }) => val.cols.to_vec(), - _ => vec!["".to_string()], +fn drop_cols_set(val: &mut Value, head: Span, drop: usize) -> Result, ShellError> { + if let Value::Record { val: record, .. } = val { + let len = record.len().saturating_sub(drop); + Ok(record.drain(len..).map(|(col, _)| col).collect()) + } else { + Err(unsupported_value_error(val, head)) } } -fn get_cellpath_columns(keep_cols: Vec, span: Span) -> Vec { - let mut output = vec![]; - for keep_col in keep_cols { - let val = Value::string(keep_col, span); - let cell_path = match CellPath::from_value(val) { - Ok(v) => v, - Err(_) => return vec![], - }; - output.push(cell_path); +fn drop_record_cols( + val: &mut Value, + head: Span, + drop_cols: &HashSet, +) -> Result<(), ShellError> { + if let Value::Record { val, .. } = val { + val.retain(|col, _| !drop_cols.contains(col)); + Ok(()) + } else { + Err(unsupported_value_error(val, head)) } - output } -fn get_keep_columns(input: Vec, mut num_of_columns_to_drop: i64) -> Vec { - let vlen: i64 = input.len() as i64; - - if num_of_columns_to_drop > vlen { - num_of_columns_to_drop = vlen; +fn unsupported_value_error(val: &Value, head: Span) -> ShellError { + ShellError::OnlySupportsThisInputType { + exp_input_type: "table or record".into(), + wrong_type: val.get_type().to_string(), + dst_span: head, + src_span: val.span(), } - - let num_of_columns_to_keep = (vlen - num_of_columns_to_drop) as usize; - input[0..num_of_columns_to_keep].to_vec() } #[cfg(test)] mod test { + use super::*; + #[test] fn test_examples() { - use super::DropColumn; - use crate::test_examples; - test_examples(DropColumn {}) + crate::test_examples(DropColumn) } } diff --git a/crates/nu-command/tests/commands/drop.rs b/crates/nu-command/tests/commands/drop.rs index ecba9177d..2f7fb9eaa 100644 --- a/crates/nu-command/tests/commands/drop.rs +++ b/crates/nu-command/tests/commands/drop.rs @@ -93,3 +93,45 @@ fn fail_on_non_iterator() { assert!(actual.err.contains("command doesn't support")); } + +#[test] +fn disjoint_columns_first_row_becomes_empty() { + let actual = nu!(pipeline( + " + [{a: 1}, {b: 2, c: 3}] + | drop column + | columns + | to nuon + " + )); + + assert_eq!(actual.out, "[b, c]"); +} + +#[test] +fn disjoint_columns() { + let actual = nu!(pipeline( + " + [{a: 1, b: 2}, {c: 3}] + | drop column + | columns + | to nuon + " + )); + + assert_eq!(actual.out, "[a, c]"); +} + +#[test] +fn record() { + let actual = nu!("{a: 1, b: 2} | drop column | to nuon"); + + assert_eq!(actual.out, "{a: 1}"); +} + +#[test] +fn more_columns_than_record_has() { + let actual = nu!("{a: 1, b: 2} | drop column 3 | to nuon"); + + assert_eq!(actual.out, "{}"); +}