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 DateTime::UNIX_EPOCH
282 .checked_add_signed(Duration::days(i64::from(days_since_epoch)))
283 .map(|value| value.date_naive())
284 .ok_or_else(|| {
285 ArrowError::CastError(format!(
286 "Could not cast `{days_since_epoch}` days into a NaiveDate"
287 ))
288 })
289}
290
291pub(crate) fn decode_timestamp_micros(data: &[u8]) -> Result<DateTime<Utc>, ArrowError> {
293 let micros_since_epoch = i64::from_le_bytes(array_from_slice(data, 0)?);
294 DateTime::from_timestamp_micros(micros_since_epoch).ok_or_else(|| {
295 ArrowError::CastError(format!(
296 "Could not cast `{micros_since_epoch}` microseconds into a DateTime<Utc>"
297 ))
298 })
299}
300
301pub(crate) fn decode_timestampntz_micros(data: &[u8]) -> Result<NaiveDateTime, ArrowError> {
303 let micros_since_epoch = i64::from_le_bytes(array_from_slice(data, 0)?);
304 DateTime::from_timestamp_micros(micros_since_epoch)
305 .ok_or_else(|| {
306 ArrowError::CastError(format!(
307 "Could not cast `{micros_since_epoch}` microseconds into a NaiveDateTime"
308 ))
309 })
310 .map(|v| v.naive_utc())
311}
312
313pub(crate) fn decode_time_ntz(data: &[u8]) -> Result<NaiveTime, ArrowError> {
314 let micros_since_epoch = u64::from_le_bytes(array_from_slice(data, 0)?);
315
316 let case_error = ArrowError::CastError(format!(
317 "Could not cast {micros_since_epoch} microseconds into a NaiveTime"
318 ));
319
320 if micros_since_epoch >= 86_400_000_000 {
321 return Err(case_error);
322 }
323
324 let nanos_since_midnight = micros_since_epoch * 1_000;
325 NaiveTime::from_num_seconds_from_midnight_opt(
326 (nanos_since_midnight / 1_000_000_000) as u32,
327 (nanos_since_midnight % 1_000_000_000) as u32,
328 )
329 .ok_or(case_error)
330}
331
332pub(crate) fn decode_timestamp_nanos(data: &[u8]) -> Result<DateTime<Utc>, ArrowError> {
334 let nanos_since_epoch = i64::from_le_bytes(array_from_slice(data, 0)?);
335
336 Ok(DateTime::from_timestamp_nanos(nanos_since_epoch))
338}
339
340pub(crate) fn decode_timestampntz_nanos(data: &[u8]) -> Result<NaiveDateTime, ArrowError> {
342 decode_timestamp_nanos(data).map(|v| v.naive_utc())
343}
344
345pub(crate) fn decode_uuid(data: &[u8]) -> Result<Uuid, ArrowError> {
347 let bytes: [u8; 16] = array_from_slice(data, 0)?;
348 Ok(Uuid::from_bytes(bytes))
349}
350
351pub(crate) fn decode_binary(data: &[u8]) -> Result<&[u8], ArrowError> {
353 let len = u32::from_le_bytes(array_from_slice(data, 0)?) as usize;
354 slice_from_slice_at_offset(data, 4, 0..len)
355}
356
357pub(crate) fn decode_long_string(data: &[u8]) -> Result<&str, ArrowError> {
359 let len = u32::from_le_bytes(array_from_slice(data, 0)?) as usize;
360 string_from_slice(data, 4, 0..len)
361}
362
363pub(crate) fn decode_short_string(
365 metadata: u8,
366 data: &[u8],
367) -> Result<ShortString<'_>, ArrowError> {
368 let len = (metadata >> 2) as usize;
369 let string = string_from_slice(data, 0, 0..len)?;
370 ShortString::try_new(string)
371}
372
373#[cfg(test)]
374mod tests {
375 use super::*;
376
377 macro_rules! test_decoder_bounds {
378 ($test_name:ident, $data:expr, $decode_fn:ident, $expected:expr) => {
379 mod $test_name {
380 use super::*;
381
382 #[test]
383 fn exact_length() {
384 let result = $decode_fn(&$data).unwrap();
385 assert_eq!(result, $expected);
386 }
387
388 #[test]
389 fn truncated_length() {
390 let truncated_data = &$data[..$data.len() - 1];
392 let result = $decode_fn(truncated_data);
393 assert!(matches!(result, Err(ArrowError::InvalidArgumentError(_))));
394 }
395 }
396 };
397 }
398
399 mod integer {
400 use super::*;
401
402 test_decoder_bounds!(test_i8, [0x2a], decode_int8, 42);
403 test_decoder_bounds!(test_i16, [0xd2, 0x04], decode_int16, 1234);
404 test_decoder_bounds!(test_i32, [0x40, 0xe2, 0x01, 0x00], decode_int32, 123456);
405 test_decoder_bounds!(
406 test_i64,
407 [0x15, 0x81, 0xe9, 0x7d, 0xf4, 0x10, 0x22, 0x11],
408 decode_int64,
409 1234567890123456789
410 );
411 }
412
413 mod decimal {
414 use super::*;
415
416 test_decoder_bounds!(
417 test_decimal4,
418 [
419 0x02, 0xd2, 0x04, 0x00, 0x00, ],
422 decode_decimal4,
423 (1234, 2)
424 );
425
426 test_decoder_bounds!(
427 test_decimal8,
428 [
429 0x02, 0xd2, 0x02, 0x96, 0x49, 0x00, 0x00, 0x00, 0x00, ],
432 decode_decimal8,
433 (1234567890, 2)
434 );
435
436 test_decoder_bounds!(
437 test_decimal16,
438 [
439 0x02, 0xd2, 0xb6, 0x23, 0xc0, 0xf4, 0x10, 0x22, 0x11, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
441 0x00, 0x00, ],
443 decode_decimal16,
444 (1234567891234567890, 2)
445 );
446 }
447
448 mod float {
449 use super::*;
450
451 test_decoder_bounds!(
452 test_float,
453 [0x06, 0x2c, 0x93, 0x4e],
454 decode_float,
455 1_234_568_000.0
456 );
457
458 test_decoder_bounds!(
459 test_double,
460 [0xc9, 0xe5, 0x87, 0xb4, 0x80, 0x65, 0xd2, 0x41],
461 decode_double,
462 1234567890.1234
463 );
464 }
465
466 mod datetime {
467 use super::*;
468
469 test_decoder_bounds!(
470 test_date,
471 [0xe2, 0x4e, 0x0, 0x0],
472 decode_date,
473 NaiveDate::from_ymd_opt(2025, 4, 16).unwrap()
474 );
475
476 #[test]
477 fn test_date_out_of_range() {
478 for days in [i32::MAX, i32::MIN] {
480 let result = decode_date(&days.to_le_bytes());
481 assert!(matches!(result, Err(ArrowError::CastError(_))));
482 }
483 }
484
485 test_decoder_bounds!(
486 test_timestamp_micros,
487 [0xe0, 0x52, 0x97, 0xdd, 0xe7, 0x32, 0x06, 0x00],
488 decode_timestamp_micros,
489 NaiveDate::from_ymd_opt(2025, 4, 16)
490 .unwrap()
491 .and_hms_milli_opt(16, 34, 56, 780)
492 .unwrap()
493 .and_utc()
494 );
495
496 test_decoder_bounds!(
497 test_timestampntz_micros,
498 [0xe0, 0x52, 0x97, 0xdd, 0xe7, 0x32, 0x06, 0x00],
499 decode_timestampntz_micros,
500 NaiveDate::from_ymd_opt(2025, 4, 16)
501 .unwrap()
502 .and_hms_milli_opt(16, 34, 56, 780)
503 .unwrap()
504 );
505
506 test_decoder_bounds!(
507 test_timestamp_nanos,
508 [0x15, 0x41, 0xa2, 0x5a, 0x36, 0xa2, 0x5b, 0x18],
509 decode_timestamp_nanos,
510 NaiveDate::from_ymd_opt(2025, 8, 14)
511 .unwrap()
512 .and_hms_nano_opt(12, 33, 54, 123456789)
513 .unwrap()
514 .and_utc()
515 );
516
517 test_decoder_bounds!(
518 test_timestamp_nanos_before_epoch,
519 [0x15, 0x41, 0x52, 0xd4, 0x94, 0xe5, 0xad, 0xfa],
520 decode_timestamp_nanos,
521 NaiveDate::from_ymd_opt(1957, 11, 7)
522 .unwrap()
523 .and_hms_nano_opt(12, 33, 54, 123456789)
524 .unwrap()
525 .and_utc()
526 );
527
528 test_decoder_bounds!(
529 test_timestampntz_nanos,
530 [0x15, 0x41, 0xa2, 0x5a, 0x36, 0xa2, 0x5b, 0x18],
531 decode_timestampntz_nanos,
532 NaiveDate::from_ymd_opt(2025, 8, 14)
533 .unwrap()
534 .and_hms_nano_opt(12, 33, 54, 123456789)
535 .unwrap()
536 );
537
538 test_decoder_bounds!(
539 test_timestampntz_nanos_before_epoch,
540 [0x15, 0x41, 0x52, 0xd4, 0x94, 0xe5, 0xad, 0xfa],
541 decode_timestampntz_nanos,
542 NaiveDate::from_ymd_opt(1957, 11, 7)
543 .unwrap()
544 .and_hms_nano_opt(12, 33, 54, 123456789)
545 .unwrap()
546 );
547 }
548
549 test_decoder_bounds!(
550 test_uuid,
551 [
552 0xf2, 0x4f, 0x9b, 0x64, 0x81, 0xfa, 0x49, 0xd1, 0xb7, 0x4e, 0x8c, 0x09, 0xa6, 0xe3,
553 0x1c, 0x56,
554 ],
555 decode_uuid,
556 Uuid::parse_str("f24f9b64-81fa-49d1-b74e-8c09a6e31c56").unwrap()
557 );
558
559 mod time {
560 use super::*;
561
562 test_decoder_bounds!(
563 test_timentz,
564 [0x53, 0x1f, 0x8e, 0xdf, 0x2, 0, 0, 0],
565 decode_time_ntz,
566 NaiveTime::from_num_seconds_from_midnight_opt(12340, 567_891_000).unwrap()
567 );
568
569 #[test]
570 fn test_decode_time_ntz_invalid() {
571 let invalid_second = u64::MAX;
572 let data = invalid_second.to_le_bytes();
573 let result = decode_time_ntz(&data);
574 assert!(matches!(result, Err(ArrowError::CastError(_))));
575 }
576 }
577
578 #[test]
579 fn test_binary_exact_length() {
580 let data = [
581 0x09, 0, 0, 0, 0x03, 0x13, 0x37, 0xde, 0xad, 0xbe, 0xef, 0xca, 0xfe,
583 ];
584 let result = decode_binary(&data).unwrap();
585 assert_eq!(
586 result,
587 [0x03, 0x13, 0x37, 0xde, 0xad, 0xbe, 0xef, 0xca, 0xfe]
588 );
589 }
590
591 #[test]
592 fn test_binary_truncated_length() {
593 let data = [
594 0x09, 0, 0, 0, 0x03, 0x13, 0x37, 0xde, 0xad, 0xbe, 0xef, 0xca,
596 ];
597 let result = decode_binary(&data);
598 assert!(matches!(result, Err(ArrowError::InvalidArgumentError(_))));
599 }
600
601 #[test]
602 fn test_short_string_exact_length() {
603 let data = b"Helloo";
604 let result = decode_short_string(1 | (5 << 2), data).unwrap();
605 assert_eq!(result.0, "Hello");
606 }
607
608 #[test]
609 fn test_short_string_truncated_length() {
610 let data = b"Hel";
611 let result = decode_short_string(1 | (5 << 2), data);
612 assert!(matches!(result, Err(ArrowError::InvalidArgumentError(_))));
613 }
614
615 #[test]
616 fn test_string_exact_length() {
617 let data = [
618 0x05, 0, 0, 0, b'H', b'e', b'l', b'l', b'o', b'o',
620 ];
621 let result = decode_long_string(&data).unwrap();
622 assert_eq!(result, "Hello");
623 }
624
625 #[test]
626 fn test_string_truncated_length() {
627 let data = [
628 0x05, 0, 0, 0, b'H', b'e', b'l',
630 ];
631 let result = decode_long_string(&data);
632 assert!(matches!(result, Err(ArrowError::InvalidArgumentError(_))));
633 }
634
635 #[test]
636 fn test_offset() {
637 assert_eq!(OffsetSizeBytes::try_new(0).unwrap(), OffsetSizeBytes::One);
638 assert_eq!(OffsetSizeBytes::try_new(1).unwrap(), OffsetSizeBytes::Two);
639 assert_eq!(OffsetSizeBytes::try_new(2).unwrap(), OffsetSizeBytes::Three);
640 assert_eq!(OffsetSizeBytes::try_new(3).unwrap(), OffsetSizeBytes::Four);
641
642 assert!(OffsetSizeBytes::try_new(4).is_err());
644 assert!(OffsetSizeBytes::try_new(255).is_err());
645 }
646
647 #[test]
648 fn unpack_u32_all_widths() {
649 let buf_one = [0x01u8, 0xAB, 0xCD];
651 assert_eq!(OffsetSizeBytes::One.unpack_u32(&buf_one, 0).unwrap(), 0x01);
652 assert_eq!(OffsetSizeBytes::One.unpack_u32(&buf_one, 2).unwrap(), 0xCD);
653
654 let buf_two = [0x34, 0x12, 0x78, 0x56];
656 assert_eq!(
657 OffsetSizeBytes::Two.unpack_u32(&buf_two, 0).unwrap(),
658 0x1234
659 );
660 assert_eq!(
661 OffsetSizeBytes::Two.unpack_u32(&buf_two, 1).unwrap(),
662 0x5678
663 );
664
665 let buf_three = [0x01, 0x02, 0x03, 0xFF, 0x00, 0x00];
667 assert_eq!(
668 OffsetSizeBytes::Three.unpack_u32(&buf_three, 0).unwrap(),
669 0x030201
670 );
671 assert_eq!(
672 OffsetSizeBytes::Three.unpack_u32(&buf_three, 1).unwrap(),
673 0x0000FF
674 );
675
676 let buf_four = [0x78, 0x56, 0x34, 0x12, 0xEF, 0xCD, 0xAB, 0x90];
678 assert_eq!(
679 OffsetSizeBytes::Four.unpack_u32(&buf_four, 0).unwrap(),
680 0x1234_5678
681 );
682 assert_eq!(
683 OffsetSizeBytes::Four.unpack_u32(&buf_four, 1).unwrap(),
684 0x90AB_CDEF
685 );
686 }
687
688 #[test]
689 fn unpack_u32_out_of_bounds() {
690 let tiny = [0x00u8]; assert!(OffsetSizeBytes::Two.unpack_u32(&tiny, 0).is_err());
692 assert!(OffsetSizeBytes::Three.unpack_u32(&tiny, 0).is_err());
693 }
694
695 #[test]
696 fn unpack_simple() {
697 let buf = [
698 0x41, 0x02, 0x00, 0x00, 0x00, 0x05, 0x00, 0x09, 0x00, ];
704
705 let width = OffsetSizeBytes::Two;
706
707 let dict_size = width.unpack_u32_at_offset(&buf, 1, 0).unwrap();
709 assert_eq!(dict_size, 2);
710
711 let first = width.unpack_u32_at_offset(&buf, 1, 1).unwrap();
713 assert_eq!(first, 0);
714
715 let second = width.unpack_u32_at_offset(&buf, 1, 2).unwrap();
716 assert_eq!(second, 5);
717
718 let third = width.unpack_u32_at_offset(&buf, 1, 3).unwrap();
719 assert_eq!(third, 9);
720
721 let err = width.unpack_u32_at_offset(&buf, 1, 4);
722 assert!(err.is_err())
723 }
724}