fory_core/row/
writer.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use super::{bit_util::calculate_bitmap_width_in_bytes, row::Row};
19use crate::buffer::Writer;
20use crate::error::Error;
21
22pub struct WriteCallbackInfo {
23    field_offset: usize,
24    data_start: usize,
25}
26
27struct FieldWriterHelper<'a, 'b> {
28    pub writer: &'a mut Writer<'b>,
29    base_offset: usize,
30    get_field_offset: Box<dyn Fn(usize) -> usize>,
31}
32
33impl<'b: 'a, 'a> FieldWriterHelper<'a, 'b> {
34    fn new(
35        writer: &'a mut Writer<'b>,
36        base_offset: usize,
37        get_field_offset: Box<dyn Fn(usize) -> usize>,
38    ) -> FieldWriterHelper<'a, 'b> {
39        FieldWriterHelper {
40            writer,
41            base_offset,
42            get_field_offset,
43        }
44    }
45
46    fn write_start(&mut self, idx: usize) -> WriteCallbackInfo {
47        let base_offset = self.base_offset;
48        let field_offset = (self.get_field_offset)(idx);
49        let writer: &mut Writer = self.writer;
50        let offset = writer.len() - base_offset;
51        writer.set_bytes(field_offset, &(offset as u32).to_le_bytes());
52        let data_start: usize = writer.len();
53        WriteCallbackInfo {
54            field_offset,
55            data_start,
56        }
57    }
58
59    fn write_end(&mut self, callback_info: WriteCallbackInfo) {
60        let writer: &mut Writer = self.writer;
61        let size: usize = writer.len() - callback_info.data_start;
62        writer.set_bytes(callback_info.field_offset + 4, &(size as u32).to_le_bytes());
63    }
64}
65
66pub struct StructWriter<'a, 'b> {
67    field_writer_helper: FieldWriterHelper<'a, 'b>,
68}
69
70impl<'a, 'b> StructWriter<'a, 'b> {
71    fn get_fixed_size(bit_map_width_in_bytes: usize, num_fields: usize) -> usize {
72        bit_map_width_in_bytes + num_fields * 8
73    }
74    pub fn new(num_fields: usize, writer: &'a mut Writer<'b>) -> StructWriter<'a, 'b> {
75        let base_offset = writer.len();
76        let bit_map_width_in_bytes = calculate_bitmap_width_in_bytes(num_fields);
77
78        let struct_writer = StructWriter {
79            field_writer_helper: FieldWriterHelper::new(
80                writer,
81                base_offset,
82                Box::new(move |idx| base_offset + bit_map_width_in_bytes + idx * 8),
83            ),
84        };
85        let fixed_size = Self::get_fixed_size(bit_map_width_in_bytes, num_fields);
86        struct_writer.field_writer_helper.writer.reserve(fixed_size);
87        struct_writer.field_writer_helper.writer.skip(fixed_size);
88        struct_writer
89    }
90
91    pub fn get_writer(&mut self) -> &mut Writer<'b> {
92        self.field_writer_helper.writer
93    }
94
95    pub fn write_start(&mut self, idx: usize) -> WriteCallbackInfo {
96        self.field_writer_helper.write_start(idx)
97    }
98
99    pub fn write_end(&mut self, callback_info: WriteCallbackInfo) {
100        self.field_writer_helper.write_end(callback_info)
101    }
102}
103
104pub struct ArrayWriter<'a, 'b> {
105    field_writer_helper: FieldWriterHelper<'a, 'b>,
106}
107
108impl<'a, 'b> ArrayWriter<'a, 'b> {
109    fn get_fixed_size(bit_map_width_in_bytes: usize, num_fields: usize) -> usize {
110        8 + bit_map_width_in_bytes + num_fields * 8
111    }
112
113    pub fn new(
114        num_fields: usize,
115        writer: &'a mut Writer<'b>,
116    ) -> Result<ArrayWriter<'a, 'b>, Error> {
117        let base_offset = writer.len();
118        let bit_map_width_in_bytes = calculate_bitmap_width_in_bytes(num_fields);
119        let array_writer = ArrayWriter {
120            field_writer_helper: FieldWriterHelper::new(
121                writer,
122                base_offset,
123                Box::new(move |idx| 8 + base_offset + bit_map_width_in_bytes + idx * 8),
124            ),
125        };
126        let fixed_size = Self::get_fixed_size(bit_map_width_in_bytes, num_fields);
127        array_writer.field_writer_helper.writer.reserve(fixed_size);
128        array_writer
129            .field_writer_helper
130            .writer
131            .write_u64(num_fields as u64);
132        array_writer.field_writer_helper.writer.skip(fixed_size - 8);
133        Ok(array_writer)
134    }
135
136    pub fn get_writer(&mut self) -> &mut Writer<'b> {
137        self.field_writer_helper.writer
138    }
139
140    pub fn write_start(&mut self, idx: usize) -> WriteCallbackInfo {
141        self.field_writer_helper.write_start(idx)
142    }
143
144    pub fn write_end(&mut self, callback_info: WriteCallbackInfo) {
145        self.field_writer_helper.write_end(callback_info)
146    }
147}
148
149pub struct MapWriter<'a, 'b> {
150    base_offset: usize,
151    writer: &'a mut Writer<'b>,
152}
153
154impl<'a, 'b> MapWriter<'a, 'b> {
155    fn get_fixed_size(&self) -> usize {
156        // key_byte_size
157        8
158    }
159
160    pub fn new(writer: &'a mut Writer<'b>) -> MapWriter<'a, 'b> {
161        let base_offset = writer.len();
162        let array_writer = MapWriter {
163            writer,
164            base_offset,
165        };
166        let fixed_size = array_writer.get_fixed_size();
167        array_writer.writer.reserve(fixed_size);
168        array_writer.writer.skip(fixed_size);
169        array_writer
170    }
171
172    pub fn get_writer(&mut self) -> &mut Writer<'b> {
173        self.writer
174    }
175
176    pub fn write_start(&mut self, _idx: usize) -> usize {
177        self.writer.len()
178    }
179
180    pub fn write_end(&mut self, data_start: usize) {
181        let size: usize = self.writer.len() - data_start;
182        self.writer
183            .set_bytes(self.base_offset, &(size as u64).to_le_bytes());
184    }
185}
186
187pub fn to_row<'a, T: Row<'a>>(v: &T) -> Result<Vec<u8>, Error> {
188    let mut buffer = vec![];
189    let mut writer = Writer::from_buffer(&mut buffer);
190    T::write(v, &mut writer)?;
191    Ok(buffer)
192}
fory_core/row/writer.rs

fory_core/row/
writer.rs
