wp_connector_utils/arrow/
format.rs

1//! Wire format selection for Arrow IPC payloads.
2//!
3//! - [`WireFormat`] — enum + `parse_strict` / `from_data_format` / `is_arrow`
4
5// ---------------------------------------------------------------------------
6// WireFormat
7// ---------------------------------------------------------------------------
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
10pub enum WireFormat {
11    #[default]
12    Ndjson,
13    ArrowStream,
14    ArrowFramed,
15}
16
17pub const SUPPORTED_DATA_FORMATS: &[&str] = &["ndjson", "arrow_ipc", "arrow_framed"];
18
19impl WireFormat {
20    pub fn from_data_format(value: Option<&str>) -> Self {
21        match value.unwrap_or("ndjson") {
22            "arrow_framed" => WireFormat::ArrowFramed,
23            "arrow_ipc" => WireFormat::ArrowStream,
24            _ => WireFormat::Ndjson,
25        }
26    }
27
28    pub fn parse_strict(value: Option<&str>) -> Result<Self, String> {
29        match value {
30            None | Some("ndjson") => Ok(WireFormat::Ndjson),
31            Some("arrow_ipc") => Ok(WireFormat::ArrowStream),
32            Some("arrow_framed") => Ok(WireFormat::ArrowFramed),
33            Some(raw) => Err(format!(
34                "data_format must be one of: {} (got '{raw}')",
35                SUPPORTED_DATA_FORMATS.join(", ")
36            )),
37        }
38    }
39
40    pub fn is_arrow(self) -> bool {
41        matches!(self, WireFormat::ArrowStream | WireFormat::ArrowFramed)
42    }
43}
44
45#[cfg(test)]
46mod tests {
47    use super::*;
48
49    #[test]
50    fn defaults_to_ndjson() {
51        assert_eq!(WireFormat::from_data_format(None), WireFormat::Ndjson);
52        assert_eq!(WireFormat::default(), WireFormat::Ndjson);
53    }
54    #[test]
55    fn parses_arrow() {
56        assert_eq!(
57            WireFormat::from_data_format(Some("arrow_ipc")),
58            WireFormat::ArrowStream
59        );
60        assert_eq!(
61            WireFormat::from_data_format(Some("arrow_framed")),
62            WireFormat::ArrowFramed
63        );
64    }
65    #[test]
66    fn unknown_falls_back() {
67        assert_eq!(WireFormat::from_data_format(Some("x")), WireFormat::Ndjson);
68    }
69    #[test]
70    fn parse_strict_rejects_unknown() {
71        assert!(WireFormat::parse_strict(Some("x")).is_err());
72    }
73    #[test]
74    fn parse_strict_accepts_known() {
75        assert_eq!(WireFormat::parse_strict(None).unwrap(), WireFormat::Ndjson);
76        assert_eq!(
77            WireFormat::parse_strict(Some("ndjson")).unwrap(),
78            WireFormat::Ndjson
79        );
80        assert_eq!(
81            WireFormat::parse_strict(Some("arrow_ipc")).unwrap(),
82            WireFormat::ArrowStream
83        );
84        assert_eq!(
85            WireFormat::parse_strict(Some("arrow_framed")).unwrap(),
86            WireFormat::ArrowFramed
87        );
88    }
89    #[test]
90    fn is_arrow_flag() {
91        assert!(!WireFormat::Ndjson.is_arrow());
92        assert!(WireFormat::ArrowStream.is_arrow());
93        assert!(WireFormat::ArrowFramed.is_arrow());
94    }
95}
wp_connector_utils/arrow/format.rs

wp_connector_utils/arrow/
format.rs
