prost_arrow/
binary.rs

1use std::any::Any;
2use std::sync::Arc;
3
4use arrow_array::builder::{ArrayBuilder, BinaryBuilder, StringBuilder};
5use arrow_array::ArrayRef;
6use arrow_schema::DataType;
7use prost::alloc::string::String;
8
9use crate::traits::{ArrowBuilder, ToArrow};
10
11impl ToArrow for String {
12    type Item = String;
13    type Builder = ArrowStringBuilder;
14
15    fn to_datatype() -> DataType {
16        DataType::Utf8
17    }
18}
19
20impl ToArrow for Vec<u8> {
21    type Item = Vec<u8>;
22    type Builder = ArrowBinaryBuilder;
23
24    fn to_datatype() -> DataType {
25        DataType::Binary
26    }
27}
28
29pub struct ArrowStringBuilder {
30    builder: StringBuilder,
31}
32
33impl ArrowBuilder<String> for ArrowStringBuilder {
34    fn new_with_capacity(capacity: usize) -> Self {
35        Self {
36            builder: StringBuilder::with_capacity(capacity, capacity * 1024),
37        }
38    }
39
40    fn append_value(&mut self, value: String) {
41        self.builder.append_value(value)
42    }
43
44    fn append_option(&mut self, value: Option<String>) {
45        self.builder.append_option(value)
46    }
47}
48
49impl ArrayBuilder for ArrowStringBuilder {
50    fn len(&self) -> usize {
51        self.builder.len()
52    }
53
54    fn finish(&mut self) -> ArrayRef {
55        Arc::new(self.builder.finish())
56    }
57
58    fn finish_cloned(&self) -> ArrayRef {
59        Arc::new(self.builder.finish_cloned())
60    }
61
62    fn as_any(&self) -> &dyn Any {
63        self
64    }
65
66    fn as_any_mut(&mut self) -> &mut dyn Any {
67        self
68    }
69
70    fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
71        self
72    }
73}
74
75pub struct ArrowBinaryBuilder {
76    builder: BinaryBuilder,
77}
78
79impl ArrowBuilder<Vec<u8>> for ArrowBinaryBuilder {
80    fn new_with_capacity(capacity: usize) -> Self {
81        Self {
82            builder: BinaryBuilder::with_capacity(capacity, capacity * 1024),
83        }
84    }
85
86    fn append_value(&mut self, value: Vec<u8>) {
87        self.builder.append_value(value)
88    }
89
90    fn append_option(&mut self, value: Option<Vec<u8>>) {
91        self.builder.append_option(value)
92    }
93}
94
95impl ArrayBuilder for ArrowBinaryBuilder {
96    fn len(&self) -> usize {
97        self.builder.len()
98    }
99
100    fn finish(&mut self) -> ArrayRef {
101        Arc::new(self.builder.finish())
102    }
103
104    fn finish_cloned(&self) -> ArrayRef {
105        Arc::new(self.builder.finish_cloned())
106    }
107
108    fn as_any(&self) -> &dyn Any {
109        self
110    }
111
112    fn as_any_mut(&mut self) -> &mut dyn Any {
113        self
114    }
115
116    fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
117        self
118    }
119}
prost_arrow/binary.rs

prost_arrow/
binary.rs
