diff --git a/crates/nu-command/src/filters/drop/column.rs b/crates/nu-command/src/filters/drop/column.rs index 9cc530223b..f1496900f1 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 ecba9177d2..2f7fb9eaaa 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, "{}"); +}