1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
use std::collections::HashMap;

use arrow::datatypes::Field as ArrowField;

/// Arrow metadata for an arrow record batch.
pub type ArrowBatchMetadata = HashMap<String, String>;

/// Arrow metadata for a column/field.
pub type ArrowFieldMetadata = HashMap<String, String>;

#[derive(thiserror::Error, Debug)]
#[error("Missing metadata {key:?}")]
pub struct MissingMetadataKey {
    pub key: String,
}

#[derive(thiserror::Error, Debug)]
#[error("Field {field_name:?} is missing metadata {metadata_key:?}")]
pub struct MissingFieldMetadata {
    pub field_name: String,
    pub metadata_key: String,
}

/// Make it more ergonomic to work with arrow metadata.
pub trait MetadataExt {
    type Error;

    fn missing_key_error(&self, key: &str) -> Self::Error;
    fn get_opt(&self, key: &str) -> Option<&str>;

    fn get_or_err(&self, key: &str) -> Result<&str, Self::Error> {
        self.get_opt(key).ok_or_else(|| self.missing_key_error(key))
    }

    fn get_bool(&self, key: &str) -> bool {
        self.get_opt(key)
            .map(|value| !matches!(value.to_lowercase().as_str(), "false" | "no"))
            .unwrap_or(false)
    }
}

impl MetadataExt for HashMap<String, String> {
    type Error = MissingMetadataKey;

    fn missing_key_error(&self, key: &str) -> Self::Error {
        MissingMetadataKey {
            key: key.to_owned(),
        }
    }

    fn get_opt(&self, key: &str) -> Option<&str> {
        self.get(key).map(|value| value.as_str())
    }
}

impl MetadataExt for ArrowField {
    type Error = MissingFieldMetadata;

    fn missing_key_error(&self, key: &str) -> Self::Error {
        MissingFieldMetadata {
            field_name: self.name().clone(),
            metadata_key: key.to_owned(),
        }
    }

    fn get_opt(&self, key: &str) -> Option<&str> {
        self.metadata().get(key).map(|v| v.as_str())
    }
}