1use crate::ShortString;
18use crate::utils::{
19 array_from_slice, overflow_error, slice_from_slice_at_offset, string_from_slice,
20};
21
22use arrow_schema::ArrowError;
23use chrono::{DateTime, Duration, NaiveDate, NaiveDateTime, NaiveTime, Utc};
24use uuid::Uuid;
25
26#[derive(Debug, Clone, Copy, PartialEq)]
34pub enum VariantBasicType {
35 Primitive = 0,
36 ShortString = 1,
37 Object = 2,
38 Array = 3,
39}
40
41#[derive(Debug, Clone, Copy, PartialEq)]
49pub enum VariantPrimitiveType {
50 Null = 0,
51 BooleanTrue = 1,
52 BooleanFalse = 2,
53 Int8 = 3,
54 Int16 = 4,
55 Int32 = 5,
56 Int64 = 6,
57 Double = 7,
58 Decimal4 = 8,
59 Decimal8 = 9,
60 Decimal16 = 10,
61 Date = 11,
62 TimestampMicros = 12,
63 TimestampNtzMicros = 13,
64 Float = 14,
65 Binary = 15,
66 String = 16,
67 Time = 17,
68 TimestampNanos = 18,
69 TimestampNtzNanos = 19,
70 Uuid = 20,
71}
72
73pub(crate) fn get_basic_type(header: u8) -> VariantBasicType {
75 let basic_type = header & 0x03; match basic_type {
78 0 => VariantBasicType::Primitive,
79 1 => VariantBasicType::ShortString,
80 2 => VariantBasicType::Object,
81 3 => VariantBasicType::Array,
82 _ => {
83 unreachable!();
86 }
87 }
88}
89
90impl TryFrom<u8> for VariantPrimitiveType {
91 type Error = ArrowError;
92
93 fn try_from(value: u8) -> Result<Self, Self::Error> {
94 match value {
95 0 => Ok(VariantPrimitiveType::Null),
96 1 => Ok(VariantPrimitiveType::BooleanTrue),
97 2 => Ok(VariantPrimitiveType::BooleanFalse),
98 3 => Ok(VariantPrimitiveType::Int8),
99 4 => Ok(VariantPrimitiveType::Int16),
100 5 => Ok(VariantPrimitiveType::Int32),
101 6 => Ok(VariantPrimitiveType::Int64),
102 7 => Ok(VariantPrimitiveType::Double),
103 8 => Ok(VariantPrimitiveType::Decimal4),
104 9 => Ok(VariantPrimitiveType::Decimal8),
105 10 => Ok(VariantPrimitiveType::Decimal16),
106 11 => Ok(VariantPrimitiveType::Date),
107 12 => Ok(VariantPrimitiveType::TimestampMicros),
108 13 => Ok(VariantPrimitiveType::TimestampNtzMicros),
109 14 => Ok(VariantPrimitiveType::Float),
110 15 => Ok(VariantPrimitiveType::Binary),
111 16 => Ok(VariantPrimitiveType::String),
112 17 => Ok(VariantPrimitiveType::Time),
113 18 => Ok(VariantPrimitiveType::TimestampNanos),
114 19 => Ok(VariantPrimitiveType::TimestampNtzNanos),
115 20 => Ok(VariantPrimitiveType::Uuid),
116 _ => Err(ArrowError::InvalidArgumentError(format!(
117 "unknown primitive type: {value}",
118 ))),
119 }
120 }
121}
122
123#[derive(Debug, Clone, Copy, PartialEq)]
127pub(crate) enum OffsetSizeBytes {
128 One = 1,
129 Two = 2,
130 Three = 3,
131 Four = 4,
132}
133
134impl OffsetSizeBytes {
135 pub(crate) fn try_new(offset_size_minus_one: u8) -> Result<Self, ArrowError> {
137 use OffsetSizeBytes::*;
138 let result = match offset_size_minus_one {
139 0 => One,
140 1 => Two,
141 2 => Three,
142 3 => Four,
143 _ => {
144 return Err(ArrowError::InvalidArgumentError(
145 "offset_size_minus_one must be 0–3".to_string(),
146 ));
147 }
148 };
149 Ok(result)
150 }
151
152 pub(crate) fn unpack_u32(&self, bytes: &[u8], index: usize) -> Result<u32, ArrowError> {
159 self.unpack_u32_at_offset(bytes, 0, index)
160 }
161
162 pub(crate) fn unpack_u32_at_offset(
172 &self,
173 bytes: &[u8],
174 byte_offset: usize, offset_index: usize, ) -> Result<u32, ArrowError> {
177 use OffsetSizeBytes::*;
178
179 let offset = offset_index
182 .checked_mul(*self as usize)
183 .and_then(|n| n.checked_add(byte_offset))
184 .ok_or_else(|| overflow_error("unpacking offset array value"))?;
185 let value = match self {
186 One => u8::from_le_bytes(array_from_slice(bytes, offset)?).into(),
187 Two => u16::from_le_bytes(array_from_slice(bytes, offset)?).into(),
188 Three => {
189 let b3_chunks: [u8; 3] = array_from_slice(bytes, offset)?;
191 let mut buf = [0u8; 4];
193 buf[..3].copy_from_slice(&b3_chunks);
194 u32::from_le_bytes(buf)
195 }
196 Four => u32::from_le_bytes(array_from_slice(bytes, offset)?),
197 };
198 Ok(value)
199 }
200}
201
202pub(crate) fn map_bytes_to_offsets(
204 buffer: &[u8],
205 offset_size: OffsetSizeBytes,
206) -> impl Iterator<Item = usize> + use<'_> {
207 buffer
208 .chunks_exact(offset_size as usize)
209 .map(move |chunk| match offset_size {
210 OffsetSizeBytes::One => chunk[0] as usize,
211 OffsetSizeBytes::Two => u16::from_le_bytes([chunk[0], chunk[1]]) as usize,
212 OffsetSizeBytes::Three => {
213 u32::from_le_bytes([chunk[0], chunk[1], chunk[2], 0]) as usize
214 }
215 OffsetSizeBytes::Four => {
216 u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]) as usize
217 }
218 })
219}
220
221pub(crate) fn get_primitive_type(metadata: u8) -> Result<VariantPrimitiveType, ArrowError> {
223 VariantPrimitiveType::try_from(metadata >> 2)
225}
226
227pub(crate) fn decode_int8(data: &[u8]) -> Result<i8, ArrowError> {
229 Ok(i8::from_le_bytes(array_from_slice(data, 0)?))
230}
231
232pub(crate) fn decode_int16(data: &[u8]) -> Result<i16, ArrowError> {
234 Ok(i16::from_le_bytes(array_from_slice(data, 0)?))
235}
236
237pub(crate) fn decode_int32(data: &[u8]) -> Result<i32, ArrowError> {
239 Ok(i32::from_le_bytes(array_from_slice(data, 0)?))
240}
241
242pub(crate) fn decode_int64(data: &[u8]) -> Result<i64, ArrowError> {
244 Ok(i64::from_le_bytes(array_from_slice(data, 0)?))
245}
246
247pub(crate) fn decode_decimal4(data: &[u8]) -> Result<(i32, u8), ArrowError> {
249 let scale = u8::from_le_bytes(array_from_slice(data, 0)?);
250 let integer = i32::from_le_bytes(array_from_slice(data, 1)?);
251 Ok((integer, scale))
252}
253
254pub(crate) fn decode_decimal8(data: &[u8]) -> Result<(i64, u8), ArrowError> {
256 let scale = u8::from_le_bytes(array_from_slice(data, 0)?);
257 let integer = i64::from_le_bytes(array_from_slice(data, 1)?);
258 Ok((integer, scale))
259}
260
261pub(crate) fn decode_decimal16(data: &[u8]) -> Result<(i128, u8), ArrowError> {
263 let scale = u8::from_le_bytes(array_from_slice(data, 0)?);
264 let integer = i128::from_le_bytes(array_from_slice(data, 1)?);
265 Ok((integer, scale))
266}
267
268pub(crate) fn decode_float(data: &[u8]) -> Result<f32, ArrowError> {
270 Ok(f32::from_le_bytes(array_from_slice(data, 0)?))
271}
272
273pub(crate) fn decode_double(data: &[u8]) -> Result<f64, ArrowError> {
275 Ok(f64::from_le_bytes(array_from_slice(data, 0)?))
276}
277
278pub(crate) fn decode_date(data: &[u8]) -> Result<NaiveDate, ArrowError> {
280 let days_since_epoch = i32::from_le_bytes(array_from_slice(data, 0)?);
281 let value = DateTime::UNIX_EPOCH + Duration::days(i64::from(days_since_epoch));
282 Ok(value.date_naive())
283}
284
285pub(crate) fn decode_timestamp_micros(data: &[u8]) -> Result<DateTime<Utc>, ArrowError> {
287 let micros_since_epoch = i64::from_le_bytes(array_from_slice(data, 0)?);
288 DateTime::from_timestamp_micros(micros_since_epoch).ok_or_else(|| {
289 ArrowError::CastError(format!(
290 "Could not cast `{micros_since_epoch}` microseconds into a DateTime<Utc>"
291 ))
292 })
293}
294
295pub(crate) fn decode_timestampntz_micros(data: &[u8]) -> Result<NaiveDateTime, ArrowError> {
297 let micros_since_epoch = i64::from_le_bytes(array_from_slice(data, 0)?);
298 DateTime::from_timestamp_micros(micros_since_epoch)
299 .ok_or_else(|| {
300 ArrowError::CastError(format!(
301 "Could not cast `{micros_since_epoch}` microseconds into a NaiveDateTime"
302 ))
303 })
304 .map(|v| v.naive_utc())
305}
306
307pub(crate) fn decode_time_ntz(data: &[u8]) -> Result<NaiveTime, ArrowError> {
308 let micros_since_epoch = u64::from_le_bytes(array_from_slice(data, 0)?);
309
310 let case_error = ArrowError::CastError(format!(
311 "Could not cast {micros_since_epoch} microseconds into a NaiveTime"
312 ));
313
314 if micros_since_epoch >= 86_400_000_000 {
315 return Err(case_error);
316 }
317
318 let nanos_since_midnight = micros_since_epoch * 1_000;
319 NaiveTime::from_num_seconds_from_midnight_opt(
320 (nanos_since_midnight / 1_000_000_000) as u32,
321 (nanos_since_midnight % 1_000_000_000) as u32,
322 )
323 .ok_or(case_error)
324}
325
326pub(crate) fn decode_timestamp_nanos(data: &[u8]) -> Result<DateTime<Utc>, ArrowError> {
328 let nanos_since_epoch = i64::from_le_bytes(array_from_slice(data, 0)?);
329
330 Ok(DateTime::from_timestamp_nanos(nanos_since_epoch))
332}
333
334pub(crate) fn decode_timestampntz_nanos(data: &[u8]) -> Result<NaiveDateTime, ArrowError> {
336 decode_timestamp_nanos(data).map(|v| v.naive_utc())
337}
338
339pub(crate) fn decode_uuid(data: &[u8]) -> Result<Uuid, ArrowError> {
341 Uuid::from_slice(&data[0..16])
342 .map_err(|_| ArrowError::CastError(format!("Cant decode uuid from {:?}", &data[0..16])))
343}
344
345pub(crate) fn decode_binary(data: &[u8]) -> Result<&[u8], ArrowError> {
347 let len = u32::from_le_bytes(array_from_slice(data, 0)?) as usize;
348 slice_from_slice_at_offset(data, 4, 0..len)
349}
350
351pub(crate) fn decode_long_string(data: &[u8]) -> Result<&str, ArrowError> {
353 let len = u32::from_le_bytes(array_from_slice(data, 0)?) as usize;
354 string_from_slice(data, 4, 0..len)
355}
356
357pub(crate) fn decode_short_string(
359 metadata: u8,
360 data: &[u8],
361) -> Result<ShortString<'_>, ArrowError> {
362 let len = (metadata >> 2) as usize;
363 let string = string_from_slice(data, 0, 0..len)?;
364 ShortString::try_new(string)
365}
366
367#[cfg(test)]
368mod tests {
369 use super::*;
370
371 macro_rules! test_decoder_bounds {
372 ($test_name:ident, $data:expr, $decode_fn:ident, $expected:expr) => {
373 mod $test_name {
374 use super::*;
375
376 #[test]
377 fn exact_length() {
378 let result = $decode_fn(&$data).unwrap();
379 assert_eq!(result, $expected);
380 }
381
382 #[test]
383 fn truncated_length() {
384 let truncated_data = &$data[..$data.len() - 1];
386 let result = $decode_fn(truncated_data);
387 assert!(matches!(result, Err(ArrowError::InvalidArgumentError(_))));
388 }
389 }
390 };
391 }
392
393 mod integer {
394 use super::*;
395
396 test_decoder_bounds!(test_i8, [0x2a], decode_int8, 42);
397 test_decoder_bounds!(test_i16, [0xd2, 0x04], decode_int16, 1234);
398 test_decoder_bounds!(test_i32, [0x40, 0xe2, 0x01, 0x00], decode_int32, 123456);
399 test_decoder_bounds!(
400 test_i64,
401 [0x15, 0x81, 0xe9, 0x7d, 0xf4, 0x10, 0x22, 0x11],
402 decode_int64,
403 1234567890123456789
404 );
405 }
406
407 mod decimal {
408 use super::*;
409
410 test_decoder_bounds!(
411 test_decimal4,
412 [
413 0x02, 0xd2, 0x04, 0x00, 0x00, ],
416 decode_decimal4,
417 (1234, 2)
418 );
419
420 test_decoder_bounds!(
421 test_decimal8,
422 [
423 0x02, 0xd2, 0x02, 0x96, 0x49, 0x00, 0x00, 0x00, 0x00, ],
426 decode_decimal8,
427 (1234567890, 2)
428 );
429
430 test_decoder_bounds!(
431 test_decimal16,
432 [
433 0x02, 0xd2, 0xb6, 0x23, 0xc0, 0xf4, 0x10, 0x22, 0x11, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
435 0x00, 0x00, ],
437 decode_decimal16,
438 (1234567891234567890, 2)
439 );
440 }
441
442 mod float {
443 use super::*;
444
445 test_decoder_bounds!(
446 test_float,
447 [0x06, 0x2c, 0x93, 0x4e],
448 decode_float,
449 1234567890.1234
450 );
451
452 test_decoder_bounds!(
453 test_double,
454 [0xc9, 0xe5, 0x87, 0xb4, 0x80, 0x65, 0xd2, 0x41],
455 decode_double,
456 1234567890.1234
457 );
458 }
459
460 mod datetime {
461 use super::*;
462
463 test_decoder_bounds!(
464 test_date,
465 [0xe2, 0x4e, 0x0, 0x0],
466 decode_date,
467 NaiveDate::from_ymd_opt(2025, 4, 16).unwrap()
468 );
469
470 test_decoder_bounds!(
471 test_timestamp_micros,
472 [0xe0, 0x52, 0x97, 0xdd, 0xe7, 0x32, 0x06, 0x00],
473 decode_timestamp_micros,
474 NaiveDate::from_ymd_opt(2025, 4, 16)
475 .unwrap()
476 .and_hms_milli_opt(16, 34, 56, 780)
477 .unwrap()
478 .and_utc()
479 );
480
481 test_decoder_bounds!(
482 test_timestampntz_micros,
483 [0xe0, 0x52, 0x97, 0xdd, 0xe7, 0x32, 0x06, 0x00],
484 decode_timestampntz_micros,
485 NaiveDate::from_ymd_opt(2025, 4, 16)
486 .unwrap()
487 .and_hms_milli_opt(16, 34, 56, 780)
488 .unwrap()
489 );
490
491 test_decoder_bounds!(
492 test_timestamp_nanos,
493 [0x15, 0x41, 0xa2, 0x5a, 0x36, 0xa2, 0x5b, 0x18],
494 decode_timestamp_nanos,
495 NaiveDate::from_ymd_opt(2025, 8, 14)
496 .unwrap()
497 .and_hms_nano_opt(12, 33, 54, 123456789)
498 .unwrap()
499 .and_utc()
500 );
501
502 test_decoder_bounds!(
503 test_timestamp_nanos_before_epoch,
504 [0x15, 0x41, 0x52, 0xd4, 0x94, 0xe5, 0xad, 0xfa],
505 decode_timestamp_nanos,
506 NaiveDate::from_ymd_opt(1957, 11, 7)
507 .unwrap()
508 .and_hms_nano_opt(12, 33, 54, 123456789)
509 .unwrap()
510 .and_utc()
511 );
512
513 test_decoder_bounds!(
514 test_timestampntz_nanos,
515 [0x15, 0x41, 0xa2, 0x5a, 0x36, 0xa2, 0x5b, 0x18],
516 decode_timestampntz_nanos,
517 NaiveDate::from_ymd_opt(2025, 8, 14)
518 .unwrap()
519 .and_hms_nano_opt(12, 33, 54, 123456789)
520 .unwrap()
521 );
522
523 test_decoder_bounds!(
524 test_timestampntz_nanos_before_epoch,
525 [0x15, 0x41, 0x52, 0xd4, 0x94, 0xe5, 0xad, 0xfa],
526 decode_timestampntz_nanos,
527 NaiveDate::from_ymd_opt(1957, 11, 7)
528 .unwrap()
529 .and_hms_nano_opt(12, 33, 54, 123456789)
530 .unwrap()
531 );
532 }
533
534 #[test]
535 fn test_uuid() {
536 let data = [
537 0xf2, 0x4f, 0x9b, 0x64, 0x81, 0xfa, 0x49, 0xd1, 0xb7, 0x4e, 0x8c, 0x09, 0xa6, 0xe3,
538 0x1c, 0x56,
539 ];
540 let result = decode_uuid(&data).unwrap();
541 assert_eq!(
542 Uuid::parse_str("f24f9b64-81fa-49d1-b74e-8c09a6e31c56").unwrap(),
543 result
544 );
545 }
546
547 mod time {
548 use super::*;
549
550 test_decoder_bounds!(
551 test_timentz,
552 [0x53, 0x1f, 0x8e, 0xdf, 0x2, 0, 0, 0],
553 decode_time_ntz,
554 NaiveTime::from_num_seconds_from_midnight_opt(12340, 567_891_000).unwrap()
555 );
556
557 #[test]
558 fn test_decode_time_ntz_invalid() {
559 let invalid_second = u64::MAX;
560 let data = invalid_second.to_le_bytes();
561 let result = decode_time_ntz(&data);
562 assert!(matches!(result, Err(ArrowError::CastError(_))));
563 }
564 }
565
566 #[test]
567 fn test_binary_exact_length() {
568 let data = [
569 0x09, 0, 0, 0, 0x03, 0x13, 0x37, 0xde, 0xad, 0xbe, 0xef, 0xca, 0xfe,
571 ];
572 let result = decode_binary(&data).unwrap();
573 assert_eq!(
574 result,
575 [0x03, 0x13, 0x37, 0xde, 0xad, 0xbe, 0xef, 0xca, 0xfe]
576 );
577 }
578
579 #[test]
580 fn test_binary_truncated_length() {
581 let data = [
582 0x09, 0, 0, 0, 0x03, 0x13, 0x37, 0xde, 0xad, 0xbe, 0xef, 0xca,
584 ];
585 let result = decode_binary(&data);
586 assert!(matches!(result, Err(ArrowError::InvalidArgumentError(_))));
587 }
588
589 #[test]
590 fn test_short_string_exact_length() {
591 let data = b"Helloo";
592 let result = decode_short_string(1 | 5 << 2, data).unwrap();
593 assert_eq!(result.0, "Hello");
594 }
595
596 #[test]
597 fn test_short_string_truncated_length() {
598 let data = b"Hel";
599 let result = decode_short_string(1 | 5 << 2, data);
600 assert!(matches!(result, Err(ArrowError::InvalidArgumentError(_))));
601 }
602
603 #[test]
604 fn test_string_exact_length() {
605 let data = [
606 0x05, 0, 0, 0, b'H', b'e', b'l', b'l', b'o', b'o',
608 ];
609 let result = decode_long_string(&data).unwrap();
610 assert_eq!(result, "Hello");
611 }
612
613 #[test]
614 fn test_string_truncated_length() {
615 let data = [
616 0x05, 0, 0, 0, b'H', b'e', b'l',
618 ];
619 let result = decode_long_string(&data);
620 assert!(matches!(result, Err(ArrowError::InvalidArgumentError(_))));
621 }
622
623 #[test]
624 fn test_offset() {
625 assert_eq!(OffsetSizeBytes::try_new(0).unwrap(), OffsetSizeBytes::One);
626 assert_eq!(OffsetSizeBytes::try_new(1).unwrap(), OffsetSizeBytes::Two);
627 assert_eq!(OffsetSizeBytes::try_new(2).unwrap(), OffsetSizeBytes::Three);
628 assert_eq!(OffsetSizeBytes::try_new(3).unwrap(), OffsetSizeBytes::Four);
629
630 assert!(OffsetSizeBytes::try_new(4).is_err());
632 assert!(OffsetSizeBytes::try_new(255).is_err());
633 }
634
635 #[test]
636 fn unpack_u32_all_widths() {
637 let buf_one = [0x01u8, 0xAB, 0xCD];
639 assert_eq!(OffsetSizeBytes::One.unpack_u32(&buf_one, 0).unwrap(), 0x01);
640 assert_eq!(OffsetSizeBytes::One.unpack_u32(&buf_one, 2).unwrap(), 0xCD);
641
642 let buf_two = [0x34, 0x12, 0x78, 0x56];
644 assert_eq!(
645 OffsetSizeBytes::Two.unpack_u32(&buf_two, 0).unwrap(),
646 0x1234
647 );
648 assert_eq!(
649 OffsetSizeBytes::Two.unpack_u32(&buf_two, 1).unwrap(),
650 0x5678
651 );
652
653 let buf_three = [0x01, 0x02, 0x03, 0xFF, 0x00, 0x00];
655 assert_eq!(
656 OffsetSizeBytes::Three.unpack_u32(&buf_three, 0).unwrap(),
657 0x030201
658 );
659 assert_eq!(
660 OffsetSizeBytes::Three.unpack_u32(&buf_three, 1).unwrap(),
661 0x0000FF
662 );
663
664 let buf_four = [0x78, 0x56, 0x34, 0x12, 0xEF, 0xCD, 0xAB, 0x90];
666 assert_eq!(
667 OffsetSizeBytes::Four.unpack_u32(&buf_four, 0).unwrap(),
668 0x1234_5678
669 );
670 assert_eq!(
671 OffsetSizeBytes::Four.unpack_u32(&buf_four, 1).unwrap(),
672 0x90AB_CDEF
673 );
674 }
675
676 #[test]
677 fn unpack_u32_out_of_bounds() {
678 let tiny = [0x00u8]; assert!(OffsetSizeBytes::Two.unpack_u32(&tiny, 0).is_err());
680 assert!(OffsetSizeBytes::Three.unpack_u32(&tiny, 0).is_err());
681 }
682
683 #[test]
684 fn unpack_simple() {
685 let buf = [
686 0x41, 0x02, 0x00, 0x00, 0x00, 0x05, 0x00, 0x09, 0x00, ];
692
693 let width = OffsetSizeBytes::Two;
694
695 let dict_size = width.unpack_u32_at_offset(&buf, 1, 0).unwrap();
697 assert_eq!(dict_size, 2);
698
699 let first = width.unpack_u32_at_offset(&buf, 1, 1).unwrap();
701 assert_eq!(first, 0);
702
703 let second = width.unpack_u32_at_offset(&buf, 1, 2).unwrap();
704 assert_eq!(second, 5);
705
706 let third = width.unpack_u32_at_offset(&buf, 1, 3).unwrap();
707 assert_eq!(third, 9);
708
709 let err = width.unpack_u32_at_offset(&buf, 1, 4);
710 assert!(err.is_err())
711 }
712}