omf/file/parquet/
writer.rs

1use std::io::{Seek, Write};
2
3use crate::{
4    Array, ArrayType, array_type,
5    data::{write_checks::*, *},
6    error::Error,
7    file::zip_container::FileType,
8    pqarray::{PqArrayWriter, PqWriteOptions},
9};
10
11use super::super::{Compression, Writer};
12
13impl From<Compression> for PqWriteOptions {
14    fn from(value: Compression) -> Self {
15        Self {
16            compression_level: value.level(),
17            ..Default::default()
18        }
19    }
20}
21
22impl<W: Write + Seek + Send> Writer<W> {
23    fn array_writer<'a>(&self) -> PqArrayWriter<'a> {
24        PqArrayWriter::new(self.compression().into())
25    }
26
27    fn array_write<A: ArrayType>(&mut self, writer: PqArrayWriter) -> Result<Array<A>, Error> {
28        let f = self.builder.open(FileType::Parquet)?;
29        let name = f.name().to_owned();
30        let length = writer.write(f)?;
31        Ok(Array::new(name, length))
32    }
33
34    /// Write an [`array_type::Scalar`](crate::array_type::Scalar) array.
35    ///
36    /// Values can be `f32` or `f64`.
37    pub fn array_scalars<I, T>(&mut self, data: I) -> Result<Array<array_type::Scalar>, Error>
38    where
39        I: IntoIterator<Item = T>,
40        T: FloatType,
41    {
42        let mut min = MinimumScalar::new();
43        let mut writer = self.array_writer();
44        writer.add("scalar", data.into_iter().map(|v| min.visit(v)))?;
45        Ok(self.array_write(writer)?.add_write_checks(min.get()))
46    }
47
48    /// Write an [`array_type::Vertex`](crate::array_type::Vertex) array.
49    pub fn array_vertices<I, T>(&mut self, data: I) -> Result<Array<array_type::Vertex>, Error>
50    where
51        I: IntoIterator<Item = [T; 3]>,
52        T: FloatType,
53    {
54        let mut writer = self.array_writer();
55        writer.add_multiple(&["x", "y", "z"], data)?;
56        self.array_write(writer)
57    }
58
59    /// Write an [`array_type::Segment`](crate::array_type::Segment) array.
60    ///
61    /// Values can be `[u8; 2]`, `[u16; 2]`, or `[u32; 2]` and all indices must be less than the
62    /// number of vertices.
63    pub fn array_segments<I>(&mut self, data: I) -> Result<Array<array_type::Segment>, Error>
64    where
65        I: IntoIterator<Item = [u32; 2]>,
66    {
67        let mut max = MaximumIndex::new();
68        let mut writer = self.array_writer();
69        writer.add_multiple(&["a", "b"], data.into_iter().map(|v| max.visit_array(v)))?;
70        Ok(self.array_write(writer)?.add_write_checks(max.get()))
71    }
72
73    /// Write an [`array_type::Triangle`](crate::array_type::Triangle) array.
74    ///
75    /// Values can be `[u8; 3]`, `[u16; 3]`, or `[u32; 3]` and all indices must be less than the
76    /// number of vertices.
77    pub fn array_triangles<I>(&mut self, data: I) -> Result<Array<array_type::Triangle>, Error>
78    where
79        I: IntoIterator<Item = [u32; 3]>,
80    {
81        let mut max = MaximumIndex::new();
82        let mut writer = self.array_writer();
83        writer.add_multiple(
84            &["a", "b", "c"],
85            data.into_iter().map(|v| max.visit_array(v)),
86        )?;
87        Ok(self.array_write(writer)?.add_write_checks(max.get()))
88    }
89
90    /// Write an [`array_type::Name`](crate::array_type::Name) array.
91    pub fn array_names<I>(&mut self, data: I) -> Result<Array<array_type::Name>, Error>
92    where
93        I: IntoIterator<Item = String>,
94    {
95        let mut writer = self.array_writer();
96        writer.add("name", data)?;
97        self.array_write(writer)
98    }
99
100    /// Write an [`array_type::Gradient`](crate::array_type::Gradient) array.
101    ///
102    /// Values are `[u8; 4]` with channels in RGBA color.
103    pub fn array_gradient<I>(&mut self, data: I) -> Result<Array<array_type::Gradient>, Error>
104    where
105        I: IntoIterator<Item = [u8; 4]>,
106    {
107        let mut writer = self.array_writer();
108        writer.add_multiple(&["r", "g", "b", "a"], data.into_iter())?;
109        self.array_write(writer)
110    }
111
112    /// Write an [`array_type::Texcoord`](crate::array_type::Texcoord) array.
113    ///
114    /// Values can be either `[f32; 2]` or `[f64; 2]` containing normalized texture coordinates.
115    pub fn array_texcoords<I, T>(&mut self, data: I) -> Result<Array<array_type::Texcoord>, Error>
116    where
117        I: IntoIterator<Item = [T; 2]>,
118        T: FloatType,
119    {
120        let mut writer = self.array_writer();
121        writer.add_multiple(&["u", "v"], data)?;
122        self.array_write(writer)
123    }
124
125    /// Write an [`array_type::Boundary`](crate::array_type::Boundary) array.
126    ///
127    /// The boundary value type `T` can be `f64`, `i64`, `chrono::NaiveDate`, or `chrono::DateTime<Utc>`.
128    pub fn array_boundaries<I, T>(&mut self, data: I) -> Result<Array<array_type::Boundary>, Error>
129    where
130        I: IntoIterator<Item = Boundary<T>>,
131        T: NumberType,
132    {
133        let mut increasing = IncreasingBoundary::new();
134        let mut writer = self.array_writer();
135        writer.add_multiple(
136            &["value", "inclusive"],
137            data.into_iter()
138                .map(|b| (increasing.visit(b.value()), b.is_inclusive())),
139        )?;
140        Ok(self.array_write(writer)?.add_write_checks(increasing.get()))
141    }
142
143    /// Write an [`array_type::RegularSubblock`](crate::array_type::RegularSubblock) array.
144    ///
145    /// The `parent_indices` and `corners` iterators must be the same length. Each row is
146    /// `[parent_i, parent_j, parent_k, min_corner_i, min_corner_j, min_corner_k,
147    /// max_corner_i, max_corner_j, max_corner_k]`. The parent and corner indices can
148    /// be different types.
149    ///
150    /// Parent indices can be `[u8; 3]`, `[u16; 3]`, or `[u32; 3]`. Sub-block corners can
151    /// separately be `[u8; 6]`, `[u16; 6]`, or `[u32; 6]` which each row storing
152    /// $(u_{min}, v_{min}, w_{min}, u_{max}, v_{max}, w_{max})$ as indices into the regular
153    /// grid within the parent block.
154    pub fn array_regular_subblocks<I>(
155        &mut self,
156        data: I,
157    ) -> Result<Array<array_type::RegularSubblock>, Error>
158    where
159        I: IntoIterator<Item = ([u32; 3], [u32; 6])>,
160    {
161        let mut parents = ParentIndices::new();
162        let mut corners = RegularCorners::new();
163        let mut writer = self.array_writer();
164        writer.add_multiple(
165            &[
166                "parent_u",
167                "parent_v",
168                "parent_w",
169                "corner_min_u",
170                "corner_min_v",
171                "corner_min_w",
172                "corner_max_u",
173                "corner_max_v",
174                "corner_max_w",
175            ],
176            data.into_iter().map(|(p, c)| {
177                parents.visit(p);
178                corners.visit(c);
179                (p[0], p[1], p[2], c[0], c[1], c[2], c[3], c[4], c[5])
180            }),
181        )?;
182        Ok(self
183            .array_write(writer)?
184            .add_write_checks(parents.get())
185            .add_write_checks(corners.get()))
186    }
187
188    /// Write an [`array_type::FreeformSubblock`](crate::array_type::FreeformSubblock) array.
189    ///
190    /// The `parent_indices` and `corners` iterators must be the same length. Each row is
191    /// `[parent_i, parent_j, parent_k, min_corner_x, min_corner_y, min_corner_z,
192    /// max_corner_x, max_corner_y, max_corner_z]`.
193    ///
194    /// Parent indices can be `[u8; 3]`, `[u16; 3]`, or `[u32; 3]`. Sub-block corners can be
195    /// `[f32; 6]` or `[f64; 6]` which each row storing
196    /// $(u_{min}, v_{min}, w_{min}, u_{max}, v_{max}, w_{max})$ in the range [0, 1] relative
197    /// to the parent block.
198    pub fn array_freeform_subblocks<I, C>(
199        &mut self,
200        data: I,
201    ) -> Result<Array<array_type::FreeformSubblock>, Error>
202    where
203        I: IntoIterator<Item = ([u32; 3], [C; 6])>,
204        C: FloatType,
205    {
206        let mut corner = FreeformCorners::new();
207        let mut parent = ParentIndices::new();
208        let mut writer = self.array_writer();
209        writer.add_multiple(
210            &[
211                "parent_u",
212                "parent_v",
213                "parent_w",
214                "corner_min_u",
215                "corner_min_v",
216                "corner_min_w",
217                "corner_max_u",
218                "corner_max_v",
219                "corner_max_w",
220            ],
221            data.into_iter().map(|(p, c)| {
222                parent.visit(p);
223                corner.visit(c);
224                (p[0], p[1], p[2], c[0], c[1], c[2], c[3], c[4], c[5])
225            }),
226        )?;
227        Ok(self
228            .array_write(writer)?
229            .add_write_checks(parent.get())
230            .add_write_checks(corner.get()))
231    }
232
233    /// Write an [`array_type::Number`](crate::array_type::Number) array.
234    ///
235    /// Values are `Option<T>` where `T` can be `f32`, `f64`, `i32`, `i64`, `chrono::NaiveDate`,
236    /// or `chrono::DateTime<Utc>`. Use `None` to represent null values rather than NaN or
237    /// any flag values like −9999.
238    pub fn array_numbers<I, T>(&mut self, data: I) -> Result<Array<array_type::Number>, Error>
239    where
240        I: IntoIterator<Item = Option<T>>,
241        T: NumberType,
242    {
243        let mut writer = self.array_writer();
244        writer.add_nullable("number", data)?;
245        self.array_write(writer)
246    }
247
248    /// Write an [`array_type::Index`](crate::array_type::Index) array.
249    ///
250    /// Values are `Option<T>` where `T` can be `Option<u8>`, `Option<u16>`, or `Option<u32>`.
251    /// Smaller types won't compress much better but will let other applications allocate less
252    /// memory when reading the array. Use `None` to represent null values.
253    pub fn array_indices<I>(&mut self, data: I) -> Result<Array<array_type::Index>, Error>
254    where
255        I: IntoIterator<Item = Option<u32>>,
256    {
257        let mut max = MaximumIndex::new();
258        let mut writer = self.array_writer();
259        writer.add_nullable("index", data.into_iter().map(|v| max.visit_opt(v)))?;
260        Ok(self.array_write(writer)?.add_write_checks(max.get()))
261    }
262
263    /// Write a [`array_type::Vector`](crate::array_type::Vector) array.
264    ///
265    /// Values are `Option<T>` where `T` can be `[f32; 2]`, `[f64; 2]`, `[f32; 3]`, or `[f64; 3]`.
266    pub fn array_vectors<I, T, V>(&mut self, data: I) -> Result<Array<array_type::Vector>, Error>
267    where
268        I: IntoIterator<Item = Option<V>>,
269        V: VectorSource<T>,
270        T: FloatType,
271    {
272        let mut writer = self.array_writer();
273        if V::IS_3D {
274            writer.add_nullable_group(
275                "vector",
276                &["x", "y", "z"],
277                data.into_iter().map(|o| o.map(V::into_3d)),
278            )?;
279        } else {
280            writer.add_nullable_group(
281                "vector",
282                &["x", "y"],
283                data.into_iter().map(|o| o.map(V::into_2d)),
284            )?;
285        }
286        self.array_write(writer)
287    }
288
289    /// Write a [`array_type::Text`](crate::array_type::Text) array.
290    pub fn array_text<I>(&mut self, data: I) -> Result<Array<array_type::Text>, Error>
291    where
292        I: IntoIterator<Item = Option<String>>,
293    {
294        let mut writer = self.array_writer();
295        writer.add_nullable("text", data)?;
296        self.array_write(writer)
297    }
298
299    /// Write a [`array_type::Boolean`](crate::array_type::Boolean) array.
300    pub fn array_booleans<I>(&mut self, data: I) -> Result<Array<array_type::Boolean>, Error>
301    where
302        I: IntoIterator<Item = Option<bool>>,
303    {
304        let mut writer = self.array_writer();
305        writer.add_nullable("bool", data)?;
306        self.array_write(writer)
307    }
308
309    /// Write a [`array_type::Color`](crate::array_type::Color) array.
310    ///
311    /// Values are `Option<[u8; 4]>` with channels in RGBA order.
312    pub fn array_colors<I>(&mut self, data: I) -> Result<Array<array_type::Color>, Error>
313    where
314        I: IntoIterator<Item = Option<[u8; 4]>>,
315    {
316        let mut writer = self.array_writer();
317        writer.add_nullable_group("color", &["r", "g", "b", "a"], data.into_iter())?;
318        self.array_write(writer)
319    }
320}
omf/file/parquet/writer.rs

omf/file/parquet/
writer.rs
