Ensure map-like row data preserves column order

Author: Abestanis

src/serializer.rs

@@ -292,9 +292,9 @@ impl<'a, 'w, W: io::Write> SerializeMap for &'a mut SeRecord<'w, W> {
 
     fn serialize_key<T: ?Sized + Serialize>(
         &mut self,
-        _key: &T,
+        key: &T,
     ) -> Result<(), Self::Error> {
-        Ok(())
+        self.wtr.check_map_key(key)
     }
 
     fn serialize_value<T: ?Sized + Serialize>(
@@ -305,6 +305,7 @@ impl<'a, 'w, W: io::Write> SerializeMap for &'a mut SeRecord<'w, W> {
     }
 
     fn end(self) -> Result<Self::Ok, Self::Error> {
+        self.wtr.on_map_end();
         Ok(())
     }
 }
@@ -742,6 +743,7 @@ impl<'a, 'w, W: io::Write> SerializeMap for &'a mut SeHeader<'w, W> {
             return Err(err);
         }
 
+        self.wtr.check_map_key(key)?;
         let mut key_serializer = SeRecord { wtr: self.wtr };
         key.serialize(&mut key_serializer)?;
         self.state = HeaderState::InStructField;
@@ -763,6 +765,7 @@ impl<'a, 'w, W: io::Write> SerializeMap for &'a mut SeHeader<'w, W> {
     }
 
     fn end(self) -> Result<Self::Ok, Self::Error> {
+        self.wtr.on_map_end();
         Ok(())
     }
 }
@@ -856,6 +859,18 @@ mod tests {
         s.serialize(&mut SeHeader::new(&mut wtr)).unwrap_err()
     }
 
+    #[derive(Debug)]
+    struct CustomOrderMap(Vec<(&'static str, f64)>);
+
+    impl Serialize for CustomOrderMap {
+        fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+        where
+            S: serde::Serializer,
+        {
+            serializer.collect_map(self.0.iter().copied())
+        }
+    }
+
     #[test]
     fn bool() {
         let got = serialize(true);
@@ -1126,6 +1141,40 @@ mod tests {
         }
     }
 
+    #[test]
+    fn ordered_map() {
+        let mut map = BTreeMap::new();
+        map.insert("a", 2.0);
+        map.insert("b", 1.0);
+
+        let got = serialize(&map);
+        assert_eq!(got, "2.0,1.0\n");
+        let (wrote, got) = serialize_header(map);
+        assert!(wrote);
+        assert_eq!(got, "a,b");
+    }
+
+    #[test]
+    fn unordered_map() {
+        let mut writer = Writer::from_writer(vec![]);
+        writer
+            .serialize(CustomOrderMap(vec![("a", 2.0), ("b", 1.0)]))
+            .unwrap();
+        writer
+            .serialize(CustomOrderMap(vec![("a", 3.0), ("b", 4.0)]))
+            .unwrap();
+        writer.flush().unwrap();
+        let csv = String::from_utf8(writer.get_ref().clone()).unwrap();
+        assert_eq!(csv, "a,b\n2.0,1.0\n3.0,4.0\n");
+        let error = writer
+            .serialize(CustomOrderMap(vec![("b", 2.0), ("a", 1.0)])) // Wrong key order
+            .unwrap_err();
+        assert!(
+            matches!(error.kind(), ErrorKind::Serialize(_)),
+            "Got unexpected error: {error}"
+        )
+    }
+
     #[test]
     fn struct_no_headers() {
         #[derive(Serialize)]
@@ -1391,6 +1440,36 @@ mod tests {
         assert_eq!(got, format!("x,y,extra1,extra2"));
     }
 
+    #[test]
+    fn flatten_map_with_different_key_order() {
+        #[derive(Serialize, Debug)]
+        struct Row {
+            x: f64,
+            y: f64,
+            #[serde(flatten)]
+            extra: CustomOrderMap,
+        }
+        let mut writer = Writer::from_writer(vec![]);
+        writer
+            .serialize(Row {
+                x: 1.0,
+                y: 2.0,
+                extra: CustomOrderMap(vec![("extra1", 3.0), ("extra2", 4.0)]),
+            })
+            .unwrap();
+        let error = writer
+            .serialize(Row {
+                x: 1.0,
+                y: 2.0,
+                extra: CustomOrderMap(vec![("extra2", 4.0), ("extra1", 3.0)]),
+            })
+            .unwrap_err();
+        assert!(
+            matches!(error.kind(), ErrorKind::Serialize(_)),
+            "Expected ErrorKind::Serialize but got '{error}'"
+        );
+    }
+
     #[test]
     fn flatten_map_different_num_entries() {
         #[derive(Clone, Serialize, Debug, PartialEq)]
@@ -1400,20 +1479,20 @@ mod tests {
             #[serde(flatten)]
             extra: BTreeMap<&'static str, f64>,
         }
-        let mut wtr = Writer::from_writer(vec![]);
+        let mut writer = Writer::from_writer(vec![]);
 
         let mut extra = BTreeMap::new();
         extra.insert("extra1", 3.0);
         extra.insert("extra2", 4.0);
         let row = Row { x: 1.0, y: 2.0, extra };
-        wtr.serialize(row).unwrap();
+        writer.serialize(row).unwrap();
 
         let mut extra = BTreeMap::new();
-        extra.insert("extra3", 3.0);
-        extra.insert("extra4", 4.0);
-        extra.insert("extra5", 5.0);
+        extra.insert("extra1", 3.0);
+        extra.insert("extra2", 4.0);
+        extra.insert("extra3", 5.0);
         let row = Row { x: 1.0, y: 2.0, extra };
-        let error = wtr.serialize(row).unwrap_err();
+        let error = writer.serialize(row).unwrap_err();
         match *error.kind() {
             ErrorKind::UnequalLengths {
                 pos: None,

src/writer.rs

@@ -540,6 +540,15 @@ pub struct Writer<W: io::Write> {
     state: WriterState,
 }
 
+/// State for tracking headers while writing.
+#[derive(Debug)]
+struct HeaderTrackingState {
+    /// The serialized headers in their expected order.
+    expected_headers: Vec<Vec<u8>>,
+    /// The index into the `expected_headers` list of the next expected header.
+    next_expected_index: usize,
+}
+
 #[derive(Debug)]
 struct WriterState {
     /// Whether the Serde serializer should attempt to write a header row.
@@ -557,6 +566,9 @@ struct WriterState {
     /// immediately after flushing the buffer. This avoids flushing the buffer
     /// twice if the inner writer panics.
     panicked: bool,
+    /// Header tracking state for map like data, to ensure that column order
+    /// is preserved across all rows.
+    header_tracking: Option<HeaderTrackingState>,
 }
 
 /// HeaderState encodes a small state machine for handling header writes.
@@ -638,6 +650,10 @@ impl<W: io::Write> Writer<W> {
                 first_field_count: None,
                 fields_written: 0,
                 panicked: false,
+                header_tracking: Some(HeaderTrackingState {
+                    expected_headers: Vec::new(),
+                    next_expected_index: 0,
+                }),
             },
         }
     }
@@ -1180,6 +1196,47 @@ impl<W: io::Write> Writer<W> {
         }
         Ok(())
     }
+
+    /// Track the `key` of a map entry. If this is not the first row, also verify that the
+    /// `key` matches the expected next key.
+    pub(crate) fn check_map_key<T: ?Sized + Serialize>(
+        &mut self,
+        key: &T,
+    ) -> Result<()> {
+        let Some(tracking) = &mut self.state.header_tracking else {
+            return Ok(());
+        };
+        let mut encoded_key_serializer = Writer::from_writer(Vec::new());
+        serialize(&mut encoded_key_serializer, key)?;
+        let encoded_key =
+            encoded_key_serializer.into_inner().map_err(|error| {
+                Error::new(ErrorKind::Serialize(format!(
+                    "Failed to serialize key to bytes: {error:?}"
+                )))
+            })?;
+        if let Some(expected_key) =
+            tracking.expected_headers.get(tracking.next_expected_index)
+        {
+            if expected_key != &encoded_key {
+                return Err(Error::new(ErrorKind::Serialize(format!(
+                    "Out of order key `{}`",
+                    String::from_utf8_lossy(&encoded_key)
+                ))));
+            }
+        } else {
+            // Even if this is not the first row, accept more keys. If the writer is flexible then adding more fields is allowed.
+            tracking.expected_headers.push(encoded_key);
+        }
+        tracking.next_expected_index += 1;
+        Ok(())
+    }
+
+    /// Reset the map key tracking at the end of a row.
+    pub(crate) fn on_map_end(&mut self) {
+        if let Some(tracking) = &mut self.state.header_tracking {
+            tracking.next_expected_index = 0;
+        }
+    }
 }
 
 impl Buffer {