From f38657e6f352d8bacfdb984bc4d4dcf6307d99f3 Mon Sep 17 00:00:00 2001
From: Ian Manske <ian.manske@pm.me>
Date: Thu, 6 Jul 2023 17:54:59 +0000
Subject: [PATCH] Fix `headers` command handling of missing values (#9603)

# Description
This fixes the `headers` command handling of missing values (issue
#9602). Previously, each row in the table would have its columns set to
be exactly equal to the first row even if it had less columns than the
first row. This would cause to values magically change their column or
cause panics in other commands if rows ended up having more columns than
values.

# Tests
Added a missing values  test for the `headers` command
---
 crates/nu-command/src/filters/headers.rs    | 46 ++++++++++++++-------
 crates/nu-command/tests/commands/headers.rs | 13 ++++++
 2 files changed, 44 insertions(+), 15 deletions(-)

diff --git a/crates/nu-command/src/filters/headers.rs b/crates/nu-command/src/filters/headers.rs
index 8b3c1ae5a1..99e3953366 100644
--- a/crates/nu-command/src/filters/headers.rs
+++ b/crates/nu-command/src/filters/headers.rs
@@ -89,28 +89,38 @@ impl Command for Headers {
         let config = engine_state.get_config();
         let metadata = input.metadata();
         let value = input.into_value(call.head);
-        let headers = extract_headers(&value, config)?;
-        let new_headers = replace_headers(value, &headers)?;
+        let (old_headers, new_headers) = extract_headers(&value, config)?;
+        let new_headers = replace_headers(value, &old_headers, &new_headers)?;
 
         Ok(new_headers.into_pipeline_data().set_metadata(metadata))
     }
 }
 
-fn replace_headers(value: Value, headers: &[String]) -> Result<Value, ShellError> {
+fn replace_headers(
+    value: Value,
+    old_headers: &[String],
+    new_headers: &[String],
+) -> Result<Value, ShellError> {
     match value {
-        Value::Record { vals, span, .. } => {
-            let vals = vals.into_iter().take(headers.len()).collect();
-            Ok(Value::Record {
-                cols: headers.to_owned(),
-                vals,
-                span,
-            })
+        Value::Record { cols, vals, span } => {
+            let (cols, vals) = cols
+                .into_iter()
+                .zip(vals)
+                .filter_map(|(col, val)| {
+                    old_headers
+                        .iter()
+                        .position(|c| c == &col)
+                        .map(|i| (new_headers[i].clone(), val))
+                })
+                .unzip();
+
+            Ok(Value::Record { cols, vals, span })
         }
         Value::List { vals, span } => {
             let vals = vals
                 .into_iter()
                 .skip(1)
-                .map(|value| replace_headers(value, headers))
+                .map(|value| replace_headers(value, old_headers, new_headers))
                 .collect::<Result<Vec<Value>, ShellError>>()?;
 
             Ok(Value::List { vals, span })
@@ -133,9 +143,12 @@ fn is_valid_header(value: &Value) -> bool {
     )
 }
 
-fn extract_headers(value: &Value, config: &Config) -> Result<Vec<String>, ShellError> {
+fn extract_headers(
+    value: &Value,
+    config: &Config,
+) -> Result<(Vec<String>, Vec<String>), ShellError> {
     match value {
-        Value::Record { vals, .. } => {
+        Value::Record { cols, vals, .. } => {
             for v in vals {
                 if !is_valid_header(v) {
                     return Err(ShellError::TypeMismatch {
@@ -146,7 +159,8 @@ fn extract_headers(value: &Value, config: &Config) -> Result<Vec<String>, ShellE
                 }
             }
 
-            Ok(vals
+            let old_headers = cols.to_vec();
+            let new_headers = vals
                 .iter()
                 .enumerate()
                 .map(|(idx, value)| {
@@ -157,7 +171,9 @@ fn extract_headers(value: &Value, config: &Config) -> Result<Vec<String>, ShellE
                         col
                     }
                 })
-                .collect::<Vec<String>>())
+                .collect::<Vec<String>>();
+
+            Ok((old_headers, new_headers))
         }
         Value::List { vals, span } => vals
             .iter()
diff --git a/crates/nu-command/tests/commands/headers.rs b/crates/nu-command/tests/commands/headers.rs
index 88cc557866..566277a20f 100644
--- a/crates/nu-command/tests/commands/headers.rs
+++ b/crates/nu-command/tests/commands/headers.rs
@@ -30,6 +30,19 @@ fn headers_adds_missing_column_name() {
     assert_eq!(actual.out, r#"["r1c1","r2c1"]"#)
 }
 
+#[test]
+fn headers_handles_missing_values() {
+    let actual = nu!(pipeline(
+        r#"
+            [{x: a, y: b}, {x: 1, y: 2}, {x: 1, z: 3}]
+            | headers
+            | to nuon --raw
+        "#
+    ));
+
+    assert_eq!(actual.out, "[{a: 1, b: 2}, {a: 1}]")
+}
+
 #[test]
 fn headers_invalid_column_type_empty_record() {
     let actual = nu!(