Skip to main content

arrow_schema/
metadata.rs

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