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 clonable 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 an iterator over the entries, sorted by key.
132    pub fn iter(&self) -> MetadataIter<'_> {
133        self.0
134            .as_deref()
135            .map(|map| map.iter())
136            .into_iter()
137            .flatten()
138    }
139
140    /// Returns an iterator over the keys, in sorted order.
141    pub fn keys(&self) -> impl Iterator<Item = &String> {
142        self.iter().map(|(key, _)| key)
143    }
144
145    /// Returns an iterator over the values, sorted by key.
146    pub fn values(&self) -> impl Iterator<Item = &String> {
147        self.iter().map(|(_, value)| value)
148    }
149}
150
151/// Iterator over the entries of a [`Metadata`], sorted by key.
152pub type MetadataIter<'a> =
153    std::iter::Flatten<std::option::IntoIter<btree_map::Iter<'a, String, String>>>;
154
155impl fmt::Debug for Metadata {
156    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
157        f.debug_map().entries(self.iter()).finish()
158    }
159}
160
161impl Index<&str> for Metadata {
162    type Output = String;
163
164    fn index(&self, key: &str) -> &String {
165        self.get(key)
166            .unwrap_or_else(|| panic!("no entry found for key {key:?}"))
167    }
168}
169
170impl From<BTreeMap<String, String>> for Metadata {
171    fn from(map: BTreeMap<String, String>) -> Self {
172        if map.is_empty() {
173            Self(None)
174        } else {
175            Self(Some(Arc::new(map)))
176        }
177    }
178}
179
180impl From<HashMap<String, String>> for Metadata {
181    fn from(map: HashMap<String, String>) -> Self {
182        map.into_iter().collect::<BTreeMap<_, _>>().into()
183    }
184}
185
186impl From<Arc<BTreeMap<String, String>>> for Metadata {
187    fn from(map: Arc<BTreeMap<String, String>>) -> Self {
188        if map.is_empty() {
189            Self(None)
190        } else {
191            Self(Some(map))
192        }
193    }
194}
195
196impl<K: Into<String>, V: Into<String>, const N: usize> From<[(K, V); N]> for Metadata {
197    fn from(entries: [(K, V); N]) -> Self {
198        entries.into_iter().collect()
199    }
200}
201
202impl From<Metadata> for BTreeMap<String, String> {
203    fn from(metadata: Metadata) -> Self {
204        match metadata.0 {
205            None => BTreeMap::new(),
206            Some(map) => Arc::try_unwrap(map).unwrap_or_else(|map| (*map).clone()),
207        }
208    }
209}
210
211impl From<Metadata> for HashMap<String, String> {
212    fn from(metadata: Metadata) -> Self {
213        metadata.into_iter().collect()
214    }
215}
216
217impl From<&Metadata> for HashMap<String, String> {
218    fn from(metadata: &Metadata) -> Self {
219        metadata
220            .iter()
221            .map(|(k, v)| (k.clone(), v.clone()))
222            .collect()
223    }
224}
225
226impl<K: Into<String>, V: Into<String>> FromIterator<(K, V)> for Metadata {
227    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
228        iter.into_iter()
229            .map(|(k, v)| (k.into(), v.into()))
230            .collect::<BTreeMap<_, _>>()
231            .into()
232    }
233}
234
235impl<K: Into<String>, V: Into<String>> Extend<(K, V)> for Metadata {
236    fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) {
237        let mut iter = iter
238            .into_iter()
239            .map(|(k, v)| (k.into(), v.into()))
240            .peekable();
241        if iter.peek().is_none() {
242            return; // Avoid cloning a shared map for an empty iterator
243        }
244        Arc::make_mut(self.0.get_or_insert_default()).extend(iter);
245    }
246}
247
248impl<'a> IntoIterator for &'a Metadata {
249    type Item = (&'a String, &'a String);
250    type IntoIter = MetadataIter<'a>;
251
252    fn into_iter(self) -> Self::IntoIter {
253        self.iter()
254    }
255}
256
257impl IntoIterator for Metadata {
258    type Item = (String, String);
259    type IntoIter = btree_map::IntoIter<String, String>;
260
261    fn into_iter(self) -> Self::IntoIter {
262        BTreeMap::from(self).into_iter()
263    }
264}
265
266impl PartialEq<HashMap<String, String>> for Metadata {
267    fn eq(&self, other: &HashMap<String, String>) -> bool {
268        self.len() == other.len()
269            && self
270                .iter()
271                .all(|(k, v)| other.get(k).is_some_and(|other_v| v == other_v))
272    }
273}
274
275impl PartialEq<BTreeMap<String, String>> for Metadata {
276    fn eq(&self, other: &BTreeMap<String, String>) -> bool {
277        self.iter().eq(other.iter())
278    }
279}
280
281#[cfg(feature = "serde")]
282mod serde_impl {
283    use super::Metadata;
284    use std::collections::BTreeMap;
285
286    impl serde_core::Serialize for Metadata {
287        fn serialize<S: serde_core::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
288            use serde_core::ser::SerializeMap as _;
289            let mut map = serializer.serialize_map(Some(self.len()))?;
290            for (key, value) in self {
291                map.serialize_entry(key, value)?;
292            }
293            map.end()
294        }
295    }
296
297    impl<'de> serde_core::Deserialize<'de> for Metadata {
298        fn deserialize<D: serde_core::Deserializer<'de>>(
299            deserializer: D,
300        ) -> Result<Self, D::Error> {
301            Ok(BTreeMap::<String, String>::deserialize(deserializer)?.into())
302        }
303    }
304}
305
306#[cfg(test)]
307mod tests {
308    use super::*;
309
310    #[test]
311    fn test_empty() {
312        let metadata = Metadata::new();
313        assert!(metadata.is_empty());
314        assert_eq!(metadata.len(), 0);
315        assert_eq!(metadata.get("key"), None);
316        assert!(!metadata.contains_key("key"));
317        assert_eq!(metadata.iter().count(), 0);
318        assert_eq!(metadata, Metadata::default());
319
320        // Empty maps don't allocate:
321        assert!(Metadata::from(HashMap::new()).0.is_none());
322        assert!(Metadata::from(BTreeMap::new()).0.is_none());
323        assert!(
324            std::iter::empty::<(String, String)>()
325                .collect::<Metadata>()
326                .0
327                .is_none()
328        );
329    }
330
331    #[test]
332    fn test_insert_get_remove() {
333        let mut metadata = Metadata::new();
334        assert_eq!(metadata.insert("a", "1"), None);
335        assert_eq!(metadata.insert("a", "2"), Some("1".to_string()));
336        assert_eq!(metadata.insert("b", "3"), None);
337
338        assert_eq!(metadata.len(), 2);
339        assert_eq!(metadata.get("a"), Some(&"2".to_string()));
340        assert_eq!(metadata["b"], "3");
341        assert!(metadata.contains_key("a"));
342        assert!(!metadata.contains_key("c"));
343
344        assert_eq!(metadata.remove("c"), None);
345        assert_eq!(metadata.remove("a"), Some("2".to_string()));
346        assert_eq!(metadata.remove("b"), Some("3".to_string()));
347
348        // Invariant: removing the last entry restores the unallocated state
349        assert!(metadata.is_empty());
350        assert!(metadata.0.is_none());
351    }
352
353    #[test]
354    fn test_clear() {
355        let mut metadata = Metadata::from([("a", "1")]);
356        metadata.clear();
357        assert!(metadata.is_empty());
358        assert!(metadata.0.is_none());
359    }
360
361    #[test]
362    #[should_panic(expected = "no entry found for key \"missing\"")]
363    fn test_index_panics() {
364        let metadata = Metadata::from([("a", "1")]);
365        let _ = &metadata["missing"];
366    }
367
368    #[test]
369    fn test_copy_on_write() {
370        let mut metadata = Metadata::from([("a", "1")]);
371        let clone = metadata.clone();
372
373        // Cloning is shallow:
374        let arc = metadata.0.as_ref().expect("non-empty");
375        assert!(Arc::ptr_eq(arc, clone.0.as_ref().expect("non-empty")));
376
377        // Mutation clones the shared map, leaving the clone untouched:
378        metadata.insert("b", "2");
379        assert_eq!(metadata.len(), 2);
380        assert_eq!(clone.len(), 1);
381
382        // Mutating an unshared map does not clone it:
383        let arc = metadata.0.as_ref().expect("non-empty").clone();
384        metadata.insert("c", "3");
385        assert!(
386            Arc::ptr_eq(&arc, metadata.0.as_ref().expect("non-empty"))
387                || Arc::strong_count(&arc) == 1
388        );
389
390        // Removing a missing key from a shared map does not clone it:
391        let clone = metadata.clone();
392        let arc = metadata.0.as_ref().expect("non-empty");
393        assert!(Arc::ptr_eq(arc, clone.0.as_ref().expect("non-empty")));
394        let mut metadata2 = metadata.clone();
395        assert_eq!(metadata2.remove("missing"), None);
396        assert!(Arc::ptr_eq(
397            metadata.0.as_ref().expect("non-empty"),
398            metadata2.0.as_ref().expect("non-empty")
399        ));
400    }
401
402    #[test]
403    fn test_deterministic_iteration_order() {
404        let metadata: Metadata = [("b", "2"), ("a", "1"), ("c", "3")].into_iter().collect();
405        let keys: Vec<&String> = metadata.keys().collect();
406        assert_eq!(keys, ["a", "b", "c"]);
407        let values: Vec<&String> = metadata.values().collect();
408        assert_eq!(values, ["1", "2", "3"]);
409    }
410
411    #[test]
412    fn test_eq_with_std_maps() {
413        let metadata = Metadata::from([("a", "1"), ("b", "2")]);
414
415        let hash_map: HashMap<String, String> = metadata.clone().into();
416        assert_eq!(metadata, hash_map);
417
418        let btree_map: BTreeMap<String, String> = metadata.clone().into();
419        assert_eq!(metadata, btree_map);
420
421        let mut different = hash_map.clone();
422        different.insert("c".to_string(), "3".to_string());
423        assert_ne!(metadata, different);
424
425        let mut different = hash_map;
426        different.insert("a".to_string(), "other".to_string());
427        assert_ne!(metadata, different);
428    }
429
430    #[test]
431    fn test_extend() {
432        let mut metadata = Metadata::new();
433        let clone = metadata.clone();
434        metadata.extend(std::iter::empty::<(String, String)>());
435        assert!(metadata.0.is_none()); // No allocation for an empty extend
436
437        metadata.extend([("a", "1"), ("b", "2")]);
438        assert_eq!(metadata.len(), 2);
439        assert!(clone.is_empty());
440    }
441
442    #[test]
443    fn test_debug() {
444        let metadata = Metadata::from([("b", "2"), ("a", "1")]);
445        assert_eq!(format!("{metadata:?}"), r#"{"a": "1", "b": "2"}"#);
446        assert_eq!(format!("{:?}", Metadata::new()), "{}");
447    }
448
449    #[test]
450    #[cfg(feature = "serde")]
451    fn test_serde_round_trip() {
452        for metadata in [Metadata::new(), Metadata::from([("a", "1"), ("b", "2")])] {
453            let serialized = postcard::to_stdvec(&metadata).expect("serialize");
454            let deserialized: Metadata = postcard::from_bytes(&serialized).expect("deserialize");
455            assert_eq!(metadata, deserialized);
456        }
457    }
458
459    #[test]
460    #[cfg(feature = "serde")]
461    fn test_serde_matches_std_map_format() {
462        // `Metadata` must serialize exactly like the maps it replaced
463        let metadata = Metadata::from([("a", "1"), ("b", "2")]);
464        let as_btree_map: BTreeMap<String, String> = metadata.clone().into();
465
466        let serialized = postcard::to_stdvec(&metadata).expect("serialize");
467        let map_serialized = postcard::to_stdvec(&as_btree_map).expect("serialize");
468        assert_eq!(serialized, map_serialized);
469
470        let deserialized: Metadata = postcard::from_bytes(&map_serialized).expect("deserialize");
471        assert_eq!(metadata, deserialized);
472    }
473}