ptars_core/
config.rs

1use std::sync::Arc;
2
3/// Time unit for timestamp, time of day and duration values.
4///
5/// Mirrors `arrow_schema::TimeUnit` so that the config carries no arrow types
6/// and can be re-exported from arrow-version-independent crates.
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
8pub enum TimeUnit {
9    /// Seconds
10    Second,
11    /// Milliseconds
12    Millisecond,
13    /// Microseconds
14    Microsecond,
15    /// Nanoseconds
16    #[default]
17    Nanosecond,
18}
19
20impl From<TimeUnit> for arrow_schema::TimeUnit {
21    fn from(value: TimeUnit) -> Self {
22        match value {
23            TimeUnit::Second => arrow_schema::TimeUnit::Second,
24            TimeUnit::Millisecond => arrow_schema::TimeUnit::Millisecond,
25            TimeUnit::Microsecond => arrow_schema::TimeUnit::Microsecond,
26            TimeUnit::Nanosecond => arrow_schema::TimeUnit::Nanosecond,
27        }
28    }
29}
30
31/// How to represent protobuf enum fields in Arrow.
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
33pub enum EnumRepr {
34    /// Store as Int32 (protobuf enum number). This is the default.
35    #[default]
36    Int32,
37    /// Store as Utf8/LargeUtf8 (enum value name as string).
38    String,
39    /// Store as Binary/LargeBinary (enum value name as bytes).
40    Binary,
41}
42
43/// Policy for stripping Confluent Schema Registry wire format prefix from messages.
44#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
45pub enum ConfluentWirePolicy {
46    /// No prefix stripping — bytes are raw protobuf wire format. This is the default.
47    #[default]
48    Raw,
49    /// Strip 5-byte Confluent header (1 magic byte + 4-byte schema ID).
50    /// Use with Avro and JSON Schema.
51    Standard,
52    /// Strip 5-byte Confluent header plus varint-encoded message index array.
53    /// Use with Protobuf schemas in Confluent Schema Registry.
54    Protobuf,
55}
56
57/// Configuration for protobuf to Arrow conversions.
58///
59/// This struct allows customizing how protobuf types are mapped to Arrow types,
60/// similar to Python's `ProtarrowConfig`.
61#[derive(Debug, Clone)]
62pub struct PtarsConfig {
63    /// Timezone for timestamp values. Default: Some("UTC")
64    /// Set to None for timezone-naive timestamps.
65    pub timestamp_tz: Option<Arc<str>>,
66
67    /// Time unit for timestamp values. Default: Nanosecond
68    pub timestamp_unit: TimeUnit,
69
70    /// Time unit for time of day values. Default: Nanosecond
71    pub time_unit: TimeUnit,
72
73    /// Time unit for duration values. Default: Nanosecond
74    pub duration_unit: TimeUnit,
75
76    /// Name for list item field. Default: "item"
77    pub list_value_name: Arc<str>,
78
79    /// Name for map value field. Default: "value"
80    pub map_value_name: Arc<str>,
81
82    /// Whether list fields can be null. Default: false
83    pub list_nullable: bool,
84
85    /// Whether map fields can be null. Default: false
86    pub map_nullable: bool,
87
88    /// Whether list element values can be null. Default: false
89    pub list_value_nullable: bool,
90
91    /// Whether map values can be null. Default: false
92    pub map_value_nullable: bool,
93
94    /// Whether to use LargeUtf8 instead of Utf8 for string fields. Default: false
95    pub use_large_string: bool,
96
97    /// Whether to use LargeBinary instead of Binary for bytes fields. Default: false
98    pub use_large_binary: bool,
99
100    /// Whether to use LargeList instead of List for repeated fields. Default: false
101    pub use_large_list: bool,
102
103    /// How to represent enum fields in Arrow. Default: Int32
104    /// When String, use_large_string controls Utf8 vs LargeUtf8.
105    /// When Binary, use_large_binary controls Binary vs LargeBinary.
106    pub enum_repr: EnumRepr,
107
108    /// Policy for stripping Confluent Schema Registry wire format prefix. Default: Raw
109    pub confluent_wire_policy: ConfluentWirePolicy,
110}
111
112impl Default for PtarsConfig {
113    fn default() -> Self {
114        Self {
115            timestamp_tz: Some(Arc::from("UTC")),
116            timestamp_unit: TimeUnit::Nanosecond,
117            time_unit: TimeUnit::Nanosecond,
118            duration_unit: TimeUnit::Nanosecond,
119            list_value_name: Arc::from("item"),
120            map_value_name: Arc::from("value"),
121            list_nullable: false,
122            map_nullable: false,
123            list_value_nullable: false,
124            map_value_nullable: false,
125            use_large_string: false,
126            use_large_binary: false,
127            use_large_list: false,
128            enum_repr: EnumRepr::default(),
129            confluent_wire_policy: ConfluentWirePolicy::default(),
130        }
131    }
132}
133
134impl PtarsConfig {
135    /// Create a new config with default values.
136    pub fn new() -> Self {
137        Self::default()
138    }
139
140    /// Set the timezone for timestamp values.
141    pub fn with_timestamp_tz(mut self, tz: Option<&str>) -> Self {
142        self.timestamp_tz = tz.map(Arc::from);
143        self
144    }
145
146    /// Set the time unit for timestamp values.
147    pub fn with_timestamp_unit(mut self, unit: TimeUnit) -> Self {
148        self.timestamp_unit = unit;
149        self
150    }
151
152    /// Set the time unit for time of day values.
153    pub fn with_time_unit(mut self, unit: TimeUnit) -> Self {
154        self.time_unit = unit;
155        self
156    }
157
158    /// Set the time unit for duration values.
159    pub fn with_duration_unit(mut self, unit: TimeUnit) -> Self {
160        self.duration_unit = unit;
161        self
162    }
163
164    /// Set the name for list item fields.
165    pub fn with_list_value_name(mut self, name: &str) -> Self {
166        self.list_value_name = Arc::from(name);
167        self
168    }
169
170    /// Set the name for map value fields.
171    pub fn with_map_value_name(mut self, name: &str) -> Self {
172        self.map_value_name = Arc::from(name);
173        self
174    }
175
176    /// Set whether list fields can be null.
177    pub fn with_list_nullable(mut self, nullable: bool) -> Self {
178        self.list_nullable = nullable;
179        self
180    }
181
182    /// Set whether map fields can be null.
183    pub fn with_map_nullable(mut self, nullable: bool) -> Self {
184        self.map_nullable = nullable;
185        self
186    }
187
188    /// Set whether list element values can be null.
189    pub fn with_list_value_nullable(mut self, nullable: bool) -> Self {
190        self.list_value_nullable = nullable;
191        self
192    }
193
194    /// Set whether map values can be null.
195    pub fn with_map_value_nullable(mut self, nullable: bool) -> Self {
196        self.map_value_nullable = nullable;
197        self
198    }
199
200    /// Set whether to use LargeUtf8 instead of Utf8 for string fields.
201    pub fn with_use_large_string(mut self, use_large: bool) -> Self {
202        self.use_large_string = use_large;
203        self
204    }
205
206    /// Set whether to use LargeBinary instead of Binary for bytes fields.
207    pub fn with_use_large_binary(mut self, use_large: bool) -> Self {
208        self.use_large_binary = use_large;
209        self
210    }
211
212    /// Set whether to use LargeList instead of List for repeated fields.
213    pub fn with_use_large_list(mut self, use_large: bool) -> Self {
214        self.use_large_list = use_large;
215        self
216    }
217
218    /// Set how enum fields are represented in Arrow.
219    pub fn with_enum_repr(mut self, repr: EnumRepr) -> Self {
220        self.enum_repr = repr;
221        self
222    }
223
224    /// Set the Confluent Schema Registry wire format policy.
225    pub fn with_confluent_wire_policy(mut self, policy: ConfluentWirePolicy) -> Self {
226        self.confluent_wire_policy = policy;
227        self
228    }
229}
230
231#[cfg(test)]
232mod tests {
233    use super::*;
234
235    #[test]
236    fn test_default_config() {
237        let config = PtarsConfig::default();
238        assert_eq!(config.timestamp_tz, Some(Arc::from("UTC")));
239        assert_eq!(config.timestamp_unit, TimeUnit::Nanosecond);
240        assert_eq!(config.time_unit, TimeUnit::Nanosecond);
241        assert_eq!(config.duration_unit, TimeUnit::Nanosecond);
242        assert_eq!(config.list_value_name.as_ref(), "item");
243        assert_eq!(config.map_value_name.as_ref(), "value");
244        assert!(!config.list_nullable);
245        assert!(!config.map_nullable);
246        assert!(!config.list_value_nullable);
247        assert!(!config.map_value_nullable);
248        assert!(!config.use_large_string);
249        assert!(!config.use_large_binary);
250        assert!(!config.use_large_list);
251        assert_eq!(config.enum_repr, EnumRepr::Int32);
252        assert_eq!(config.confluent_wire_policy, ConfluentWirePolicy::Raw);
253    }
254
255    #[test]
256    fn test_new_config() {
257        let config = PtarsConfig::new();
258        assert_eq!(config.timestamp_tz, Some(Arc::from("UTC")));
259    }
260
261    #[test]
262    fn test_with_timestamp_tz() {
263        let config = PtarsConfig::new().with_timestamp_tz(Some("America/New_York"));
264        assert_eq!(config.timestamp_tz, Some(Arc::from("America/New_York")));
265
266        let config = PtarsConfig::new().with_timestamp_tz(None);
267        assert_eq!(config.timestamp_tz, None);
268    }
269
270    #[test]
271    fn test_with_timestamp_unit() {
272        let config = PtarsConfig::new().with_timestamp_unit(TimeUnit::Microsecond);
273        assert_eq!(config.timestamp_unit, TimeUnit::Microsecond);
274    }
275
276    #[test]
277    fn test_with_time_unit() {
278        let config = PtarsConfig::new().with_time_unit(TimeUnit::Millisecond);
279        assert_eq!(config.time_unit, TimeUnit::Millisecond);
280    }
281
282    #[test]
283    fn test_with_duration_unit() {
284        let config = PtarsConfig::new().with_duration_unit(TimeUnit::Second);
285        assert_eq!(config.duration_unit, TimeUnit::Second);
286    }
287
288    #[test]
289    fn test_with_list_value_name() {
290        let config = PtarsConfig::new().with_list_value_name("element");
291        assert_eq!(config.list_value_name.as_ref(), "element");
292    }
293
294    #[test]
295    fn test_with_map_value_name() {
296        let config = PtarsConfig::new().with_map_value_name("val");
297        assert_eq!(config.map_value_name.as_ref(), "val");
298    }
299
300    #[test]
301    fn test_with_list_nullable() {
302        let config = PtarsConfig::new().with_list_nullable(true);
303        assert!(config.list_nullable);
304    }
305
306    #[test]
307    fn test_with_map_nullable() {
308        let config = PtarsConfig::new().with_map_nullable(true);
309        assert!(config.map_nullable);
310    }
311
312    #[test]
313    fn test_with_list_value_nullable() {
314        let config = PtarsConfig::new().with_list_value_nullable(true);
315        assert!(config.list_value_nullable);
316    }
317
318    #[test]
319    fn test_with_map_value_nullable() {
320        let config = PtarsConfig::new().with_map_value_nullable(true);
321        assert!(config.map_value_nullable);
322    }
323
324    #[test]
325    fn test_with_use_large_string() {
326        let config = PtarsConfig::new().with_use_large_string(true);
327        assert!(config.use_large_string);
328    }
329
330    #[test]
331    fn test_with_use_large_binary() {
332        let config = PtarsConfig::new().with_use_large_binary(true);
333        assert!(config.use_large_binary);
334    }
335
336    #[test]
337    fn test_with_use_large_list() {
338        let config = PtarsConfig::new().with_use_large_list(true);
339        assert!(config.use_large_list);
340    }
341
342    #[test]
343    fn test_builder_chaining() {
344        let config = PtarsConfig::new()
345            .with_timestamp_tz(Some("Europe/London"))
346            .with_timestamp_unit(TimeUnit::Millisecond)
347            .with_time_unit(TimeUnit::Microsecond)
348            .with_duration_unit(TimeUnit::Second)
349            .with_list_value_name("elem")
350            .with_map_value_name("v")
351            .with_list_nullable(true)
352            .with_map_nullable(true)
353            .with_list_value_nullable(true)
354            .with_map_value_nullable(true)
355            .with_use_large_string(true)
356            .with_use_large_binary(true)
357            .with_use_large_list(true)
358            .with_enum_repr(EnumRepr::String)
359            .with_confluent_wire_policy(ConfluentWirePolicy::Protobuf);
360
361        assert_eq!(config.timestamp_tz, Some(Arc::from("Europe/London")));
362        assert_eq!(config.timestamp_unit, TimeUnit::Millisecond);
363        assert_eq!(config.time_unit, TimeUnit::Microsecond);
364        assert_eq!(config.duration_unit, TimeUnit::Second);
365        assert_eq!(config.list_value_name.as_ref(), "elem");
366        assert_eq!(config.map_value_name.as_ref(), "v");
367        assert!(config.list_nullable);
368        assert!(config.map_nullable);
369        assert!(config.list_value_nullable);
370        assert!(config.map_value_nullable);
371        assert!(config.use_large_string);
372        assert!(config.use_large_binary);
373        assert!(config.use_large_list);
374        assert_eq!(config.enum_repr, EnumRepr::String);
375        assert_eq!(config.confluent_wire_policy, ConfluentWirePolicy::Protobuf);
376    }
377}
ptars_core/config.rs

ptars_core/
config.rs
