diff --git a/crates/nu-command/src/filters/flatten.rs b/crates/nu-command/src/filters/flatten.rs index 8b3af912e..59cfb8a44 100644 --- a/crates/nu-command/src/filters/flatten.rs +++ b/crates/nu-command/src/filters/flatten.rs @@ -81,7 +81,7 @@ impl Command for Flatten { }, Example { description: "Flatten inner table", - example: "{ a: b, d: [ 1 2 3 4 ], e: [ 4 3 ] } | flatten d --all", + example: "{ a: b, d: [ 1 2 3 4 ], e: [ 4 3 ] } | flatten d --all", result: Some(Value::list( vec![ Value::test_record(record! { @@ -126,232 +126,197 @@ fn flatten( call: &Call, input: PipelineData, ) -> Result { - let tag = call.head; let columns: Vec = call.rest(engine_state, stack, 0)?; let metadata = input.metadata(); let flatten_all = call.has_flag("all"); input .flat_map( - move |item| flat_value(&columns, &item, tag, flatten_all), + move |item| flat_value(&columns, item, flatten_all), engine_state.ctrlc.clone(), ) .map(|x| x.set_metadata(metadata)) } -enum TableInside<'a> { +enum TableInside { // handle for a column which contains a single list(but not list of records) // it contains (column, span, values in the column, column index). - Entries(&'a str, &'a Span, Vec<&'a Value>, usize), + Entries(String, Vec, usize), // handle for a column which contains a table, we can flatten the inner column to outer level - // `columns` means that for the given row, it contains `len(columns)` nested rows, and each nested row contains a list of column name. - // Likely, `values` means that for the given row, it contains `len(values)` nested rows, and each nested row contains a list of values. - // + // `records` is the nested/inner table to flatten to the outer level // `parent_column_name` is handled for conflicting column name, the nested table may contains columns which has the same name // to outer level, for that case, the output column name should be f"{parent_column_name}_{inner_column_name}". // `parent_column_index` is the column index in original table. FlattenedRows { - columns: Vec>, - _span: &'a Span, - values: Vec>, - parent_column_name: &'a str, + records: Vec, + parent_column_name: String, parent_column_index: usize, }, } -fn flat_value(columns: &[CellPath], item: &Value, name_tag: Span, all: bool) -> Vec { +fn flat_value(columns: &[CellPath], item: Value, all: bool) -> Vec { let tag = item.span(); - if item.as_record().is_ok() { - let mut out = IndexMap::::new(); - let mut inner_table = None; + match item { + Value::Record { val, .. } => { + let mut out = IndexMap::::new(); + let mut inner_table = None; - let record = match item { - Value::Record { val, .. } => val, - // Propagate errors by explicitly matching them before the final case. - Value::Error { .. } => return vec![item.clone()], - other => { - return vec![Value::error( - ShellError::OnlySupportsThisInputType { - exp_input_type: "record".into(), - wrong_type: other.get_type().to_string(), - dst_span: name_tag, - src_span: other.span(), - }, - name_tag, - )]; - } - }; + for (column_index, (column, value)) in val.into_iter().enumerate() { + let column_requested = columns.iter().find(|c| c.into_string() == column); + let need_flatten = { columns.is_empty() || column_requested.is_some() }; + let span = value.span(); - let s = item.span(); - - for (column_index, (column, value)) in record.iter().enumerate() { - let column_requested = columns.iter().find(|c| c.into_string() == *column); - let need_flatten = { columns.is_empty() || column_requested.is_some() }; - let span = value.span(); - - match value { - Value::Record { val, .. } => { - if need_flatten { - val.iter().for_each(|(col, val)| { - if out.contains_key(col) { - out.insert(format!("{column}_{col}"), val.clone()); - } else { - out.insert(col.to_string(), val.clone()); - } - }) - } else if out.contains_key(column) { - out.insert(format!("{column}_{column}"), value.clone()); - } else { - out.insert(column.to_string(), value.clone()); - } - } - Value::List { vals, .. } if all && vals.iter().all(|f| f.as_record().is_ok()) => { - if need_flatten && inner_table.is_some() { - return vec![Value::error( ShellError::UnsupportedInput { msg: "can only flatten one inner list at a time. tried flattening more than one column with inner lists... but is flattened already".to_string(), input: "value originates from here".into(), msg_span: s, input_span: span }, span) - ]; - } - // it's a table (a list of record, we can flatten inner record) - let mut records = vec![]; - - for v in vals { - if let Ok(r) = v.as_record() { - records.push(r) - } - } - - if need_flatten { - let cols = records.iter().map(|r| r.cols.clone()); - let vals = records.iter().map(|r| r.vals.clone()); - - inner_table = Some(TableInside::FlattenedRows { - columns: cols.collect(), - _span: &s, - values: vals.collect(), - parent_column_name: column, - parent_column_index: column_index, - }); - } else if out.contains_key(column) { - out.insert(format!("{column}_{column}"), value.clone()); - } else { - out.insert(column.to_string(), value.clone()); - } - } - Value::List { vals: values, .. } => { - if need_flatten && inner_table.is_some() { - return vec![Value::error( ShellError::UnsupportedInput { msg: "can only flatten one inner list at a time. tried flattening more than one column with inner lists... but is flattened already".to_string(), input: "value originates from here".into(), msg_span: s, input_span: span }, span) - ]; - } - - if !columns.is_empty() { - let cell_path = column_requested.and_then(|x| match x.members.first() { - Some(PathMember::String { val, span: _, .. }) => Some(val), - _ => None, - }); - - if let Some(r) = cell_path { - inner_table = Some(TableInside::Entries( - r, - &s, - values.iter().collect::>(), - column_index, - )); - } else { - out.insert(column.to_string(), value.clone()); - } - } else { - inner_table = Some(TableInside::Entries( - column, - &s, - values.iter().collect::>(), - column_index, - )); - } - } - _ => { - out.insert(column.to_string(), value.clone()); - } - } - } - - let mut expanded = vec![]; - match inner_table { - Some(TableInside::Entries(column, _, entries, parent_column_index)) => { - for entry in entries { - let base = out.clone(); - let mut record = Record::new(); - let mut index = 0; - for (col, val) in base.into_iter() { - // meet the flattened column, push them to result record first - // this can avoid output column order changed. - if index == parent_column_index { - record.push(column, entry.clone()); - } - record.push(col, val); - index += 1; - } - // the flattened column may be the last column in the original table. - if index == parent_column_index { - record.push(column, entry.clone()); - } - expanded.push(Value::record(record, tag)); - } - } - Some(TableInside::FlattenedRows { - columns, - _span, - values, - parent_column_name, - parent_column_index, - }) => { - for (inner_cols, inner_vals) in columns.into_iter().zip(values) { - let base = out.clone(); - let mut record = Record::new(); - let mut index = 0; - - for (base_col, base_val) in base.into_iter() { - // meet the flattened column, push them to result record first - // this can avoid output column order changed. - if index == parent_column_index { - for (col, val) in inner_cols.iter().zip(inner_vals.iter()) { - if record.contains(col) { - record.push(format!("{parent_column_name}_{col}"), val.clone()); + match value { + Value::Record { val, .. } => { + if need_flatten { + for (col, val) in val { + if out.contains_key(&col) { + out.insert(format!("{column}_{col}"), val); } else { - record.push(col, val.clone()); - }; + out.insert(col, val); + } } + } else if out.contains_key(&column) { + out.insert(format!("{column}_{column}"), Value::record(val, span)); + } else { + out.insert(column, Value::record(val, span)); + } + } + Value::List { vals, .. } => { + if need_flatten && inner_table.is_some() { + return vec![Value::error( + ShellError::UnsupportedInput { + msg: "can only flatten one inner list at a time. tried flattening more than one column with inner lists... but is flattened already".into(), + input: "value originates from here".into(), + msg_span: tag, + input_span: span + }, + span, + )]; } - record.push(base_col, base_val); - index += 1; - } + if all && vals.iter().all(|f| f.as_record().is_ok()) { + // it's a table (a list of record, we can flatten inner record) + if need_flatten { + let records = vals + .into_iter() + .filter_map(|v| { + if let Value::Record { val, .. } = v { + Some(val) + } else { + None + } + }) + .collect(); - // the flattened column may be the last column in the original table. - if index == parent_column_index { - for (col, val) in inner_cols.iter().zip(inner_vals.iter()) { - if record.contains(col) { - record.push(format!("{parent_column_name}_{col}"), val.clone()); + inner_table = Some(TableInside::FlattenedRows { + records, + parent_column_name: column, + parent_column_index: column_index, + }); + } else if out.contains_key(&column) { + out.insert(format!("{column}_{column}"), Value::list(vals, span)); } else { - record.push(col, val.clone()); + out.insert(column, Value::list(vals, span)); } + } else if !columns.is_empty() { + let cell_path = + column_requested.and_then(|x| match x.members.first() { + Some(PathMember::String { val, .. }) => Some(val), + _ => None, + }); + + if let Some(r) = cell_path { + inner_table = + Some(TableInside::Entries(r.clone(), vals, column_index)); + } else { + out.insert(column, Value::list(vals, span)); + } + } else { + inner_table = Some(TableInside::Entries(column, vals, column_index)); } } - expanded.push(Value::record(record, tag)); + _ => { + out.insert(column, value); + } } } - None => { - expanded.push(Value::record(out.into_iter().collect(), tag)); + + let mut expanded = vec![]; + match inner_table { + Some(TableInside::Entries(column, entries, parent_column_index)) => { + for entry in entries { + let base = out.clone(); + let mut record = Record::new(); + let mut index = 0; + for (col, val) in base.into_iter() { + // meet the flattened column, push them to result record first + // this can avoid output column order changed. + if index == parent_column_index { + record.push(column.clone(), entry.clone()); + } + record.push(col, val); + index += 1; + } + // the flattened column may be the last column in the original table. + if index == parent_column_index { + record.push(column.clone(), entry); + } + expanded.push(Value::record(record, tag)); + } + } + Some(TableInside::FlattenedRows { + records, + parent_column_name, + parent_column_index, + }) => { + for inner_record in records { + let base = out.clone(); + let mut record = Record::new(); + let mut index = 0; + + for (base_col, base_val) in base { + // meet the flattened column, push them to result record first + // this can avoid output column order changed. + if index == parent_column_index { + for (col, val) in &inner_record { + if record.contains(col) { + record.push( + format!("{parent_column_name}_{col}"), + val.clone(), + ); + } else { + record.push(col, val.clone()); + }; + } + } + + record.push(base_col, base_val); + index += 1; + } + + // the flattened column may be the last column in the original table. + if index == parent_column_index { + for (col, val) in inner_record { + if record.contains(&col) { + record.push(format!("{parent_column_name}_{col}"), val); + } else { + record.push(col, val); + } + } + } + expanded.push(Value::record(record, tag)); + } + } + None => { + expanded.push(Value::record(out.into_iter().collect(), tag)); + } } + expanded } - expanded - } else if item.as_list().is_ok() { - if let Value::List { vals, .. } = item { - vals.to_vec() - } else { - vec![] - } - } else { - vec![item.clone()] + Value::List { vals, .. } => vals, + item => vec![item], } }