arrow_schema/
metadata.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5use std::collections::{BTreeMap, HashMap, btree_map};
6use std::fmt;
7use std::ops::Index;
8use std::sync::Arc;
9
10/// A cheaply cloneable map of key-value metadata, used by
11/// [`Field`](crate::Field) and [`Schema`](crate::Schema).
12///
13/// Cloning a `Metadata` is always cheap, as the underlying map is
14/// reference-counted. Mutating a shared `Metadata` (e.g. via
15/// [`Metadata::insert`]) will clone the underlying map if (and only if)
16/// it is shared (copy-on-write).
17///
18/// The entries are stored in a [`BTreeMap`], so iteration order is
19/// deterministic (sorted by key).
20///
21/// # Example
22/// ```
23/// # use arrow_schema::Metadata;
24/// let mut metadata = Metadata::new();
25/// metadata.insert("key", "value");
26///
27/// let clone = metadata.clone(); // cheap
28/// assert_eq!(clone.get("key"), Some(&"value".to_string()));
29///
30/// // Mutating one does not affect the other:
31/// metadata.insert("key2", "value2");
32/// assert_eq!(metadata.len(), 2);
33/// assert_eq!(clone.len(), 1);
34/// ```
35#[derive(Clone, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
36pub struct Metadata(
37    // Invariant: the inner map is never empty (`None` encodes the empty map).
38    // This ensures the derived implementations of `PartialEq`, `Ord`, `Hash`, …
39    // treat an empty `Metadata` consistently, and the empty case never allocates.
40    // We use `BTreeMap` for deterministic iteration order.
41    Option<Arc<BTreeMap<String, String>>>,
42);
43
44impl Metadata {
45    /// Creates an empty [`Metadata`].
46    ///
47    /// This does not allocate.
48    pub const fn new() -> Self {
49        Self(None)
50    }
51
52    /// Returns the number of entries.
53    pub fn len(&self) -> usize {
54        self.0.as_ref().map_or(0, |map| map.len())
55    }
56
57    /// Returns `true` if there are no entries.
58    pub fn is_empty(&self) -> bool {
59        self.0.is_none()
60    }
61
62    /// Returns a reference to the value corresponding to the key.
63    pub fn get(&self, key: &str) -> Option<&String> {
64        self.0.as_ref()?.get(key)
65    }
66
67    /// Returns `true` if the map contains a value for the specified key.
68    pub fn contains_key(&self, key: &str) -> bool {
69        self.0.as_ref().is_some_and(|map| map.contains_key(key))
70    }
71
72    /// Inserts a key-value pair, returning the old value of the key, if any.
73    ///
74    /// If the map already contained this key, the value is replaced.
75    ///
76    /// Clones the underlying map if (and only if) it is shared.
77    pub fn insert(&mut self, key: impl Into<String>, value: impl Into<String>) -> Option<String> {
78        Arc::make_mut(self.0.get_or_insert_default()).insert(key.into(), value.into())
79    }
80
81    /// Inserts a key-value pair and returns `self`, for builder-style chaining.
82    ///
83    /// If the map already contained this key, the value is replaced.
84    ///
85    /// # Example
86    /// ```
87    /// # use arrow_schema::Metadata;
88    /// let metadata = Metadata::new().with("a", "1").with("b", "2");
89    /// assert_eq!(metadata.len(), 2);
90    /// ```
91    #[must_use]
92    pub fn with(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
93        self.insert(key, value);
94        self
95    }
96
97    /// Removes a key from the map, returning the value at the key
98    /// if the key was previously in the map.
99    ///
100    /// Clones the underlying map if (and only if) it is shared and contains the key.
101    pub fn remove(&mut self, key: &str) -> Option<String> {
102        let map = self.0.as_mut()?;
103        if !map.contains_key(key) {
104            return None;
105        }
106        let removed = Arc::make_mut(map).remove(key);
107        if map.is_empty() {
108            self.0 = None;
109        }
110        removed
111    }
112
113    /// Removes all entries.
114    pub fn clear(&mut self) {
115        self.0 = None;
116    }
117
118    /// Returns a reference to the underlying [`Arc`], or `None` if the map is empty.
119    ///
120    /// Useful for pointer-based deduplication (e.g. memory accounting) where
121    /// callers need to detect clones that share the same allocation.
122    pub fn as_arc(&self) -> Option<&Arc<BTreeMap<String, String>>> {
123        self.0.as_ref()
124    }
125
126    /// Consumes `self` and returns the underlying [`Arc`], or `None` if the map is empty.
127    ///
128    /// Useful for zero-copy conversion into wrapper types that also store an
129    /// `Arc<BTreeMap<String, String>>` internally.
130    pub fn into_arc(self) -> Option<Arc<BTreeMap<String, String>>> {
131        self.0
132    }
133
134    /// Returns an iterator over the entries, sorted by key.
135    pub fn iter(&self) -> MetadataIter<'_> {
136        self.0
137            .as_deref()
138            .map(|map| map.iter())
139            .into_iter()
140            .flatten()
141    }
142
143    /// Returns an iterator over the keys, in sorted order.
144    pub fn keys(&self) -> impl Iterator<Item = &String> {
145        self.iter().map(|(key, _)| key)
146    }
147
148    /// Returns an iterator over the values, sorted by key.
149    pub fn values(&self) -> impl Iterator<Item = &String> {
150        self.iter().map(|(_, value)| value)
151    }
152
153    /// Retains only the entries for which `f(&key, &mut value)` returns `true`.
154    ///
155    /// Entries for which `f` returns `false` are removed. Retained entries
156    /// remain in their original (sorted) order.
157    ///
158    /// Clones the underlying map if (and only if) it is shared.
159    pub fn retain<F>(&mut self, mut f: F)
160    where
161        F: FnMut(&String, &mut String) -> bool,
162    {
163        let Some(map) = self.0.as_mut() else { return };
164        Arc::make_mut(map).retain(&mut f);
165        if map.is_empty() {
166            self.0 = None;
167        }
168    }
169}
170
171/// Iterator over the entries of a [`Metadata`], sorted by key.
172pub type MetadataIter<'a> =
173    std::iter::Flatten<std::option::IntoIter<btree_map::Iter<'a, String, String>>>;
174
175impl fmt::Debug for Metadata {
176    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
177        f.debug_map().entries(self.iter()).finish()
178    }
179}
180
181impl Index<&str> for Metadata {
182    type Output = String;
183
184    fn index(&self, key: &str) -> &String {
185        self.get(key)
186            .unwrap_or_else(|| panic!("no entry found for key {key:?}"))
187    }
188}
189
190impl From<BTreeMap<String, String>> for Metadata {
191    fn from(map: BTreeMap<String, String>) -> Self {
192        if map.is_empty() {
193            Self(None)
194        } else {
195            Self(Some(Arc::new(map)))
196        }
197    }
198}
199
200impl From<HashMap<String, String>> for Metadata {
201    fn from(map: HashMap<String, String>) -> Self {
202        map.into_iter().collect::<BTreeMap<_, _>>().into()
203    }
204}
205
206impl From<Arc<BTreeMap<String, String>>> for Metadata {
207    fn from(map: Arc<BTreeMap<String, String>>) -> Self {
208        if map.is_empty() {
209            Self(None)
210        } else {
211            Self(Some(map))
212        }
213    }
214}
215
216impl<K: Into<String>, V: Into<String>, const N: usize> From<[(K, V); N]> for Metadata {
217    fn from(entries: [(K, V); N]) -> Self {
218        entries.into_iter().collect()
219    }
220}
221
222impl From<Metadata> for BTreeMap<String, String> {
223    fn from(metadata: Metadata) -> Self {
224        match metadata.0 {
225            None => BTreeMap::new(),
226            Some(map) => Arc::try_unwrap(map).unwrap_or_else(|map| (*map).clone()),
227        }
228    }
229}
230
231impl From<Metadata> for HashMap<String, String> {
232    fn from(metadata: Metadata) -> Self {
233        metadata.into_iter().collect()
234    }
235}
236
237impl From<&Metadata> for HashMap<String, String> {
238    fn from(metadata: &Metadata) -> Self {
239        metadata
240            .iter()
241            .map(|(k, v)| (k.clone(), v.clone()))
242            .collect()
243    }
244}
245
246impl<K: Into<String>, V: Into<String>> FromIterator<(K, V)> for Metadata {
247    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
248        iter.into_iter()
249            .map(|(k, v)| (k.into(), v.into()))
250            .collect::<BTreeMap<_, _>>()
251            .into()
252    }
253}
254
255impl<K: Into<String>, V: Into<String>> Extend<(K, V)> for Metadata {
256    fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) {
257        let mut iter = iter
258            .into_iter()
259            .map(|(k, v)| (k.into(), v.into()))
260            .peekable();
261        if iter.peek().is_none() {
262            return; // Avoid cloning a shared map for an empty iterator
263        }
264        Arc::make_mut(self.0.get_or_insert_default()).extend(iter);
265    }
266}
267
268impl<'a> IntoIterator for &'a Metadata {
269    type Item = (&'a String, &'a String);
270    type IntoIter = MetadataIter<'a>;
271
272    fn into_iter(self) -> Self::IntoIter {
273        self.iter()
274    }
275}
276
277impl IntoIterator for Metadata {
278    type Item = (String, String);
279    type IntoIter = btree_map::IntoIter<String, String>;
280
281    fn into_iter(self) -> Self::IntoIter {
282        BTreeMap::from(self).into_iter()
283    }
284}
285
286impl PartialEq<HashMap<String, String>> for Metadata {
287    fn eq(&self, other: &HashMap<String, String>) -> bool {
288        self.len() == other.len()
289            && self
290                .iter()
291                .all(|(k, v)| other.get(k).is_some_and(|other_v| v == other_v))
292    }
293}
294
295impl PartialEq<BTreeMap<String, String>> for Metadata {
296    fn eq(&self, other: &BTreeMap<String, String>) -> bool {
297        self.iter().eq(other.iter())
298    }
299}
300
301#[cfg(feature = "serde")]
302mod serde_impl {
303    use super::Metadata;
304    use std::collections::BTreeMap;
305
306    impl serde_core::Serialize for Metadata {
307        fn serialize<S: serde_core::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
308            use serde_core::ser::SerializeMap as _;
309            let mut map = serializer.serialize_map(Some(self.len()))?;
310            for (key, value) in self {
311                map.serialize_entry(key, value)?;
312            }
313            map.end()
314        }
315    }
316
317    impl<'de> serde_core::Deserialize<'de> for Metadata {
318        fn deserialize<D: serde_core::Deserializer<'de>>(
319            deserializer: D,
320        ) -> Result<Self, D::Error> {
321            Ok(BTreeMap::<String, String>::deserialize(deserializer)?.into())
322        }
323    }
324}
325
326#[cfg(test)]
327mod tests {
328    use super::*;
329
330    #[test]
331    fn test_empty() {
332        let metadata = Metadata::new();
333        assert!(metadata.is_empty());
334        assert_eq!(metadata.len(), 0);
335        assert_eq!(metadata.get("key"), None);
336        assert!(!metadata.contains_key("key"));
337        assert_eq!(metadata.iter().count(), 0);
338        assert_eq!(metadata, Metadata::default());
339
340        // Empty maps don't allocate:
341        assert!(Metadata::from(HashMap::new()).0.is_none());
342        assert!(Metadata::from(BTreeMap::new()).0.is_none());
343        assert!(
344            std::iter::empty::<(String, String)>()
345                .collect::<Metadata>()
346                .0
347                .is_none()
348        );
349    }
350
351    #[test]
352    fn test_insert_get_remove() {
353        let mut metadata = Metadata::new();
354        assert_eq!(metadata.insert("a", "1"), None);
355        assert_eq!(metadata.insert("a", "2"), Some("1".to_string()));
356        assert_eq!(metadata.insert("b", "3"), None);
357
358        assert_eq!(metadata.len(), 2);
359        assert_eq!(metadata.get("a"), Some(&"2".to_string()));
360        assert_eq!(metadata["b"], "3");
361        assert!(metadata.contains_key("a"));
362        assert!(!metadata.contains_key("c"));
363
364        assert_eq!(metadata.remove("c"), None);
365        assert_eq!(metadata.remove("a"), Some("2".to_string()));
366        assert_eq!(metadata.remove("b"), Some("3".to_string()));
367
368        // Invariant: removing the last entry restores the unallocated state
369        assert!(metadata.is_empty());
370        assert!(metadata.0.is_none());
371    }
372
373    #[test]
374    fn test_clear() {
375        let mut metadata = Metadata::from([("a", "1")]);
376        metadata.clear();
377        assert!(metadata.is_empty());
378        assert!(metadata.0.is_none());
379    }
380
381    #[test]
382    #[should_panic(expected = "no entry found for key \"missing\"")]
383    fn test_index_panics() {
384        let metadata = Metadata::from([("a", "1")]);
385        let _ = &metadata["missing"];
386    }
387
388    #[test]
389    fn test_copy_on_write() {
390        let mut metadata = Metadata::from([("a", "1")]);
391        let clone = metadata.clone();
392
393        // Cloning is shallow:
394        let arc = metadata.0.as_ref().expect("non-empty");
395        assert!(Arc::ptr_eq(arc, clone.0.as_ref().expect("non-empty")));
396
397        // Mutation clones the shared map, leaving the clone untouched:
398        metadata.insert("b", "2");
399        assert_eq!(metadata.len(), 2);
400        assert_eq!(clone.len(), 1);
401
402        // Mutating an unshared map does not clone it:
403        let arc = metadata.0.as_ref().expect("non-empty").clone();
404        metadata.insert("c", "3");
405        assert!(
406            Arc::ptr_eq(&arc, metadata.0.as_ref().expect("non-empty"))
407                || Arc::strong_count(&arc) == 1
408        );
409
410        // Removing a missing key from a shared map does not clone it:
411        let clone = metadata.clone();
412        let arc = metadata.0.as_ref().expect("non-empty");
413        assert!(Arc::ptr_eq(arc, clone.0.as_ref().expect("non-empty")));
414        let mut metadata2 = metadata.clone();
415        assert_eq!(metadata2.remove("missing"), None);
416        assert!(Arc::ptr_eq(
417            metadata.0.as_ref().expect("non-empty"),
418            metadata2.0.as_ref().expect("non-empty")
419        ));
420    }
421
422    #[test]
423    fn test_deterministic_iteration_order() {
424        let metadata: Metadata = [("b", "2"), ("a", "1"), ("c", "3")].into_iter().collect();
425        let keys: Vec<&String> = metadata.keys().collect();
426        assert_eq!(keys, ["a", "b", "c"]);
427        let values: Vec<&String> = metadata.values().collect();
428        assert_eq!(values, ["1", "2", "3"]);
429    }
430
431    #[test]
432    fn test_eq_with_std_maps() {
433        let metadata = Metadata::from([("a", "1"), ("b", "2")]);
434
435        let hash_map: HashMap<String, String> = metadata.clone().into();
436        assert_eq!(metadata, hash_map);
437
438        let btree_map: BTreeMap<String, String> = metadata.clone().into();
439        assert_eq!(metadata, btree_map);
440
441        let mut different = hash_map.clone();
442        different.insert("c".to_string(), "3".to_string());
443        assert_ne!(metadata, different);
444
445        let mut different = hash_map;
446        different.insert("a".to_string(), "other".to_string());
447        assert_ne!(metadata, different);
448    }
449
450    #[test]
451    fn test_extend() {
452        let mut metadata = Metadata::new();
453        let clone = metadata.clone();
454        metadata.extend(std::iter::empty::<(String, String)>());
455        assert!(metadata.0.is_none()); // No allocation for an empty extend
456
457        metadata.extend([("a", "1"), ("b", "2")]);
458        assert_eq!(metadata.len(), 2);
459        assert!(clone.is_empty());
460    }
461
462    #[test]
463    fn test_as_arc_into_arc() {
464        let empty = Metadata::new();
465        assert!(empty.as_arc().is_none());
466        assert!(empty.into_arc().is_none());
467
468        // Non-empty metadata exposes the same Arc via as_arc / into_arc
469        let metadata = Metadata::from([("a", "1")]);
470        let arc_ref = metadata.as_arc().expect("non-empty");
471        assert_eq!(arc_ref.get("a").map(String::as_str), Some("1"));
472
473        let clone = metadata.clone();
474        // Clones share the same allocation
475        assert!(Arc::ptr_eq(
476            metadata.as_arc().unwrap(),
477            clone.as_arc().unwrap()
478        ));
479
480        let arc = metadata.into_arc().expect("non-empty");
481        assert_eq!(arc.get("a").map(String::as_str), Some("1"));
482        // into_arc gives back the same Arc (no copy)
483        assert!(Arc::ptr_eq(&arc, clone.as_arc().unwrap()));
484    }
485
486    #[test]
487    fn test_debug() {
488        let metadata = Metadata::from([("b", "2"), ("a", "1")]);
489        assert_eq!(format!("{metadata:?}"), r#"{"a": "1", "b": "2"}"#);
490        assert_eq!(format!("{:?}", Metadata::new()), "{}");
491    }
492
493    #[test]
494    fn test_retain() {
495        let mut metadata =
496            Metadata::from([("a", "1"), ("b", "2"), ("c", "3"), ("d", "4"), ("e", "5")]);
497
498        metadata.retain(|k, _| -> bool { k >= &String::from("c") });
499
500        let result_map = Metadata::from([("c", "3"), ("d", "4"), ("e", "5")]);
501        assert_eq!(metadata, result_map)
502    }
503
504    #[test]
505    fn test_retain_empty() {
506        let mut metadata = Metadata::new();
507        metadata.retain(|_, _| -> bool { true });
508        assert!(metadata.is_empty())
509    }
510
511    #[test]
512    fn test_retain_all_removed() {
513        let mut metadata = Metadata::from([("a", "1"), ("b", "2"), ("c", "3")]);
514        metadata.retain(|_, _| false);
515        assert!(metadata.is_empty());
516        assert!(metadata.0.is_none());
517    }
518
519    #[test]
520    fn test_retain_copy_on_write() {
521        let mut metadata = Metadata::from([("a", "1"), ("b", "2"), ("c", "3")]);
522        let clone = metadata.clone();
523
524        // both share the same Arc before any mutation:
525        assert!(Arc::ptr_eq(
526            metadata.0.as_ref().expect("non-empty"),
527            clone.0.as_ref().expect("non-empty"),
528        ));
529
530        // retain clones the shared map, leaving the clone untouched:
531        metadata.retain(|k, _| k != "a");
532        assert_eq!(metadata.len(), 2);
533        assert_eq!(clone.len(), 3);
534        assert!(!metadata.contains_key("a"));
535        assert!(clone.contains_key("a"));
536    }
537
538    #[test]
539    #[cfg(feature = "serde")]
540    fn test_serde_round_trip() {
541        for metadata in [Metadata::new(), Metadata::from([("a", "1"), ("b", "2")])] {
542            let serialized = postcard::to_stdvec(&metadata).expect("serialize");
543            let deserialized: Metadata = postcard::from_bytes(&serialized).expect("deserialize");
544            assert_eq!(metadata, deserialized);
545        }
546    }
547
548    #[test]
549    #[cfg(feature = "serde")]
550    fn test_serde_matches_std_map_format() {
551        // `Metadata` must serialize exactly like the maps it replaced
552        let metadata = Metadata::from([("a", "1"), ("b", "2")]);
553        let as_btree_map: BTreeMap<String, String> = metadata.clone().into();
554
555        let serialized = postcard::to_stdvec(&metadata).expect("serialize");
556        let map_serialized = postcard::to_stdvec(&as_btree_map).expect("serialize");
557        assert_eq!(serialized, map_serialized);
558
559        let deserialized: Metadata = postcard::from_bytes(&map_serialized).expect("deserialize");
560        assert_eq!(metadata, deserialized);
561    }
562}
arrow_schema/metadata.rs

arrow_schema/
metadata.rs
