1use std::iter;
5
6use arrow_array::{Array, FixedSizeBinaryArray, LargeBinaryArray, LargeStringArray};
7use arrow_buffer::BooleanBuffer;
8use reifydb_value::{
9 encoding::LeBytes,
10 value::{
11 Value,
12 column_view::{ColumnView, ViewData},
13 container::{
14 any_array,
15 decimal_array::DecimalView,
16 dictionary_array, digest_array,
17 temporal_array::{dates, datetimes, durations, times},
18 uuid_array::{identity_ids, uuid4s, uuid7s},
19 wide_int_array::wides,
20 },
21 date::Date,
22 datetime::DateTime,
23 dictionary::DictionaryEntryId,
24 duration::Duration,
25 identity::IdentityId,
26 time::Time,
27 uuid::{Uuid4, Uuid7},
28 value_type::ValueType,
29 },
30};
31
32use crate::{error::EncodeError, tag::ValueKind, value::encode_value_into};
33
34macro_rules! encode_fixed {
35 ($container:expr, $ty:expr, $elem:ty) => {
36 encode_fixed!(slice: $container.values(), $ty, $elem)
37 };
38 (slice: $slice:expr, $ty:expr, $elem:ty) => {{
39 let slice: &[$elem] = $slice;
40 let mut buf = Vec::with_capacity(slice.len() * <$elem as LeBytes>::ENCODED_SIZE);
41 for v in slice {
42 buf.extend_from_slice(LeBytes::to_le_bytes(v).as_ref());
43 }
44 PlainEncoded {
45 data: buf,
46 offsets: vec![],
47 type_code: ValueKind::of_type(&$ty).byte(),
48 }
49 }};
50}
51
52pub struct PlainEncoded {
53 pub data: Vec<u8>,
54
55 pub offsets: Vec<u8>,
56
57 pub type_code: u8,
58}
59
60pub fn encode_plain(col: &ColumnView<'_>) -> Result<PlainEncoded, EncodeError> {
61 let result = match &col.data {
62 ViewData::Bool(c) => PlainEncoded {
63 data: encode_bitvec(c.values()),
64 offsets: vec![],
65 type_code: ValueKind::Boolean.byte(),
66 },
67 ViewData::Float4(c) => encode_fixed!(c, ValueType::Float4, f32),
68 ViewData::Float8(c) => encode_fixed!(c, ValueType::Float8, f64),
69 ViewData::Int1(c) => encode_fixed!(c, ValueType::Int1, i8),
70 ViewData::Int2(c) => encode_fixed!(c, ValueType::Int2, i16),
71 ViewData::Int4(c) => encode_fixed!(c, ValueType::Int4, i32),
72 ViewData::Int8(c) => encode_fixed!(c, ValueType::Int8, i64),
73 ViewData::Int16(c) => encode_fixed!(slice: &wides::<i128>(c), ValueType::Int16, i128),
74 ViewData::Uint1(c) => encode_fixed!(c, ValueType::Uint1, u8),
75 ViewData::Uint2(c) => encode_fixed!(c, ValueType::Uint2, u16),
76 ViewData::Uint4(c) => encode_fixed!(c, ValueType::Uint4, u32),
77 ViewData::Uint8(c) => encode_fixed!(c, ValueType::Uint8, u64),
78 ViewData::Uint16(c) => encode_fixed!(slice: &wides::<u128>(c), ValueType::Uint16, u128),
79 ViewData::Date(c) => {
80 let slice: &[Date] = dates(c);
81 let mut buf = Vec::with_capacity(slice.len() * Date::ENCODED_SIZE);
82 for v in slice {
83 buf.extend_from_slice(v.to_le_bytes().as_ref());
84 }
85 PlainEncoded {
86 data: buf,
87 offsets: vec![],
88 type_code: ValueKind::Date.byte(),
89 }
90 }
91 ViewData::DateTime(c) => {
92 let slice: &[DateTime] = datetimes(c);
93 let mut buf = Vec::with_capacity(slice.len() * DateTime::ENCODED_SIZE);
94 for v in slice {
95 buf.extend_from_slice(v.to_le_bytes().as_ref());
96 }
97 PlainEncoded {
98 data: buf,
99 offsets: vec![],
100 type_code: ValueKind::DateTime.byte(),
101 }
102 }
103 ViewData::Time(c) => {
104 let slice: &[Time] = times(c);
105 let mut buf = Vec::with_capacity(slice.len() * Time::ENCODED_SIZE);
106 for v in slice {
107 buf.extend_from_slice(v.to_le_bytes().as_ref());
108 }
109 PlainEncoded {
110 data: buf,
111 offsets: vec![],
112 type_code: ValueKind::Time.byte(),
113 }
114 }
115 ViewData::Duration(c) => {
116 let slice: &[Duration] = durations(c);
117 let mut buf = Vec::with_capacity(slice.len() * Duration::ENCODED_SIZE);
118 for v in slice {
119 buf.extend_from_slice(v.to_le_bytes().as_ref());
120 }
121 PlainEncoded {
122 data: buf,
123 offsets: vec![],
124 type_code: ValueKind::Duration.byte(),
125 }
126 }
127 ViewData::IdentityId(c) => {
128 let slice: &[IdentityId] = identity_ids(c);
129 let mut buf = Vec::with_capacity(slice.len() * IdentityId::ENCODED_SIZE);
130 for v in slice {
131 buf.extend_from_slice(v.to_le_bytes().as_ref());
132 }
133 PlainEncoded {
134 data: buf,
135 offsets: vec![],
136 type_code: ValueKind::IdentityId.byte(),
137 }
138 }
139 ViewData::Uuid4(c) => {
140 let slice: &[Uuid4] = uuid4s(c);
141 let mut buf = Vec::with_capacity(slice.len() * Uuid4::ENCODED_SIZE);
142 for v in slice {
143 buf.extend_from_slice(v.to_le_bytes().as_ref());
144 }
145 PlainEncoded {
146 data: buf,
147 offsets: vec![],
148 type_code: ValueKind::Uuid4.byte(),
149 }
150 }
151 ViewData::Uuid7(c) => {
152 let slice: &[Uuid7] = uuid7s(c);
153 let mut buf = Vec::with_capacity(slice.len() * Uuid7::ENCODED_SIZE);
154 for v in slice {
155 buf.extend_from_slice(v.to_le_bytes().as_ref());
156 }
157 PlainEncoded {
158 data: buf,
159 offsets: vec![],
160 type_code: ValueKind::Uuid7.byte(),
161 }
162 }
163 ViewData::Utf8 {
164 container,
165 ..
166 } => encode_varlen_strings(container, ValueType::Utf8),
167 ViewData::Enum {
168 container,
169 ..
170 } => encode_varlen_strings(container, ValueType::Utf8),
171 ViewData::Blob {
172 container,
173 ..
174 } => encode_varlen_blobs(container, ValueType::Blob),
175 ViewData::Decimal(c) => encode_unscaled(*c, ValueKind::Decimal),
176 ViewData::Any {
177 container,
178 ..
179 } => {
180 let mut data = Vec::new();
181 for val in any_array::values(container) {
182 encode_value_into(&val, &mut data)?;
183 }
184 return Ok(PlainEncoded {
185 data,
186 offsets: vec![],
187 type_code: ValueKind::Any.byte(),
188 });
189 }
190 ViewData::DictionaryId {
191 container,
192 ..
193 } => encode_dictionary_ids(container),
194 ViewData::Digest {
195 container,
196 ..
197 } => encode_varlen(
198 container.len(),
199 |i| match digest_array::get(container, i) {
200 Some(digest) => digest.encode(),
201 None => Vec::new(),
202 },
203 col.base_type(),
204 ),
205 ViewData::None {
206 array,
207 } => {
208 let mut data = Vec::new();
209 for _ in 0..array.len() {
210 encode_value_into(&Value::none(), &mut data)?;
211 }
212 return Ok(PlainEncoded {
213 data,
214 offsets: vec![],
215 type_code: ValueKind::Any.byte(),
216 });
217 }
218 };
219 Ok(result)
220}
221
222fn encode_unscaled(array: DecimalView<'_>, kind: ValueKind) -> PlainEncoded {
223 let data = match array {
224 DecimalView::Decimal128(a) => a.values().iter().flat_map(|v| v.to_le_bytes()).collect(),
225 DecimalView::Decimal256(a) => a.values().iter().flat_map(|v| v.to_le_bytes()).collect(),
226 };
227 PlainEncoded {
228 data,
229 offsets: vec![],
230 type_code: kind.byte(),
231 }
232}
233
234fn encode_varlen(count: usize, get_bytes: impl Fn(usize) -> Vec<u8>, ty: ValueType) -> PlainEncoded {
235 let mut offsets = Vec::with_capacity((count + 1) * 4);
236 let mut data = Vec::new();
237 let mut offset: u32 = 0;
238 offsets.extend_from_slice(&offset.to_le_bytes());
239 for i in 0..count {
240 let bytes = get_bytes(i);
241 data.extend_from_slice(&bytes);
242 offset += bytes.len() as u32;
243 offsets.extend_from_slice(&offset.to_le_bytes());
244 }
245 PlainEncoded {
246 data,
247 offsets,
248 type_code: ValueKind::of_type(&ty).byte(),
249 }
250}
251
252fn encode_varlen_strings(c: &LargeStringArray, ty: ValueType) -> PlainEncoded {
253 encode_varlen(c.len(), |i| c.value(i).as_bytes().to_vec(), ty)
254}
255
256fn encode_varlen_blobs(c: &LargeBinaryArray, ty: ValueType) -> PlainEncoded {
257 encode_varlen(c.len(), |i| c.value(i).to_vec(), ty)
258}
259
260fn encode_dictionary_ids(c: &FixedSizeBinaryArray) -> PlainEncoded {
261 let mut buf = Vec::new();
262 if !c.is_empty() {
263 let mut disc = 1u8;
264 for i in 0..c.len() {
265 if let Some(id) = dictionary_array::get(c, i) {
266 let d = match id {
267 DictionaryEntryId::U1(_) => 1u8,
268 DictionaryEntryId::U2(_) => 2u8,
269 DictionaryEntryId::U4(_) => 4u8,
270 DictionaryEntryId::U8(_) => 8u8,
271 DictionaryEntryId::U16(_) => 16u8,
272 };
273 if d > disc {
274 disc = d;
275 }
276 }
277 }
278
279 buf.push(disc);
280 for i in 0..c.len() {
281 if let Some(id) = dictionary_array::get(c, i) {
282 match disc {
283 1 => buf.push(id.to_u128() as u8),
284 2 => buf.extend_from_slice(&(id.to_u128() as u16).to_le_bytes()),
285 4 => buf.extend_from_slice(&(id.to_u128() as u32).to_le_bytes()),
286 8 => buf.extend_from_slice(&(id.to_u128() as u64).to_le_bytes()),
287 16 => buf.extend_from_slice(&id.to_u128().to_le_bytes()),
288 _ => unreachable!(),
289 }
290 } else {
291 buf.extend(iter::repeat_n(0, disc as usize));
292 }
293 }
294 }
295 PlainEncoded {
296 data: buf,
297 offsets: vec![],
298 type_code: ValueKind::DictionaryId.byte(),
299 }
300}
301
302pub fn encode_bitvec(bv: &BooleanBuffer) -> Vec<u8> {
303 let len = bv.len();
304 let byte_count = len.div_ceil(8);
305 let mut buf = Vec::with_capacity(byte_count);
306 for i in 0..byte_count {
307 let mut byte = 0u8;
308 for bit in 0..8 {
309 let idx = i * 8 + bit;
310 if idx < len && bv.value(idx) {
311 byte |= 1 << bit;
312 }
313 }
314 buf.push(byte);
315 }
316 buf
317}